Sampler: Don't consider query timeout as a mismatch

This commit fixed a few mirror issues:
- If a timeout is detected, the query should not be considered as a
mismatch.
- Add a config option to specify the timeout for sampler queries.
- Register query sampler with proper connection attributes.

Closes: REA-6228
Closes: REA-6244
Fixes: #1576
Fixes: #1573

Release-Note-Core: Allow configuring the timeout for sampler queries.

Change-Id: I8d47d40f56ae0802e0c906da589135771172b7ab
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/11132
Reviewed-by: Jason Brown <jason.b@readyset.io>
Tested-by: Buildkite CI

Author: altmannmarcelo

readyset-adapter/src/sampler.rs

@@ -180,7 +180,7 @@ async fn connect_rs(config: &UpstreamConfig, rs_addr: SocketAddr) -> Option<Data
 
 async fn connect_upstream(config: &UpstreamConfig) -> Option<DatabaseConnection> {
     debug!("Connecting to upstream");
-    let url: DatabaseURL = if let Some(url) = &config.upstream_db_url {
+    let mut url: DatabaseURL = if let Some(url) = &config.upstream_db_url {
         match url.parse() {
             Ok(url) => url,
             Err(error) => {
@@ -191,6 +191,7 @@ async fn connect_upstream(config: &UpstreamConfig) -> Option<DatabaseConnection>
     } else {
         return None;
     };
+    url.set_program_or_application_name(READYSET_QUERY_SAMPLER.to_string());
     let verification = match ServerCertVerification::from(config).await {
         Ok(verification) => verification,
         Err(error) => {
@@ -299,6 +300,7 @@ impl Sampler {
         timeout: Duration,
         conn: &mut Option<DatabaseConnection>,
         q: &str,
+        name: String,
     ) -> Result<QueryResults, DatabaseError> {
         let fut = async {
             match conn.as_mut() {
@@ -310,7 +312,7 @@ impl Sampler {
             Ok(res) => res,
             Err(_) => {
                 rate_limit(true, SAMPLER_LOG_SAMPLER, || {
-                    warn!("Sampler upstream read timeout");
+                    debug!(?name, "Sampler upstream read timeout");
                 });
                 Err(DatabaseError::UpstreamQueryTimeout)
             }
@@ -348,9 +350,13 @@ impl Sampler {
         );
 
         // Readyset execution (materialize rows for optional tracing on mismatch)
-        let rs_res =
-            Self::query_with_timeout(self.config.upstream_timeout, &mut self.rs_conn, &entry.q)
-                .await;
+        let rs_res = Self::query_with_timeout(
+            self.config.upstream_timeout,
+            &mut self.rs_conn,
+            &entry.q,
+            get_cache_name(&entry.event),
+        )
+        .await;
         let (rs_hash_opt, rs_rows_opt) = match rs_res {
             Ok(rs_rows) => {
                 let mut rows_vec: Vec<Vec<DfValue>> = match <Vec<Vec<DfValue>>>::try_from(rs_rows) {
@@ -363,6 +369,12 @@ impl Sampler {
                 let hash = Some(normalize_and_hash(&mut rows_vec, has_order_by));
                 (hash, Some(rows_vec))
             }
+            Err(DatabaseError::UpstreamQueryTimeout) => {
+                rate_limit(true, SAMPLER_LOG_SAMPLER, || {
+                    debug!("Sampler readyset query timed out; discarding sample");
+                });
+                return;
+            }
             Err(e) => {
                 debug!(error = %e, "Sampler readyset error");
                 self.reconnect_rs_if_closed().await;
@@ -384,6 +396,7 @@ impl Sampler {
             self.config.upstream_timeout,
             &mut self.upstream_conn,
             &entry.q,
+            get_cache_name(&entry.event),
         )
         .await;
         let (up_hash_opt, up_rows_opt) = match up_res {
@@ -398,6 +411,12 @@ impl Sampler {
                 let hash = Some(normalize_and_hash(&mut rows_vec, has_order_by));
                 (hash, Some(rows_vec))
             }
+            Err(DatabaseError::UpstreamQueryTimeout) => {
+                rate_limit(true, SAMPLER_LOG_SAMPLER, || {
+                    debug!("Sampler upstream query timed out; discarding sample");
+                });
+                return;
+            }
             Err(e) => {
                 debug!(error = %e, "Sampler upstream error");
                 self.reconnect_upstream_if_closed().await;

readyset/src/lib.rs

@@ -458,6 +458,10 @@ pub struct Options {
     #[arg(long, env = "SAMPLER_MAX_QUERIES_PER_SECOND", default_value = "100")]
     sampler_max_queries_per_second: u64,
 
+    /// Specifies the timeout for sampler queries.
+    #[arg(long, env = "SAMPLER_QUERY_TIMEOUT_MS", default_value = "1000")]
+    sampler_query_timeout_ms: u64,
+
     /// If set, Readyset will only verify if the current configuration is valid and then exit.
     ///
     /// On success, the exit code will be zero with a message printed to stdout.  On failure, the
@@ -1301,7 +1305,7 @@ where
             readyset_adapter::sampler::SamplerConfig {
                 sample_rate: options.sampler_sample_rate,
                 queue_capacity: options.sampler_queue_capacity,
-                upstream_timeout: Duration::from_secs(1),
+                upstream_timeout: Duration::from_millis(options.sampler_query_timeout_ms),
                 max_retry_attempts: options.sampler_max_retries,
                 retry_delay: Duration::from_millis(options.sampler_retry_delay_ms),
                 max_qps: options.sampler_max_queries_per_second,