Merge pull request #1457 from wprzytula/request-timeout-improvements

wprzytula · web-flow · commit 78fb4de47e03 · 2025-10-28T15:40:16.000+01:00
Request timeout improvements
diff --git a/scylla/src/client/execution_profile.rs b/scylla/src/client/execution_profile.rs
@@ -240,7 +240,8 @@ pub struct ExecutionProfileBuilder {
 }
 
 impl ExecutionProfileBuilder {
-    /// Changes client-side timeout.
+    /// Changes client-side timeout for executing statements.
+    /// If set to None, the driver will wait indefinitely for a response from the server.
     /// The default is 30 seconds.
     ///
     /// # Example
diff --git a/scylla/src/client/pager.rs b/scylla/src/client/pager.rs
@@ -144,7 +144,7 @@ struct PagerWorker<'a, QueryFunc, SpanCreatorFunc> {
     page_query: QueryFunc,
 
     load_balancing_policy: Arc<dyn LoadBalancingPolicy>,
-    statement_info: RoutingInfo<'a>,
+    routing_info: RoutingInfo<'a>,
     query_is_idempotent: bool,
     query_consistency: Consistency,
     retry_session: Box<dyn RetrySession>,
@@ -170,7 +170,7 @@ where
     // Contract: this function MUST send at least one item through self.sender
     async fn work(mut self, cluster_state: Arc<ClusterState>) -> PageSendAttemptedProof {
         let load_balancer = Arc::clone(&self.load_balancing_policy);
-        let statement_info = self.statement_info.clone();
+        let statement_info = self.routing_info.clone();
         let query_plan =
             load_balancing::Plan::new(load_balancer.as_ref(), &statement_info, &cluster_state);
 
@@ -359,7 +359,7 @@ where
                 self.log_attempt_success();
                 self.log_request_success();
                 self.load_balancing_policy
-                    .on_request_success(&self.statement_info, elapsed, node);
+                    .on_request_success(&self.routing_info, elapsed, node);
 
                 request_span.record_raw_rows_fields(&rows);
 
@@ -396,7 +396,7 @@ where
                 #[cfg(feature = "metrics")]
                 self.metrics.inc_failed_paged_queries();
                 self.load_balancing_policy.on_request_failure(
-                    &self.statement_info,
+                    &self.routing_info,
                     elapsed,
                     node,
                     &err,
@@ -424,7 +424,7 @@ where
                 let err =
                     RequestAttemptError::UnexpectedResponse(response.response.to_response_kind());
                 self.load_balancing_policy.on_request_failure(
-                    &self.statement_info,
+                    &self.routing_info,
                     elapsed,
                     node,
                     &err,
@@ -786,7 +786,7 @@ If you are using this API, you are probably doing something wrong."
             let worker = PagerWorker {
                 sender: sender.into(),
                 page_query,
-                statement_info: routing_info,
+                routing_info,
                 query_is_idempotent: statement.config.is_idempotent,
                 query_consistency: consistency,
                 load_balancing_policy,
@@ -914,7 +914,7 @@ If you are using this API, you are probably doing something wrong."
             let worker = PagerWorker {
                 sender: sender.into(),
                 page_query,
-                statement_info,
+                routing_info: statement_info,
                 query_is_idempotent: config.prepared.config.is_idempotent,
                 query_consistency: consistency,
                 load_balancing_policy,
diff --git a/scylla/src/client/session.rs b/scylla/src/client/session.rs
@@ -2031,7 +2031,14 @@ impl Session {
                 |_: tokio::time::error::Elapsed| {
                     #[cfg(feature = "metrics")]
                     self.metrics.inc_request_timeouts();
-                    Err(RequestError::RequestTimeout(timeout))
+
+                    let timeout_error = RequestError::RequestTimeout(timeout);
+                    trace!(
+                        parent: request_span.span(),
+                        error = %timeout_error,
+                        "Request timed out"
+                    );
+                    Err(timeout_error)
                 },
             ),
             None => runner.await,
@@ -2146,15 +2153,15 @@ impl Session {
                 };
 
                 // Use retry policy to decide what to do next
-                let query_info = RequestInfo {
+                let request_info = RequestInfo {
                     error: &request_error,
                     is_idempotent: context.is_idempotent,
                     consistency: context
                         .consistency_set_on_statement
                         .unwrap_or(execution_profile.consistency),
                 };
 
-                let retry_decision = context.retry_session.decide_should_retry(query_info);
+                let retry_decision = context.retry_session.decide_should_retry(request_info);
                 trace!(
                     parent: &span,
                     retry_decision = ?retry_decision
diff --git a/scylla/tests/integration/statements/request_timeout.rs b/scylla/tests/integration/statements/request_timeout.rs
@@ -1,75 +1,144 @@
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 
 use assert_matches::assert_matches;
 use scylla::{
-    client::execution_profile::ExecutionProfile, errors::ExecutionError, statement::Statement,
+    client::{execution_profile::ExecutionProfile, session_builder::SessionBuilder},
+    errors::ExecutionError,
+    statement::{
+        Statement,
+        batch::{Batch, BatchType},
+    },
+};
+use scylla_proxy::{
+    Condition, ProxyError, Reaction as _, RequestOpcode, RequestReaction, RequestRule, WorkerError,
 };
 
-use crate::utils::{create_new_session_builder, setup_tracing};
+use crate::utils::{setup_tracing, test_with_3_node_cluster};
 
-#[ignore = "works on remote Scylla instances only (local ones are too fast)"]
+#[cfg_attr(scylla_cloud_tests, ignore)]
 #[tokio::test]
 async fn test_request_timeout() {
     setup_tracing();
 
-    let fast_timeouting_profile_handle = ExecutionProfile::builder()
-        .request_timeout(Some(Duration::from_millis(1)))
-        .build()
-        .into_handle();
-
-    {
-        let session = create_new_session_builder().build().await.unwrap();
-
-        let mut query: Statement = Statement::new("SELECT * FROM system_schema.tables");
-        query.set_request_timeout(Some(Duration::from_millis(1)));
-        match session.query_unpaged(query, &[]).await {
-            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
-            Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+    let res = test_with_3_node_cluster(
+        scylla_proxy::ShardAwareness::QueryNode,
+        |proxy_uris, translation_map, mut running_proxy| async move {
+            let short_timeout = Duration::from_millis(1);
+            let very_long_timeout = Duration::from_secs(10000);
+            let query_str = "SELECT host_id FROM system.local WHERE key='local'";
+
+            let create_session_builder = || {
+                SessionBuilder::new()
+                    .known_node(proxy_uris[0].as_str())
+                    .address_translator(Arc::new(translation_map.clone()))
+            };
+
+            running_proxy.running_nodes.iter_mut().for_each(|node| {
+                node.change_request_rules(Some(vec![
+                    RequestRule(
+                        Condition::any([
+                            Condition::RequestOpcode(RequestOpcode::Query),
+                            Condition::RequestOpcode(RequestOpcode::Execute),
+                            Condition::RequestOpcode(RequestOpcode::Batch),
+                        ])
+                            .and(Condition::not(Condition::ConnectionRegisteredAnyEvent)),
+                        RequestReaction::delay(10 * short_timeout)
+                    )
+                ]));
+            });
+
+            let session = create_session_builder()
+                .build()
+                .await
+                .unwrap();
+
+            // Case 1: per-statement timeouts.
+            {
+                let mut query: Statement = Statement::new(query_str);
+                query.set_request_timeout(Some(short_timeout));
+                match session.query_unpaged(query, &[]).await {
+                    Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                }
+
+                let mut batch: Batch = Batch::new(BatchType::Logged);
+                batch.set_request_timeout(Some(short_timeout));
+                match session.batch(&batch, &[][..] as &[()]).await {
+                    Ok(_) => panic!("the batch should have failed due to a client-side timeout"),
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                }
+
+                let mut prepared = session
+                    .prepare(query_str)
+                    .await
+                    .unwrap();
+
+                prepared.set_request_timeout(Some(short_timeout));
+                match session.execute_unpaged(&prepared, &[]).await {
+                    Ok(_) => {
+                        panic!("the prepared query should have failed due to a client-side timeout")
+                    }
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                };
+            }
+
+            // Case 2: tight session-level timeout, overridden by per-statement long timeouts.
+            {
+                let fast_timeouting_profile = ExecutionProfile::builder()
+                    .request_timeout(Some(short_timeout))
+                    .build();
+
+                // Although this `clone()` looks suspicious, it is necessary to get an owned handle to the profile
+                // in order to call `map_to_another_profile()`, which requires a mutable reference.
+                // The profile handle itself is just a reference-counted pointer, so cloning it is cheap,
+                // and the remapping affects the session too.
+                session.get_default_execution_profile_handle().clone().map_to_another_profile(fast_timeouting_profile);
+
+                let mut query = Statement::new(query_str);
+
+                match session.query_unpaged(query.clone(), &[]).await {
+                    Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                };
+
+                query.set_request_timeout(Some(very_long_timeout));
+
+                session.query_unpaged(query, &[]).await.expect(
+                    "the query should have not failed, because no client-side timeout was specified",
+                );
+
+                let mut batch: Batch = Batch::new(BatchType::Logged);
+                match session.batch(&batch, &[][..] as &[()]).await {
+                    Ok(_) => panic!("the batch should have failed due to a client-side timeout"),
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                }
+
+                batch.set_request_timeout(Some(very_long_timeout));
+                session.batch(&batch, &[][..] as &[()]).await.expect("the batch should have not failed, because no client-side timeout was specified");
+
+                let mut prepared = session
+                    .prepare(query_str)
+                    .await
+                    .unwrap();
+
+                match session.execute_unpaged(&prepared, &[]).await {
+                    Ok(_) => {
+                        panic!("the prepared query should have failed due to a client-side timeout")
+                    }
+                    Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
+                };
+
+                prepared.set_request_timeout(Some(very_long_timeout));
+
+                session.execute_unpaged(&prepared, &[]).await.expect("the prepared query should have not failed, because no client-side timeout was specified");
+            }
+            running_proxy
         }
+    ).await;
 
-        let mut prepared = session
-            .prepare("SELECT * FROM system_schema.tables")
-            .await
-            .unwrap();
-
-        prepared.set_request_timeout(Some(Duration::from_millis(1)));
-        match session.execute_unpaged(&prepared, &[]).await {
-            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
-            Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
-        };
-    }
-    {
-        let timeouting_session = create_new_session_builder()
-            .default_execution_profile_handle(fast_timeouting_profile_handle)
-            .build()
-            .await
-            .unwrap();
-
-        let mut query = Statement::new("SELECT * FROM system_schema.tables");
-
-        match timeouting_session.query_unpaged(query.clone(), &[]).await {
-            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
-            Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
-        };
-
-        query.set_request_timeout(Some(Duration::from_secs(10000)));
-
-        timeouting_session.query_unpaged(query, &[]).await.expect(
-            "the query should have not failed, because no client-side timeout was specified",
-        );
-
-        let mut prepared = timeouting_session
-            .prepare("SELECT * FROM system_schema.tables")
-            .await
-            .unwrap();
-
-        match timeouting_session.execute_unpaged(&prepared, &[]).await {
-            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
-            Err(e) => assert_matches!(e, ExecutionError::RequestTimeout(_)),
-        };
-
-        prepared.set_request_timeout(Some(Duration::from_secs(10000)));
-
-        timeouting_session.execute_unpaged(&prepared, &[]).await.expect("the prepared query should have not failed, because no client-side timeout was specified");
+    match res {
+        Ok(()) => (),
+        Err(ProxyError::Worker(WorkerError::DriverDisconnected(_))) => (),
+        Err(err) => panic!("{}", err),
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -240,7 +240,8 @@ pub struct ExecutionProfileBuilder {`
`240`	`240`	`}`
`241`	`241`
`242`	`242`	`impl ExecutionProfileBuilder {`
`243`		`- /// Changes client-side timeout.`
	`243`	`+ /// Changes client-side timeout for executing statements.`
	`244`	`+ /// If set to None, the driver will wait indefinitely for a response from the server.`
`244`	`245`	`/// The default is 30 seconds.`
`245`	`246`	`///`
`246`	`247`	`/// # Example`