pager: encapsulate timeout logic

PageQueryTimeouter is introduced to encapsulate timeout logic for
paging queries. It was clear to me that such encapsulation in a new type
improves code readability and maintainability, as the timeout logic is
now clearly separated from the rest of the paging logic.

Author: wprzytula

scylla/src/client/pager.rs

@@ -8,7 +8,6 @@ use std::ops::ControlFlow;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
-use std::time::Duration;
 
 use futures::Stream;
 use scylla_cql::Consistency;
@@ -135,6 +134,50 @@ use crate::response::Coordinator;
 
 type PageSendAttemptedProof = SendAttemptedProof<Result<ReceivedPage, NextPageError>>;
 
+mod timeouter {
+    use std::time::Duration;
+
+    use tokio::time::Instant;
+
+    /// Encapsulation of a timeout for paging queries.
+    pub(super) struct PageQueryTimeouter {
+        timeout: Duration,
+        timeout_instant: Instant,
+    }
+
+    impl PageQueryTimeouter {
+        /// Creates a new PageQueryTimeouter with the given timeout duration,
+        /// starting from now.
+        pub(super) fn new(timeout: Duration) -> Self {
+            Self {
+                timeout,
+                timeout_instant: Instant::now() + timeout,
+            }
+        }
+
+        /// Returns the timeout duration.
+        pub(super) fn timeout_duration(&self) -> Duration {
+            self.timeout
+        }
+
+        /// Returns the instant at which the timeout will elapse.
+        ///
+        /// This can be used with `tokio::time::timeout_at`.
+        pub(super) fn deadline(&self) -> Instant {
+            self.timeout_instant
+        }
+
+        /// Resets the timeout countdown.
+        ///
+        /// This should be called right before beginning first page fetch
+        /// and after each successful page fetch.
+        pub(super) fn reset(&mut self) {
+            self.timeout_instant = Instant::now() + self.timeout;
+        }
+    }
+}
+use timeouter::PageQueryTimeouter;
+
 // PagerWorker works in the background to fetch pages
 // QueryPager receives them through a channel
 struct PagerWorker<'a, QueryFunc, SpanCreatorFunc> {
@@ -149,7 +192,7 @@ struct PagerWorker<'a, QueryFunc, SpanCreatorFunc> {
     query_is_idempotent: bool,
     query_consistency: Consistency,
     retry_session: Box<dyn RetrySession>,
-    timeout: Option<Duration>,
+    timeouter: Option<PageQueryTimeouter>,
     #[cfg(feature = "metrics")]
     metrics: Arc<Metrics>,
 
@@ -180,8 +223,7 @@ where
         let mut current_consistency: Consistency = self.query_consistency;
 
         self.log_request_start();
-        let start_instant = tokio::time::Instant::now();
-        let mut timeout_instant = self.timeout.map(|t| start_instant + t);
+        self.timeouter.as_mut().map(PageQueryTimeouter::reset);
 
         'nodes_in_plan: for (node, shard) in query_plan {
             let span = trace_span!(parent: &self.parent_span, "Executing query", node = %node.address, shard = %shard);
@@ -216,13 +258,7 @@ where
                     Result<PageSendAttemptedProof, RequestAttemptError>,
                     RequestTimeoutError,
                 > = self
-                    .query_pages(
-                        &connection,
-                        current_consistency,
-                        node,
-                        coordinator.clone(),
-                        timeout_instant.as_mut(),
-                    )
+                    .query_pages(&connection, current_consistency, node, coordinator.clone())
                     .instrument(span.clone())
                     .await;
 
@@ -326,7 +362,6 @@ where
         consistency: Consistency,
         node: NodeRef<'_>,
         coordinator: Coordinator,
-        mut timeout_instant: Option<&mut tokio::time::Instant>,
     ) -> Result<Result<PageSendAttemptedProof, RequestAttemptError>, RequestTimeoutError> {
         loop {
             let request_span = (self.span_creator)();
@@ -337,7 +372,6 @@ where
                     node,
                     coordinator.clone(),
                     &request_span,
-                    timeout_instant.as_ref().map(|instant| **instant),
                 )
                 .instrument(request_span.span().clone())
                 .await
@@ -350,11 +384,7 @@ where
                 Ok(Ok(ControlFlow::Continue(()))) => {
                     // Successfully queried one page, and there are more to fetch.
                     // Reset the timeout_instant for the next page fetch.
-                    if let Some(timeout) = self.timeout
-                        && let Some(ref mut instant) = timeout_instant
-                    {
-                        **instant = tokio::time::Instant::now() + timeout;
-                    }
+                    self.timeouter.as_mut().map(PageQueryTimeouter::reset);
                 }
                 Ok(Err(request_attempt_error)) => {
                     return Ok(Err(request_attempt_error));
@@ -373,7 +403,6 @@ where
         node: NodeRef<'_>,
         coordinator: Coordinator,
         request_span: &RequestSpan,
-        timeout_instant: Option<tokio::time::Instant>,
     ) -> Result<
         Result<ControlFlow<PageSendAttemptedProof, ()>, RequestAttemptError>,
         RequestTimeoutError,
@@ -394,20 +423,19 @@ where
                 .await
                 .and_then(QueryResponse::into_non_error_query_response)
         };
-        let query_response = match (self.timeout, timeout_instant) {
-            (Some(timeout), Some(instant)) => {
-                match tokio::time::timeout_at(instant, runner).await {
-                Ok(res) => res,
-                Err(_) /* tokio::time::error::Elapsed */ => {
-                    #[cfg(feature = "metrics")]
-                    self.metrics.inc_request_timeouts();
-                    return Err(RequestTimeoutError(timeout));
+        let query_response = match self.timeouter {
+            Some(ref timeouter) => {
+                match tokio::time::timeout_at(timeouter.deadline(), runner).await {
+                    Ok(res) => res,
+                    Err(_) /* tokio::time::error::Elapsed */ => {
+                        #[cfg(feature = "metrics")]
+                        self.metrics.inc_request_timeouts();
+                        return Err(RequestTimeoutError(timeouter.timeout_duration()));
+                    }
                 }
             }
-            }
 
-            (None, None) => runner.await,
-            _ => unreachable!("timeout_instant must be Some iff self.timeout is Some"),
+            None => runner.await,
         };
 
         let elapsed = query_start.elapsed();
@@ -802,9 +830,10 @@ If you are using this API, you are probably doing something wrong."
             .serial_consistency
             .unwrap_or(execution_profile.serial_consistency);
 
-        let timeout = statement
+        let timeouter = statement
             .get_request_timeout()
-            .or(execution_profile.request_timeout);
+            .or(execution_profile.request_timeout)
+            .map(PageQueryTimeouter::new);
 
         let page_size = statement.get_validated_page_size();
 
@@ -862,7 +891,7 @@ If you are using this API, you are probably doing something wrong."
                 query_consistency: consistency,
                 load_balancing_policy,
                 retry_session,
-                timeout,
+                timeouter,
                 #[cfg(feature = "metrics")]
                 metrics,
                 paging_state: PagingState::start(),
@@ -895,10 +924,11 @@ If you are using this API, you are probably doing something wrong."
             .serial_consistency
             .unwrap_or(config.execution_profile.serial_consistency);
 
-        let timeout = config
+        let timeouter = config
             .prepared
             .get_request_timeout()
-            .or(config.execution_profile.request_timeout);
+            .or(config.execution_profile.request_timeout)
+            .map(PageQueryTimeouter::new);
 
         let page_size = config.prepared.get_validated_page_size();
 
@@ -996,7 +1026,7 @@ If you are using this API, you are probably doing something wrong."
                 query_consistency: consistency,
                 load_balancing_policy,
                 retry_session,
-                timeout,
+                timeouter,
                 #[cfg(feature = "metrics")]
                 metrics: config.metrics,
                 paging_state: PagingState::start(),