Extended RetryDecision with CL for retries

From now on, RetryPolicies can decide not only about whether the retry
should be attempted, but also about the consistency level of such
an attempt.

Author: wprzytula

scylla/src/transport/connection.rs

@@ -378,23 +378,55 @@ impl Connection {
         values: impl ValueList,
     ) -> Result<QueryResult, QueryError> {
         let query: Query = query.into();
-        self.query(&query, &values, None).await?.into_query_result()
+        let consistency = query
+            .config
+            .determine_consistency(self.config.default_consistency);
+        self.query_single_page_with_consistency(query, &values, consistency)
+            .await
+    }
+
+    pub async fn query_single_page_with_consistency(
+        &self,
+        query: impl Into<Query>,
+        values: impl ValueList,
+        consistency: Consistency,
+    ) -> Result<QueryResult, QueryError> {
+        let query: Query = query.into();
+        self.query_with_consistency(&query, &values, consistency, None)
+            .await?
+            .into_query_result()
     }
 
     pub async fn query(
         &self,
         query: &Query,
         values: impl ValueList,
         paging_state: Option<Bytes>,
+    ) -> Result<QueryResponse, QueryError> {
+        self.query_with_consistency(
+            query,
+            values,
+            query
+                .config
+                .determine_consistency(self.config.default_consistency),
+            paging_state,
+        )
+        .await
+    }
+
+    pub async fn query_with_consistency(
+        &self,
+        query: &Query,
+        values: impl ValueList,
+        consistency: Consistency,
+        paging_state: Option<Bytes>,
     ) -> Result<QueryResponse, QueryError> {
         let serialized_values = values.serialized()?;
 
         let query_frame = query::Query {
             contents: &query.contents,
             parameters: query::QueryParameters {
-                consistency: query
-                    .config
-                    .determine_consistency(self.config.default_consistency),
+                consistency,
                 serial_consistency: query.get_serial_consistency(),
                 values: &serialized_values,
                 page_size: query.get_page_size(),
@@ -412,6 +444,22 @@ impl Connection {
         &self,
         query: &Query,
         values: impl ValueList,
+    ) -> Result<QueryResult, QueryError> {
+        self.query_all_with_consistency(
+            query,
+            values,
+            query
+                .config
+                .determine_consistency(self.config.default_consistency),
+        )
+        .await
+    }
+
+    pub async fn query_all_with_consistency(
+        &self,
+        query: &Query,
+        values: impl ValueList,
+        consistency: Consistency,
     ) -> Result<QueryResult, QueryError> {
         if query.get_page_size().is_none() {
             // Page size should be set when someone wants to use paging
@@ -426,10 +474,14 @@ impl Connection {
         let serialized_values = values.serialized()?;
         let mut paging_state: Option<Bytes> = None;
 
+        query
+            .config
+            .determine_consistency(self.config.default_consistency);
+
         loop {
             // Send next paged query
             let mut cur_result: QueryResult = self
-                .query(query, &serialized_values, paging_state)
+                .query_with_consistency(query, &serialized_values, consistency, paging_state)
                 .await?
                 .into_query_result()?;
 
@@ -462,15 +514,31 @@ impl Connection {
         prepared_statement: &PreparedStatement,
         values: impl ValueList,
         paging_state: Option<Bytes>,
+    ) -> Result<QueryResponse, QueryError> {
+        self.execute_with_consistency(
+            prepared_statement,
+            values,
+            prepared_statement
+                .config
+                .determine_consistency(self.config.default_consistency),
+            paging_state,
+        )
+        .await
+    }
+
+    pub async fn execute_with_consistency(
+        &self,
+        prepared_statement: &PreparedStatement,
+        values: impl ValueList,
+        consistency: Consistency,
+        paging_state: Option<Bytes>,
     ) -> Result<QueryResponse, QueryError> {
         let serialized_values = values.serialized()?;
 
         let execute_frame = execute::Execute {
             id: prepared_statement.get_id().to_owned(),
             parameters: query::QueryParameters {
-                consistency: prepared_statement
-                    .config
-                    .determine_consistency(self.config.default_consistency),
+                consistency,
                 serial_consistency: prepared_statement.get_serial_consistency(),
                 values: &serialized_values,
                 page_size: prepared_statement.get_page_size(),
@@ -544,10 +612,27 @@ impl Connection {
         }
     }
 
+    #[allow(dead_code)]
     pub async fn batch(
         &self,
         batch: &Batch,
         values: impl BatchValues,
+    ) -> Result<BatchResult, QueryError> {
+        self.batch_with_consistency(
+            batch,
+            values,
+            batch
+                .config
+                .determine_consistency(self.config.default_consistency),
+        )
+        .await
+    }
+
+    pub async fn batch_with_consistency(
+        &self,
+        batch: &Batch,
+        values: impl BatchValues,
+        consistency: Consistency,
     ) -> Result<BatchResult, QueryError> {
         let statements_count = batch.statements.len();
         if statements_count != values.len() {
@@ -569,9 +654,7 @@ impl Connection {
             statements_count,
             values,
             batch_type: batch.get_type(),
-            consistency: batch
-                .config
-                .determine_consistency(self.config.default_consistency),
+            consistency,
             serial_consistency: batch.get_serial_consistency(),
             timestamp: batch.get_timestamp(),
         };

scylla/src/transport/iterator.rs

@@ -145,8 +145,12 @@ impl RowIterator {
 
             let choose_connection = |node: Arc<Node>| async move { node.random_connection().await };
 
-            let page_query = |connection: Arc<Connection>, paging_state: Option<Bytes>| async move {
-                connection.query(query_ref, values_ref, paging_state).await
+            let page_query = |connection: Arc<Connection>,
+                              consistency: Consistency,
+                              paging_state: Option<Bytes>| async move {
+                connection
+                    .query_with_consistency(query_ref, values_ref, consistency, paging_state)
+                    .await
             };
 
             let worker = RowIteratorWorker {
@@ -210,9 +214,11 @@ impl RowIterator {
                 }
             };
 
-            let page_query = |connection: Arc<Connection>, paging_state: Option<Bytes>| async move {
+            let page_query = |connection: Arc<Connection>,
+                              consistency: Consistency,
+                              paging_state: Option<Bytes>| async move {
                 connection
-                    .execute(prepared_ref, values_ref, paging_state)
+                    .execute_with_consistency(prepared_ref, values_ref, consistency, paging_state)
                     .await
             };
 
@@ -291,14 +297,15 @@ impl<ConnFunc, ConnFut, QueryFunc, QueryFut> RowIteratorWorker<'_, ConnFunc, Que
 where
     ConnFunc: Fn(Arc<Node>) -> ConnFut,
     ConnFut: Future<Output = Result<Arc<Connection>, QueryError>>,
-    QueryFunc: Fn(Arc<Connection>, Option<Bytes>) -> QueryFut,
+    QueryFunc: Fn(Arc<Connection>, Consistency, Option<Bytes>) -> QueryFut,
     QueryFut: Future<Output = Result<QueryResponse, QueryError>>,
 {
     async fn work(mut self, cluster_data: Arc<ClusterData>) {
         let query_plan = self.load_balancer.plan(&self.statement_info, &cluster_data);
 
         let mut last_error: QueryError =
             QueryError::ProtocolError("Empty query plan - driver bug!");
+        let mut current_consistency: Consistency = self.query_consistency;
 
         'nodes_in_plan: for node in query_plan {
             let span = trace_span!("Executing query", node = node.address.to_string().as_str());
@@ -324,8 +331,10 @@ where
             'same_node_retries: loop {
                 trace!(parent: &span, "Execution started");
                 // Query pages until an error occurs
-                let queries_result: Result<(), QueryError> =
-                    self.query_pages(&connection).instrument(span.clone()).await;
+                let queries_result: Result<(), QueryError> = self
+                    .query_pages(&connection, current_consistency)
+                    .instrument(span.clone())
+                    .await;
 
                 last_error = match queries_result {
                     Ok(()) => {
@@ -355,12 +364,14 @@ where
                     retry_decision = format!("{:?}", retry_decision).as_str()
                 );
                 match retry_decision {
-                    RetryDecision::RetrySameNode => {
+                    RetryDecision::RetrySameNode(cl) => {
                         self.metrics.inc_retries_num();
+                        current_consistency = cl;
                         continue 'same_node_retries;
                     }
-                    RetryDecision::RetryNextNode => {
+                    RetryDecision::RetryNextNode(cl) => {
                         self.metrics.inc_retries_num();
+                        current_consistency = cl;
                         continue 'nodes_in_plan;
                     }
                     RetryDecision::DontRetry => break 'nodes_in_plan,
@@ -373,7 +384,11 @@ where
     }
 
     // Given a working connection query as many pages as possible until the first error
-    async fn query_pages(&mut self, connection: &Arc<Connection>) -> Result<(), QueryError> {
+    async fn query_pages(
+        &mut self,
+        connection: &Arc<Connection>,
+        consistency: Consistency,
+    ) -> Result<(), QueryError> {
         loop {
             self.metrics.inc_total_paged_queries();
             let query_start = std::time::Instant::now();
@@ -383,7 +398,8 @@ where
                 "Sending"
             );
             let query_response: QueryResponse =
-                (self.page_query)(connection.clone(), self.paging_state.clone()).await?;
+                (self.page_query)(connection.clone(), consistency, self.paging_state.clone())
+                    .await?;
 
             match query_response.response {
                 Response::Result(result::Result::Rows(mut rows)) => {