Repository: scylladb/scylla-rust-driver
Pull request #522 (2022-08-19 09:54:45)

Implemented client-side per-session and per-statement configurable timeouts
Two levels of client-side timeouts were added: per Session and per Statement (PreparedStatement and Query). If no timeout is specified for a statement, a global per-session one is used. Once a timeout elapses, an `RequestTimeout` error is returned.

Fixes: #304

## Pre-review checklist
<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->
- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


wprzytula: @psarna @cvybhu Maybe you could help me with understanding why `test_client_timeout` fails in nondeterministic way.
cvybhu: > @psarna @cvybhu Maybe you could help me with understanding why `test_client_timeout` fails in nondeterministic way.

`tokio::Timeout` checks if the future has completed before checking if the timeout has elapsed.
With a local node and fast network it's possible that the query completes in a single `Poll` maybe?

https://github.com/tokio-rs/tokio/blob/de81985762a242c77361a6ab9de198372ca85987/tokio/src/time/timeout.rs#L175
cvybhu: Or maybe not, I added `tokio::task::yield_now().await;` at the beginning of runner block and it still fails.
So there are at least two Polls of the timeout future and it doesn't fail 0_o
cvybhu: Yeah something like this fails as well:
```rust
async fn yields_once() {
    tokio::task::yield_now().await;
}

async fn instant_timeout_test() {
    let res = tokio::time::timeout(std::time::Duration::from_secs(0), yields_once()).await;
    assert!(res.is_err());
}
```

Although sleeping for 1ms is enough to trigger timeout.

This might be an issue with tokio, but I'm not sure if it's a bug or a feature.
Tokio time keeps mentioning that its precision can be > 1ms, for example for `sleep`:

> The implementation is platform specific, and some platforms (specifically Windows) will provide timers with a larger resolution than 1 ms.

https://docs.rs/tokio/latest/tokio/time/fn.sleep.html
cvybhu: > Tokio time keeps mentioning that its precision can be > 1ms, for example for sleep:

I think we can just assume that tokio doesn't support timeouts of < 1ms and wait with the tests until proxy is merged.
wprzytula: > > Tokio time keeps mentioning that its precision can be > 1ms, for example for sleep:
> 
> I think we can just assume that tokio doesn't support timeouts of < 1ms and wait with the tests until proxy is merged.

That's right, with proxy we will easily delay the response for a couple of millis and set the timeout for 1 ms - this should work.
psarna: @wprzytula for now, let's just verify manually. If you create a 0ms or 1ms timeout, it should properly time out when sent to Scylla booted on a remote machine (e.g. one of the office servers).

Patch:
diff --git a/.github/workflows/docs-pages.yaml b/.github/workflows/docs-pages.yaml
index 9a6c1520..0b6c366a 100644
--- a/.github/workflows/docs-pages.yaml
+++ b/.github/workflows/docs-pages.yaml
@@ -27,8 +27,6 @@ jobs:
         run: make -C docs setupenv
       - name: Build docs
         run: make -C docs multiversion
-      - name: Build redirects
-        run: make -C docs redirects
       - name: Deploy docs to GitHub Pages
         run: ./docs/_utils/deploy.sh
         env:
diff --git a/docs/Makefile b/docs/Makefile
index b612ad95..eb22febb 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -71,12 +71,6 @@ multiversion: setup
 	@echo
 	@echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
 
-.PHONY: redirects
-redirects: setup
-	$(POETRY) run redirects-cli fromfile --yaml-file _utils/redirects.yaml --output-dir $(BUILDDIR)/dirhtml
-	@echo
-	@echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
-
 # Preview commands
 .PHONY: preview
 preview: setup
diff --git a/docs/_utils/redirections.yaml b/docs/_utils/redirections.yaml
new file mode 100644
index 00000000..7692003b
--- /dev/null
+++ b/docs/_utils/redirections.yaml
@@ -0,0 +1,6 @@
+### a dictionary of redirections
+#old path: new path
+#
+
+# removing redirection html script files
+# test: /
diff --git a/docs/_utils/redirects.yaml b/docs/_utils/redirects.yaml
deleted file mode 100644
index e69de29b..00000000
diff --git a/docs/pyproject.toml b/docs/pyproject.toml
index c6678dae..cf4124a8 100644
--- a/docs/pyproject.toml
+++ b/docs/pyproject.toml
@@ -9,8 +9,7 @@ python = "^3.7"
 pyyaml = "6.0"
 pygments = "2.11.2"
 recommonmark = "0.7.1"
-redirects_cli ="~0.1.2"
-sphinx-scylladb-theme = "~1.3.1"
+sphinx-scylladb-theme = "~1.2.1"
 sphinx-sitemap = "2.2.0"
 sphinx-autobuild = "2021.3.14"
 Sphinx = "4.3.2"
diff --git a/docs/source/conf.py b/docs/source/conf.py
index d8f258c6..a389e5aa 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -3,6 +3,7 @@ import os
 import sys
 from datetime import date
 
+import recommonmark
 from recommonmark.transform import AutoStructify
 from pygments.lexers.configs import TOMLLexer
 from pygments.lexers.rust import RustLexer
@@ -19,9 +20,9 @@ BRANCHES = ['main']
 # Set the latest version.
 LATEST_VERSION = 'main'
 # Set which versions are not released yet.
-UNSTABLE_VERSIONS = []
+UNSTABLE_VERSIONS = ['']
 # Set which versions are deprecated
-DEPRECATED_VERSIONS = []
+DEPRECATED_VERSIONS = ['']
 
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
@@ -79,6 +80,11 @@ notfound_template =  '404.html'
 # Prefix added to all the URLs generated in the 404 page.
 notfound_urls_prefix = ''
 
+# -- Options for redirect extension ---------------------------------------
+
+# Read a YAML dictionary of redirections and generate an HTML file for each
+redirects_file = '_utils/redirections.yaml'
+
 # -- Options for multiversion extension ----------------------------------
 
 # Whitelist pattern for tags (set to None to ignore all tags)
diff --git a/scylla-cql/Cargo.toml b/scylla-cql/Cargo.toml
index adb66ffa..734b4eab 100644
--- a/scylla-cql/Cargo.toml
+++ b/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"
diff --git a/scylla-cql/src/errors.rs b/scylla-cql/src/errors.rs
index 175b66eb..6384166d 100644
--- a/scylla-cql/src/errors.rs
+++ b/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),
         }
     }
 }
diff --git a/scylla-cql/src/frame/response/mod.rs b/scylla-cql/src/frame/response/mod.rs
index e53ae874..e9fee538 100644
--- a/scylla-cql/src/frame/response/mod.rs
+++ b/scylla-cql/src/frame/response/mod.rs
@@ -5,7 +5,7 @@ pub mod event;
 pub mod result;
 pub mod supported;
 
-use crate::{errors::QueryError, frame::frame_errors::ParseError};
+use crate::frame::frame_errors::ParseError;
 use num_enum::TryFromPrimitive;
 
 pub use error::Error;
@@ -57,29 +57,4 @@ impl Response {
 
         Ok(response)
     }
-
-    pub fn into_non_error_response(self) -> Result<NonErrorResponse, QueryError> {
-        Ok(match self {
-            Response::Error(err) => return Err(QueryError::from(err)),
-            Response::Ready => NonErrorResponse::Ready,
-            Response::Result(res) => NonErrorResponse::Result(res),
-            Response::Authenticate(auth) => NonErrorResponse::Authenticate(auth),
-            Response::AuthSuccess(auth_succ) => NonErrorResponse::AuthSuccess(auth_succ),
-            Response::AuthChallenge(auth_chal) => NonErrorResponse::AuthChallenge(auth_chal),
-            Response::Supported(sup) => NonErrorResponse::Supported(sup),
-            Response::Event(eve) => NonErrorResponse::Event(eve),
-        })
-    }
-}
-
-// A Response which can not be Response::Error
-#[derive(Debug)]
-pub enum NonErrorResponse {
-    Ready,
-    Result(result::Result),
-    Authenticate(authenticate::Authenticate),
-    AuthSuccess(authenticate::AuthSuccess),
-    AuthChallenge(authenticate::AuthChallenge),
-    Supported(Supported),
-    Event(event::Event),
 }
diff --git a/scylla/Cargo.toml b/scylla/Cargo.toml
index 52fcf7b8..c98b9dcc 100644
--- a/scylla/Cargo.toml
+++ b/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"
diff --git a/scylla/src/statement/mod.rs b/scylla/src/statement/mod.rs
index a5eae61e..bd8513bf 100644
--- a/scylla/src/statement/mod.rs
+++ b/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,6 +34,7 @@ impl Default for StatementConfig {
             speculative_execution_policy: None,
             tracing: false,
             timestamp: None,
+            request_timeout: None,
         }
     }
 }
@@ -40,16 +42,12 @@ impl Default for StatementConfig {
 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
         }
     }
 }
diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs
index d28c8071..d51a70c9 100644
--- a/scylla/src/statement/prepared_statement.rs
+++ b/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 {
diff --git a/scylla/src/statement/query.rs b/scylla/src/statement/query.rs
index 0252b7c7..59346072 100644
--- a/scylla/src/statement/query.rs
+++ b/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 {
diff --git a/scylla/src/tracing.rs b/scylla/src/tracing.rs
index dbdff599..63822b23 100644
--- a/scylla/src/tracing.rs
+++ b/scylla/src/tracing.rs
@@ -1,5 +1,4 @@
 use crate::statement::Consistency;
-use itertools::Itertools;
 use std::collections::HashMap;
 use std::net::IpAddr;
 use std::num::NonZeroU32;
@@ -49,17 +48,6 @@ pub struct GetTracingConfig {
     pub consistency: Consistency,
 }
 
-impl TracingInfo {
-    /// Returns a list of unique nodes involved in the query
-    pub fn nodes(&self) -> Vec<IpAddr> {
-        self.events
-            .iter()
-            .filter_map(|e| e.source)
-            .unique()
-            .collect()
-    }
-}
-
 impl Default for GetTracingConfig {
     fn default() -> GetTracingConfig {
         GetTracingConfig {
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 0f5a3e0a..66ace683 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -32,7 +32,7 @@ use crate::batch::{Batch, BatchStatement};
 use crate::frame::{
     self,
     request::{self, batch, execute, query, register, Request},
-    response::{event::Event, result, NonErrorResponse, Response, ResponseOpcode},
+    response::{event::Event, result, Response, ResponseOpcode},
     server_event_type::EventType,
     value::{BatchValues, ValueList},
     FrameParams, SerializedRequest,
@@ -142,13 +142,6 @@ pub struct QueryResponse {
     pub warnings: Vec<String>,
 }
 
-// A QueryResponse in which response can not be Response::Error
-pub struct NonErrorQueryResponse {
-    pub response: NonErrorResponse,
-    pub tracing_id: Option<Uuid>,
-    pub warnings: Vec<String>,
-}
-
 /// Result of Session::batch(). Contains no rows, only some useful information.
 pub struct BatchResult {
     /// Warnings returned by the database
@@ -158,42 +151,29 @@ pub struct BatchResult {
 }
 
 impl QueryResponse {
-    pub fn into_non_error_query_response(self) -> Result<NonErrorQueryResponse, QueryError> {
-        Ok(NonErrorQueryResponse {
-            response: self.response.into_non_error_response()?,
-            tracing_id: self.tracing_id,
-            warnings: self.warnings,
-        })
-    }
-
-    pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
-        self.into_non_error_query_response()?.into_query_result()
-    }
-}
-
-impl NonErrorQueryResponse {
     pub fn as_set_keyspace(&self) -> Option<&result::SetKeyspace> {
         match &self.response {
-            NonErrorResponse::Result(result::Result::SetKeyspace(sk)) => Some(sk),
+            Response::Result(result::Result::SetKeyspace(sk)) => Some(sk),
             _ => None,
         }
     }
 
     pub fn as_schema_change(&self) -> Option<&result::SchemaChange> {
         match &self.response {
-            NonErrorResponse::Result(result::Result::SchemaChange(sc)) => Some(sc),
+            Response::Result(result::Result::SchemaChange(sc)) => Some(sc),
             _ => None,
         }
     }
 
     pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
         let (rows, paging_state, col_specs) = match self.response {
-            NonErrorResponse::Result(result::Result::Rows(rs)) => (
+            Response::Error(err) => return Err(err.into()),
+            Response::Result(result::Result::Rows(rs)) => (
                 Some(rs.rows),
                 rs.metadata.paging_state,
                 rs.metadata.col_specs,
             ),
-            NonErrorResponse::Result(_) => (None, None, vec![]),
+            Response::Result(_) => (None, None, vec![]),
             _ => {
                 return Err(QueryError::ProtocolError(
                     "Unexpected server response, expected Result or Error",
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 6b65d5bb..a1f030b9 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -14,7 +14,6 @@ use tokio::time::timeout;
 use tracing::{debug, error, trace, trace_span, Instrument};
 use uuid::Uuid;
 
-use super::connection::NonErrorQueryResponse;
 use super::connection::QueryResponse;
 use super::errors::{BadQuery, NewSessionError, QueryError};
 use crate::cql_to_rust::FromRow;
@@ -64,6 +63,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 +121,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 +142,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 +181,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 +380,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 {
@@ -471,7 +477,6 @@ impl Session {
                         connection
                             .query(query_ref, values_ref, paging_state_ref.clone())
                             .await
-                            .and_then(QueryResponse::into_non_error_query_response)
                     }
                 },
             )
@@ -486,7 +491,7 @@ impl Session {
 
     async fn handle_set_keyspace_response(
         &self,
-        response: &NonErrorQueryResponse,
+        response: &QueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(set_keyspace) = response.as_set_keyspace() {
             debug!(
@@ -503,7 +508,7 @@ impl Session {
     async fn handle_auto_await_schema_agreement(
         &self,
         contents: &str,
-        response: &NonErrorQueryResponse,
+        response: &QueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(timeout) = self.auto_await_schema_agreement_timeout {
             if response.as_schema_change().is_some()
@@ -750,7 +755,7 @@ impl Session {
             "Request",
             prepared_id = format!("{:X}", prepared.get_id()).as_str()
         );
-        let response: NonErrorQueryResponse = self
+        let response = self
             .run_query(
                 statement_info,
                 &prepared.config,
@@ -764,7 +769,6 @@ impl Session {
                     connection
                         .execute(prepared, values_ref, paging_state_ref.clone())
                         .await
-                        .and_then(QueryResponse::into_non_error_query_response)
                 },
             )
             .instrument(span)
@@ -1107,84 +1111,99 @@ impl Session {
     where
         ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
         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,
         }
     }
 
@@ -1200,7 +1219,6 @@ impl Session {
     where
         ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
         QueryFut: Future<Output = Result<ResT, QueryError>>,
-        ResT: AllowedRunQueryResTType,
     {
         let mut last_error: Option<QueryError> = None;
 
@@ -1308,7 +1326,6 @@ impl Session {
     ) -> Result<ResT, QueryError>
     where
         QueryFut: Future<Output = Result<ResT, QueryError>>,
-        ResT: AllowedRunQueryResTType,
     {
         let info = Statement::default();
         let config = StatementConfig {
@@ -1401,18 +1418,3 @@ async fn resolve_hostname(hostname: &str) -> Result<SocketAddr, NewSessionError>
 
     ret.ok_or(failed_err)
 }
-
-// run_query, execute_query, etc have a template type called ResT.
-// There was a bug where ResT was set to QueryResponse, which could
-// be an error response. This was not caught by retry policy which
-// assumed all errors would come from analyzing Result<ResT, QueryError>.
-// This trait is a guard to make sure that this mistake doesn't
-// happen again.
-// When using run_query make sure that the ResT type is NOT able
-// to contain any errors.
-// See https://github.com/scylladb/scylla-rust-driver/issues/501
-pub trait AllowedRunQueryResTType {}
-
-impl AllowedRunQueryResTType for Uuid {}
-impl AllowedRunQueryResTType for BatchResult {}
-impl AllowedRunQueryResTType for NonErrorQueryResponse {}
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index e711e0cc..ff5f5f1d 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -541,6 +541,27 @@ impl SessionBuilder {
         self.config.auto_await_schema_agreement_timeout = None;
         self
     }
+
+    /// Changes client-side timeout
+    /// The default is 30 seconds.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use std::time::Duration;
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .request_timeout(Some(Duration::from_millis(500)))
+    ///     .build() // Turns SessionBuilder into Session
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn request_timeout(mut self, duration: Option<std::time::Duration>) -> Self {
+        self.config.request_timeout = duration;
+        self
+    }
 }
 
 /// Creates a [`SessionBuilder`] with default configuration, same as [`SessionBuilder::new`]
@@ -707,6 +728,21 @@ mod tests {
         assert!(builder.config.fetch_schema_metadata);
     }
 
+    #[test]
+    fn request_timeout() {
+        let mut builder = SessionBuilder::new();
+        assert_eq!(
+            builder.config.request_timeout,
+            Some(std::time::Duration::from_secs(30))
+        );
+
+        builder = builder.request_timeout(Some(std::time::Duration::from_secs(10)));
+        assert_eq!(
+            builder.config.request_timeout,
+            Some(std::time::Duration::from_secs(10))
+        );
+    }
+
     #[test]
     fn all_features() {
         let mut builder = SessionBuilder::new();


Test patch:
diff --git a/.github/workflows/authenticate_test.yml b/.github/workflows/authenticate_test.yml
index 620f7e9e..5f2b271d 100644
--- a/.github/workflows/authenticate_test.yml
+++ b/.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
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index c875dd97..f95d8401 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -14,6 +14,7 @@ use crate::transport::topology::{CollectionType, ColumnKind, CqlType, NativeType
 use crate::CachingSession;
 use crate::QueryResult;
 use crate::{IntoTypedRows, Session, SessionBuilder};
+use assert_matches::assert_matches;
 use bytes::Bytes;
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
@@ -919,7 +920,6 @@ async fn test_get_tracing_info(session: &Session, ks: String) {
     // Getting tracing info from session using this uuid works
     let tracing_info: TracingInfo = session.get_tracing_info(&tracing_id).await.unwrap();
     assert!(!tracing_info.events.is_empty());
-    assert!(!tracing_info.nodes().is_empty());
 }
 
 async fn test_tracing_query_iter(session: &Session, ks: String) {
@@ -1182,6 +1182,74 @@ async fn test_timestamp() {
     assert_eq!(results, expected_results);
 }
 
+#[ignore = "works on remote Scylla instances only (local ones are too fast)"]
+#[tokio::test]
+async fn test_request_timeout() {
+    use std::time::Duration;
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+
+    {
+        let session = SessionBuilder::new()
+            .known_node(uri.as_str())
+            .build()
+            .await
+            .unwrap();
+
+        let mut query: Query = Query::new("SELECT * FROM system_schema.tables");
+        query.set_request_timeout(Some(Duration::from_millis(1)));
+        match session.query(query, &[]).await {
+            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        }
+
+        let mut prepared = session
+            .prepare("SELECT * FROM system_schema.tables")
+            .await
+            .unwrap();
+
+        prepared.set_request_timeout(Some(Duration::from_millis(1)));
+        match session.execute(&prepared, &[]).await {
+            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+    }
+    {
+        let timeouting_session = SessionBuilder::new()
+            .known_node(uri)
+            .request_timeout(Some(Duration::from_millis(1)))
+            .build()
+            .await
+            .unwrap();
+
+        let mut query = Query::new("SELECT * FROM system_schema.tables");
+
+        match timeouting_session.query(query.clone(), &[]).await {
+            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+
+        query.set_request_timeout(Some(Duration::from_secs(10000)));
+
+        timeouting_session.query(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(&prepared, &[]).await {
+            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+
+        prepared.set_request_timeout(Some(Duration::from_secs(10000)));
+
+        timeouting_session.execute(&prepared, &[]).await.expect("the prepared query should have not failed, because no client-side timeout was specified");
+    }
+}
+
 #[tokio::test]
 async fn test_prepared_config() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());