SE: Tests for exhausted plan

This commit introduces 2 tests for the fix done in previous commit.
Tests use Tokio clock pausing to deterministically test the execute
function.
First test case verifies that no new fiber is spawned after exhaustion,
and that we wait for existing fibers.
Second checks that we return immediately if we are the last fiber.

Author: Lorak-mmk

scylla/src/policies/speculative_execution.rs

@@ -144,7 +144,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>>>,
@@ -195,3 +195,131 @@ 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()
+        )
+    }
+}