Merge pull request #1401 from Lorak-mmk/fix-se-plan-exhaustion

Fix plan exhaustion handling in SE and introduce SE unit tests.

Author: Lorak-mmk

scylla/src/policies/speculative_execution.rs

@@ -21,6 +21,7 @@ pub struct Context {
 
 /// The policy that decides if the driver will send speculative queries to the
 /// next targets when the current target takes too long to respond.
+// TODO(2.0): Consider renaming the methods to get rid of "retry" naming.
 pub trait SpeculativeExecutionPolicy: std::fmt::Debug + Send + Sync {
     /// The maximum number of speculative executions that will be triggered
     /// for a given request (does not include the initial request)
@@ -144,7 +145,7 @@ const EMPTY_PLAN_ERROR: RequestError = RequestError::EmptyPlan;
 pub(crate) async fn execute<QueryFut, ResT>(
     policy: &dyn SpeculativeExecutionPolicy,
     context: &Context,
-    query_runner_generator: impl Fn(bool) -> QueryFut,
+    mut query_runner_generator: impl FnMut(bool) -> QueryFut,
 ) -> Result<(ResT, Coordinator), RequestError>
 where
     QueryFut: Future<Output = Option<Result<(ResT, Coordinator), RequestError>>>,
@@ -180,6 +181,11 @@ where
                     } else {
                         last_error = Some(r)
                     }
+                } else {
+                    // The only case where None is returned is when execution plan was exhausted.
+                    // If so, there is no reason to start any more fibers.
+                    // We can't always return - there may still be fibers running.
+                    retries_remaining = 0;
                 }
                 if async_tasks.is_empty() && retries_remaining == 0 {
                     return last_error.unwrap_or({
@@ -190,3 +196,166 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    // Important to start tests with paused clock. If starting unpaused, and calling `tokio::time::pause()`, then
+    // things like `sleep` will advance the timer not fully accurately (I have no idea why), causing
+    // few ms added clock advancement at the end of the test.
+    // Starting paused is done with `#[tokio::test(flavor = "current_thread", start_paused = true)]`.
+    // Pausing can only be done with current_thread executor.
+
+    #[cfg(feature = "metrics")]
+    use std::sync::Arc;
+    use std::sync::LazyLock;
+    use std::time::Duration;
+
+    use assert_matches::assert_matches;
+
+    use crate::errors::{RequestAttemptError, RequestError};
+    #[cfg(feature = "metrics")]
+    use crate::observability::metrics::Metrics;
+    use crate::policies::speculative_execution::{Context, SimpleSpeculativeExecutionPolicy};
+    use crate::response::Coordinator;
+
+    static EMPTY_CONTEXT: LazyLock<Context> = LazyLock::new(|| Context {
+        #[cfg(feature = "metrics")]
+        metrics: Arc::new(Metrics::new()),
+    });
+
+    static IGNORABLE_ERROR: Option<Result<((), Coordinator), RequestError>> = Some(Err(
+        RequestError::LastAttemptError(RequestAttemptError::UnableToAllocStreamId),
+    ));
+
+    #[tokio::test(flavor = "current_thread", start_paused = true)]
+    async fn test_exhausted_plan_with_running_fibers() {
+        let policy = SimpleSpeculativeExecutionPolicy {
+            max_retry_count: 5,
+            retry_interval: Duration::from_secs(1),
+        };
+
+        let generator = {
+            // Index of the fiber, 0 for first execution.
+            let mut counter = 0;
+            move |_first: bool| {
+                let future = {
+                    let fiber_idx = counter;
+                    async move {
+                        if fiber_idx < 4 {
+                            tokio::time::sleep(Duration::from_secs(5)).await;
+                            IGNORABLE_ERROR.clone()
+                        } else if fiber_idx == 4 {
+                            None
+                        } else {
+                            panic!("Too many speculative executions - expected 4");
+                        }
+                    }
+                };
+                counter += 1;
+                future
+            }
+        };
+
+        let now = tokio::time::Instant::now();
+        let res = super::execute(&policy, &EMPTY_CONTEXT, generator).await;
+        assert_matches!(
+            res,
+            Err(RequestError::LastAttemptError(
+                RequestAttemptError::UnableToAllocStreamId
+            ))
+        );
+        // t - now
+        // First execution is started at t
+        // Speculative executions - at t+1, t+2, t+3, t+4
+        // The one at t+4 will return first, with None, preventing starting new one at t+5.
+        // Then execute should wait on spawned fibers. Last one will be the one spawned at t+3, finishing at t+8.
+        assert_eq!(
+            tokio::time::Instant::now(),
+            now.checked_add(Duration::from_secs(8)).unwrap()
+        )
+    }
+
+    #[tokio::test(flavor = "current_thread", start_paused = true)]
+    async fn test_exhausted_plan_last_running_fiber() {
+        let policy = SimpleSpeculativeExecutionPolicy {
+            max_retry_count: 5,
+            // Each attempt will finish before next starts
+            retry_interval: Duration::from_secs(6),
+        };
+
+        let generator = {
+            // Index of the fiber, 0 for first execution.
+            let mut counter = 0;
+            move |_first: bool| {
+                let future = {
+                    let fiber_idx = counter;
+                    async move {
+                        if fiber_idx < 4 {
+                            tokio::time::sleep(Duration::from_secs(5)).await;
+                            IGNORABLE_ERROR.clone()
+                        } else if fiber_idx == 4 {
+                            None
+                        } else {
+                            panic!("Too many speculative executions - expected 4");
+                        }
+                    }
+                };
+                counter += 1;
+                future
+            }
+        };
+
+        let now = tokio::time::Instant::now();
+        let res = super::execute(&policy, &EMPTY_CONTEXT, generator).await;
+        assert_matches!(
+            res,
+            Err(RequestError::LastAttemptError(
+                RequestAttemptError::UnableToAllocStreamId
+            ))
+        );
+        // t - now
+        // First execution is started at t
+        // Speculative executions - at t+6, t+12, t+18, t+24
+        // Each execution finishes before next starts. The one at t+24 finishes instantly with
+        // None, so the next one should not be started.
+        assert_eq!(
+            tokio::time::Instant::now(),
+            now.checked_add(Duration::from_secs(24)).unwrap()
+        )
+    }
+
+    // Regresion test for https://github.com/scylladb/scylla-rust-driver/issues/1085
+    #[tokio::test(flavor = "current_thread", start_paused = true)]
+    async fn test_se_panic_on_ignorable_errors() {
+        let policy = SimpleSpeculativeExecutionPolicy {
+            max_retry_count: 5,
+            // Each attempt will finish before next starts
+            retry_interval: Duration::from_secs(1),
+        };
+
+        let generator = {
+            move |_first: bool| async move {
+                tokio::time::sleep(Duration::from_secs(5)).await;
+                IGNORABLE_ERROR.clone()
+            }
+        };
+
+        let now = tokio::time::Instant::now();
+        let res = super::execute(&policy, &EMPTY_CONTEXT, generator).await;
+        assert_matches!(
+            res,
+            Err(RequestError::LastAttemptError(
+                RequestAttemptError::UnableToAllocStreamId
+            ))
+        );
+        // t - now
+        // First execution is started at t
+        // Speculative executions - at t+1, t+2, t+3, t+4, t+5
+        // Each execution sleeps 5 seconds and returns ignorable error.
+        // Last execution should finish at t+10.
+        assert_eq!(
+            tokio::time::Instant::now(),
+            now.checked_add(Duration::from_secs(10)).unwrap()
+        )
+    }
+}