readyset-adapter: Add Query Sampler

This commit adds a Query Sampler thread.

The sampler thread is responsible for probabilistically sampling
queries that were originally sent to Readyset. The intention is to
pro-actively validate the correctness of the Readyset Server.

The sampler thread will run queries against the upstream database and
compare the normalized results with the results returned by the
Readyset Server.

Queries are passed from adapter to sampler via a bounded channel,
which will drop queries if the channel is full.

The sampler thread has two main queues:
- A queue of queries that have been sent from the adapter. Those are
randomly sampled.
- A retry queue. In case of a mismatch, we add queries to a second
queue to be retried.

Re-try logic is implemented by a new struct that keeps track of some
metadata about the query, such as last upstream result hash, the
number or retry attempts, and the query itself.

The retry queue is a double ended queue, since we need to be able
to add elements in the end of the queue and pop elements in the front.
Elements are added in retry order. The retry is based on a timestamp
calculated from the current time + the retry delay.

The retry queue is bounded, and will drop entries if it is full.

We added the following metrics:
- `readyset_query_sampler_queue_len`: The number of queries in the
sampler queue.
- `readyset_query_sampler_queries_sampled`: The number of queries
sampled.
- `readyset_query_sampler_queries_mismatched`: The number of
queries that were sampled but returned different results.
- `readyset_query_sampler_retry_queue_len`: The number of entries in
the retry queue.
- `readyset_query_sampler_retry_queue_full`: The number of times the
retry queue is full.
- `readyset_query_sampler_max_qps_hit`: The number of times the
sampler rate limiter is hit.
- `readyset_query_sampler_reconnects`: The number of times the
sampler reconnects to the upstream database.

We also added new configuration options:
- `SAMPLER_SAMPLE_RATE`: The probability of sampling a query.
- `SAMPLER_RETRY_DELAY_MS`: The delay between retries.
- `SAMPLER_QUEUE_CAPACITY`: The maximum number of queries in the
sampler queue.
- `SAMPLER_MAX_RETRIES`: The maximum number of retries for a query.
- `SAMPLER_MAX_QUERIES_PER_SECOND`: The maximum number of queries to
sample per second to avoid overwhelming the upstream database.

Note: The current implementation is limited to Text Protocol queries.
A further CL will add support for Binary Protocol queries.

Fixes REA-5328.

Release-Note-Core: Added a query sampler thread to the Readyset
  Adapter. This thread will sample queries and compare the results
  with the results returned by the upstream database and report possible
  mismatches.

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

Author: altmannmarcelo

Cargo.lock

@@ -77,7 +77,7 @@ resolver = "2"
 consulrs = { git = "https://github.com/readysettech/consulrs.git", branch = "allow-disabling-rustls-tls-2" }
 
 # overridden in the patch section above
-mysql_async = "0.36.0"
+mysql_async = "0.36.1"
 
 # some of these are overridden in the patch section above
 postgres = "0.19.9"

Cargo.toml

@@ -302,6 +302,15 @@ impl DatabaseConnection {
         }
     }
 
+    /// Returns true if the connection is closed
+    pub fn is_closed(&self) -> bool {
+        match self {
+            Self::PostgreSQL(client, _) => client.is_closed(),
+            Self::PostgreSQLPool(client) => client.is_closed(),
+            Self::MySQL(conn) => conn.is_disconnected(),
+        }
+    }
+
     /// Returns the SQL dialect associated with the underlying connection type.
     pub fn dialect(&self) -> Dialect {
         match self {

database-utils/src/connection.rs

@@ -33,6 +33,12 @@ pub enum DatabaseError {
 
     #[error("Error interpolating parameters into query: {0}")]
     Interpolation(Box<dyn std::error::Error + Send + Sync>),
+
+    #[error("Upstream connection is None")]
+    UpstreamConnectionNone,
+
+    #[error("Upstream query timeout")]
+    UpstreamQueryTimeout,
 }
 
 impl DatabaseError {

database-utils/src/error.rs

@@ -37,6 +37,8 @@ lru = { workspace = true }
 crossbeam-skiplist = { workspace = true }
 slab = { workspace = true }
 xxhash-rust = { workspace = true }
+rand = { workspace = true }
+streaming-iterator = { workspace = true }
 
 readyset-adapter-types = { path = "../readyset-adapter-types/" }
 readyset-alloc = { path = "../readyset-alloc/" }
@@ -46,6 +48,8 @@ readyset-data = { path = "../readyset-data/" }
 readyset-server = { path = "../readyset-server" }
 dataflow-expression = { path = "../dataflow-expression" }
 readyset-tracing = { path = "../readyset-tracing" }
+readyset-dataflow = { path = "../readyset-dataflow" }
+
 readyset-client-metrics = { path = "../readyset-client-metrics" }
 readyset-telemetry-reporter = { path = "../readyset-telemetry-reporter" }
 readyset-sql = { path = "../readyset-sql" }

readyset-adapter/Cargo.toml

@@ -294,6 +294,7 @@ pub struct BackendBuilder {
     metrics_handle: Option<MetricsHandle>,
     connections: Option<Arc<SkipSet<SocketAddr>>>,
     allow_cache_ddl: bool,
+    sampler_tx: Option<tokio::sync::mpsc::Sender<(QueryExecutionEvent, String)>>,
 }
 
 impl Default for BackendBuilder {
@@ -316,6 +317,7 @@ impl Default for BackendBuilder {
             metrics_handle: None,
             connections: None,
             allow_cache_ddl: true,
+            sampler_tx: None,
         }
     }
 }
@@ -380,6 +382,7 @@ impl BackendBuilder {
             status_reporter,
             allow_cache_ddl: self.allow_cache_ddl,
             adapter_start_time,
+            sampler_tx: self.sampler_tx,
             _query_handler: PhantomData,
         }
     }
@@ -474,6 +477,15 @@ impl BackendBuilder {
         self.metrics_handle = metrics_handle;
         self
     }
+
+    /// Set the sender used to enqueue original queries for background sampling/verification
+    pub fn sampler_tx(
+        mut self,
+        tx: Option<tokio::sync::mpsc::Sender<(QueryExecutionEvent, String)>>,
+    ) -> Self {
+        self.sampler_tx = tx;
+        self
+    }
 }
 
 /// A [`PreparedStatement`] stores the data needed for an immediate execution of a prepared
@@ -588,6 +600,9 @@ where
     /// The time at which the adapter started.
     adapter_start_time: SystemTime,
 
+    /// Optional sender to enqueue original queries for background sampling/verification
+    sampler_tx: Option<tokio::sync::mpsc::Sender<(QueryExecutionEvent, String)>>,
+
     _query_handler: PhantomData<Handler>,
 }
 
@@ -3124,7 +3139,7 @@ where
                     self.noria_should_try_select(&mut view_request);
 
                 if noria_should_try {
-                    Self::query_adhoc_select(
+                    let result = Self::query_adhoc_select(
                         &mut self.noria,
                         self.upstream.as_mut(),
                         &self.settings,
@@ -3135,7 +3150,16 @@ where
                         &mut event,
                         processed_query_params.unwrap(),
                     )
-                    .await
+                    .await;
+
+                    // Enqueue the original query for background sampling if enabled. We clone here
+                    // for simplicity and to avoid borrow issues in the hot path; bounded channel
+                    // will drop on overflow.
+                    if let Some(tx) = &self.sampler_tx {
+                        let _ = tx.try_send((event.clone(), query.to_string()));
+                    }
+
+                    result
                 } else {
                     Self::query_fallback(self.upstream.as_mut(), query, &mut event).await
                 }

readyset-adapter/src/backend.rs

@@ -6,6 +6,7 @@ pub mod migration_handler;
 pub mod proxied_queries_reporter;
 mod query_handler;
 pub mod query_status_cache;
+pub mod sampler;
 mod status_reporter;
 pub mod upstream_database;
 mod utils;