session: made run_query() subject of timeouts
The whole run_query() was made a one big async block, which is now
polled together with a timeout using tokio timeout. Once a timeout
elapses, the query is cancelled and QueryError::RequestTimeout
is returned.
@@ -63,7 +63,6 @@ pub struct Session {
metrics: Arc<Metrics>,
default_consistency: Consistency,
auto_await_schema_agreement_timeout: Option<Duration>,
- #[allow(dead_code)]
request_timeout: Option<Duration>,
}
@@ -1113,6 +1112,7 @@ impl Session {
ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
QueryFut: Future<Output = Result<ResT, QueryError>>,
{
+ let runner = async {
let cluster_data = self.cluster.get_data();
let query_plan = self.load_balancer.plan(&statement_info, &cluster_data);
@@ -1190,6 +1190,21 @@ impl Session {
"Empty query plan - driver bug!",
))),
}
+ };
+
+ let effective_timeout = statement_config.request_timeout.or(self.request_timeout);
+ match effective_timeout {
+ Some(timeout) => tokio::time::timeout(timeout, runner)
+ .await
+ .unwrap_or_else(|e| {
+ Err(QueryError::RequestTimeout(format!(
+ "Request took longer than {}ms: {}",
+ timeout.as_millis(),
+ e
+ )))
+ }),
+ None => runner.await,
+ }
}
async fn execute_query<ConnFut, QueryFut, ResT>(
