Merge pull request #522 from wprzytula/timeouts

Implemented client-side per-session and per-statement configurable timeouts

Author: psarna

.github/workflows/authenticate_test.yml

@@ -24,4 +24,4 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Run tests
-      run: cargo test --verbose --  --ignored
+      run: cargo test --verbose authenticate_superuser --  --ignored

scylla-cql/Cargo.toml

@@ -14,7 +14,7 @@ scylla-macros = { version = "0.1.1", path = "../scylla-macros"}
 byteorder = "1.3.4"
 bytes = "1.0.1"
 num_enum = "0.5"
-tokio = { version = "1.12", features = ["io-util"] }
+tokio = { version = "1.12", features = ["io-util", "time"] }
 snap = "1.0"
 uuid = "1.0"
 thiserror = "1.0"

scylla-cql/src/errors.rs

@@ -40,6 +40,10 @@ pub enum QueryError {
 
     #[error("Unable to allocate stream id")]
     UnableToAllocStreamId,
+
+    /// Client timeout occurred before any response arrived
+    #[error("Request timeout: {0}")]
+    RequestTimeout(String),
 }
 
 /// An error sent from the database in response to a query
@@ -292,6 +296,11 @@ pub enum NewSessionError {
 
     #[error("Unable to allocate stream id")]
     UnableToAllocStreamId,
+
+    /// Client timeout occurred before a response arrived for some query
+    /// during `Session` creation.
+    #[error("Client timeout: {0}")]
+    RequestTimeout(String),
 }
 
 /// Invalid keyspace name given to `Session::use_keyspace()`
@@ -340,6 +349,12 @@ impl From<FrameError> for QueryError {
     }
 }
 
+impl From<tokio::time::error::Elapsed> for QueryError {
+    fn from(timer_error: tokio::time::error::Elapsed) -> QueryError {
+        QueryError::RequestTimeout(format!("{}", timer_error))
+    }
+}
+
 impl From<std::io::Error> for NewSessionError {
     fn from(io_error: std::io::Error) -> NewSessionError {
         NewSessionError::IoError(Arc::new(io_error))
@@ -359,6 +374,7 @@ impl From<QueryError> for NewSessionError {
                 NewSessionError::TooManyOrphanedStreamIds(ids)
             }
             QueryError::UnableToAllocStreamId => NewSessionError::UnableToAllocStreamId,
+            QueryError::RequestTimeout(msg) => NewSessionError::RequestTimeout(msg),
         }
     }
 }

scylla/Cargo.toml

@@ -47,6 +47,7 @@ smallvec = "1.8.0"
 [dev-dependencies]
 criterion = "0.3"
 tracing-subscriber = "0.3.14"
+assert_matches = "1.5.0"
 
 [[bench]]
 name = "benchmark"

scylla/src/statement/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use crate::transport::retry_policy::RetryPolicy;
 use crate::transport::speculative_execution::SpeculativeExecutionPolicy;
@@ -21,6 +21,7 @@ pub struct StatementConfig {
 
     pub tracing: bool,
     pub timestamp: Option<i64>,
+    pub request_timeout: Option<Duration>,
 }
 
 impl Default for StatementConfig {
@@ -33,23 +34,20 @@ impl Default for StatementConfig {
             speculative_execution_policy: None,
             tracing: false,
             timestamp: None,
+            request_timeout: None,
         }
     }
 }
 
 impl Clone for StatementConfig {
     fn clone(&self) -> Self {
         Self {
-            consistency: self.consistency,
-            serial_consistency: self.serial_consistency,
-            is_idempotent: self.is_idempotent,
             retry_policy: self
                 .retry_policy
                 .as_ref()
                 .map(|policy| policy.clone_boxed()),
             speculative_execution_policy: self.speculative_execution_policy.clone(),
-            tracing: self.tracing,
-            timestamp: self.timestamp,
+            ..*self
         }
     }
 }

scylla/src/statement/prepared_statement.rs

@@ -1,6 +1,7 @@
 use bytes::{BufMut, Bytes, BytesMut};
 use smallvec::{smallvec, SmallVec};
 use std::convert::TryInto;
+use std::time::Duration;
 use thiserror::Error;
 use uuid::Uuid;
 
@@ -247,6 +248,19 @@ impl PreparedStatement {
         self.config.timestamp
     }
 
+    /// Sets the client-side timeout for this statement.
+    /// If not None, the driver will stop waiting for the request
+    /// to finish after `timeout` passed.
+    /// Otherwise, default session client timeout will be applied.
+    pub fn set_request_timeout(&mut self, timeout: Option<Duration>) {
+        self.config.request_timeout = timeout
+    }
+
+    /// Gets client timeout associated with this query
+    pub fn get_request_timeout(&self) -> Option<Duration> {
+        self.config.request_timeout
+    }
+
     /// Sets the name of the partitioner used for this statement.
     pub(crate) fn set_partitioner_name(&mut self, partitioner_name: Option<&str>) {
         self.partitioner_name = match partitioner_name {

scylla/src/statement/query.rs

@@ -1,6 +1,7 @@
 use super::StatementConfig;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::transport::retry_policy::RetryPolicy;
+use std::time::Duration;
 
 /// CQL query statement.
 ///
@@ -117,6 +118,19 @@ impl Query {
     pub fn get_timestamp(&self) -> Option<i64> {
         self.config.timestamp
     }
+
+    /// Sets the client-side timeout for this statement.
+    /// If not None, the driver will stop waiting for the request
+    /// to finish after `timeout` passed.
+    /// Otherwise, default session client timeout will be applied.
+    pub fn set_request_timeout(&mut self, timeout: Option<Duration>) {
+        self.config.request_timeout = timeout
+    }
+
+    /// Gets client timeout associated with this query
+    pub fn get_request_timeout(&self) -> Option<Duration> {
+        self.config.request_timeout
+    }
 }
 
 impl From<String> for Query {

scylla/src/transport/session.rs

@@ -64,6 +64,7 @@ pub struct Session {
     metrics: Arc<Metrics>,
     default_consistency: Consistency,
     auto_await_schema_agreement_timeout: Option<Duration>,
+    request_timeout: Option<Duration>,
 }
 
 /// This implementation deliberately omits some details from Cluster in order
@@ -121,7 +122,7 @@ pub struct SessionConfig {
     pub auth_password: Option<String>,
 
     pub schema_agreement_interval: Duration,
-    pub connect_timeout: std::time::Duration,
+    pub connect_timeout: Duration,
 
     /// Size of the per-node connection pool, i.e. how many connections the driver should keep to each node.
     /// The default is `PerShard(1)`, which is the recommended setting for Scylla clusters.
@@ -142,6 +143,10 @@ pub struct SessionConfig {
     /// Controls the timeout for the automatic wait for schema agreement after sending a schema-altering statement.
     /// If `None`, the automatic schema agreement is disabled.
     pub auto_await_schema_agreement_timeout: Option<Duration>,
+
+    /// Controls the client-side timeout for queries.
+    /// If `None`, the queries have no timeout (the driver will block indefinitely).
+    pub request_timeout: Option<Duration>,
 }
 
 /// Describes database server known on Session startup.
@@ -177,13 +182,14 @@ impl SessionConfig {
             ssl_context: None,
             auth_username: None,
             auth_password: None,
-            connect_timeout: std::time::Duration::from_secs(5),
+            connect_timeout: Duration::from_secs(5),
             connection_pool_size: Default::default(),
             disallow_shard_aware_port: false,
             default_consistency: Consistency::LocalQuorum,
             fetch_schema_metadata: true,
             keepalive_interval: None,
             auto_await_schema_agreement_timeout: Some(std::time::Duration::from_secs(60)),
+            request_timeout: Some(Duration::from_secs(30)),
         }
     }
 
@@ -375,6 +381,7 @@ impl Session {
             metrics: Arc::new(Metrics::new()),
             default_consistency: config.default_consistency,
             auto_await_schema_agreement_timeout: config.auto_await_schema_agreement_timeout,
+            request_timeout: config.request_timeout,
         };
 
         if let Some(keyspace_name) = config.used_keyspace {
@@ -1109,82 +1116,98 @@ impl Session {
         QueryFut: Future<Output = Result<ResT, QueryError>>,
         ResT: AllowedRunQueryResTType,
     {
-        let cluster_data = self.cluster.get_data();
-        let query_plan = self.load_balancer.plan(&statement_info, &cluster_data);
-
-        // If a speculative execution policy is used to run query, query_plan has to be shared
-        // between different async functions. This struct helps to wrap query_plan in mutex so it
-        // can be shared safely.
-        struct SharedPlan<I>
-        where
-            I: Iterator<Item = Arc<Node>>,
-        {
-            iter: std::sync::Mutex<I>,
-        }
-
-        impl<I> Iterator for &SharedPlan<I>
-        where
-            I: Iterator<Item = Arc<Node>>,
-        {
-            type Item = Arc<Node>;
-
-            fn next(&mut self) -> Option<Self::Item> {
-                self.iter.lock().unwrap().next()
+        let runner = async {
+            let cluster_data = self.cluster.get_data();
+            let query_plan = self.load_balancer.plan(&statement_info, &cluster_data);
+
+            // If a speculative execution policy is used to run query, query_plan has to be shared
+            // between different async functions. This struct helps to wrap query_plan in mutex so it
+            // can be shared safely.
+            struct SharedPlan<I>
+            where
+                I: Iterator<Item = Arc<Node>>,
+            {
+                iter: std::sync::Mutex<I>,
             }
-        }
 
-        let retry_policy = match &statement_config.retry_policy {
-            Some(policy) => policy,
-            None => &self.retry_policy,
-        };
+            impl<I> Iterator for &SharedPlan<I>
+            where
+                I: Iterator<Item = Arc<Node>>,
+            {
+                type Item = Arc<Node>;
 
-        #[allow(clippy::unnecessary_lazy_evaluations)]
-        let speculative_policy = statement_config
-            .speculative_execution_policy
-            .as_ref()
-            .or_else(|| self.speculative_execution_policy.as_ref());
+                fn next(&mut self) -> Option<Self::Item> {
+                    self.iter.lock().unwrap().next()
+                }
+            }
 
-        match speculative_policy {
-            Some(speculative) if statement_config.is_idempotent => {
-                let shared_query_plan = SharedPlan {
-                    iter: std::sync::Mutex::new(query_plan),
-                };
+            let retry_policy = match &statement_config.retry_policy {
+                Some(policy) => policy,
+                None => &self.retry_policy,
+            };
 
-                let execute_query_generator = || {
-                    self.execute_query(
-                        &shared_query_plan,
+            #[allow(clippy::unnecessary_lazy_evaluations)]
+            let speculative_policy = statement_config
+                .speculative_execution_policy
+                .as_ref()
+                .or_else(|| self.speculative_execution_policy.as_ref());
+
+            match speculative_policy {
+                Some(speculative) if statement_config.is_idempotent => {
+                    let shared_query_plan = SharedPlan {
+                        iter: std::sync::Mutex::new(query_plan),
+                    };
+
+                    let execute_query_generator = || {
+                        self.execute_query(
+                            &shared_query_plan,
+                            statement_config.is_idempotent,
+                            statement_config.consistency,
+                            retry_policy.new_session(),
+                            &choose_connection,
+                            &do_query,
+                        )
+                    };
+
+                    let context = speculative_execution::Context {
+                        metrics: self.metrics.clone(),
+                    };
+
+                    speculative_execution::execute(
+                        speculative.as_ref(),
+                        &context,
+                        execute_query_generator,
+                    )
+                    .await
+                }
+                _ => self
+                    .execute_query(
+                        query_plan,
                         statement_config.is_idempotent,
                         statement_config.consistency,
                         retry_policy.new_session(),
                         &choose_connection,
                         &do_query,
                     )
-                };
-
-                let context = speculative_execution::Context {
-                    metrics: self.metrics.clone(),
-                };
-
-                speculative_execution::execute(
-                    speculative.as_ref(),
-                    &context,
-                    execute_query_generator,
-                )
-                .await
+                    .await
+                    .unwrap_or(Err(QueryError::ProtocolError(
+                        "Empty query plan - driver bug!",
+                    ))),
             }
-            _ => self
-                .execute_query(
-                    query_plan,
-                    statement_config.is_idempotent,
-                    statement_config.consistency,
-                    retry_policy.new_session(),
-                    &choose_connection,
-                    &do_query,
-                )
+        };
+
+        let effective_timeout = statement_config.request_timeout.or(self.request_timeout);
+        match effective_timeout {
+            Some(timeout) => tokio::time::timeout(timeout, runner)
                 .await
-                .unwrap_or(Err(QueryError::ProtocolError(
-                    "Empty query plan - driver bug!",
-                ))),
+                .unwrap_or_else(|e| {
+                    Err(QueryError::RequestTimeout(format!(
+                        "Request took longer than {}ms: {}",
+                        timeout.as_millis(),
+                        e
+                    )))
+                }),
+            None => runner.await,
         }
     }