statement/transport: re-add per-statement setters for retry policy

The set_retry_policy and get_retry_policy methods of Query,
PreparedStatement and Batch were removed #592, but are re-added here due
to user request and because they will be useful for cpp-rust-driver.

One important difference with the old {set,get}_retry_policy is that the
new methods take and Arc<dyn RetryPolicy> instead of Box<dyn
RetryPolicy>. This was done in order not to introduce additional
allocations when a query/statement/batch is cloned.

The retry_policy field was added directly to
Batch/Query/PreparedStatement instead of StatementConfig. Unfortunately,
StatementConfig is public along with all of its fields, so adding more
fields to it will be a breaking change. This should be fixed
during/after #701.

Tested manually by setting a custom retry policy that prints a message
when it is used.

Fixes: #692

Author: piodul

scylla/src/statement/batch.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use crate::history::HistoryListener;
+use crate::retry_policy::RetryPolicy;
 use crate::statement::{prepared_statement::PreparedStatement, query::Query};
 use crate::transport::execution_profile::ExecutionProfileHandle;
 
@@ -15,6 +16,9 @@ pub use crate::frame::request::batch::BatchType;
 pub struct Batch {
     pub(crate) config: StatementConfig,
 
+    // TODO: Move this after #701 is fixed
+    retry_policy: Option<Arc<dyn RetryPolicy>>,
+
     pub statements: Vec<BatchStatement>,
     batch_type: BatchType,
 }
@@ -108,6 +112,18 @@ impl Batch {
         self.config.timestamp
     }
 
+    /// Set the retry policy for this batch, overriding the one from execution profile if not None.
+    #[inline]
+    pub fn set_retry_policy(&mut self, retry_policy: Option<Arc<dyn RetryPolicy>>) {
+        self.retry_policy = retry_policy;
+    }
+
+    /// Get the retry policy set for the batch.
+    #[inline]
+    pub fn get_retry_policy(&self) -> Option<&Arc<dyn RetryPolicy>> {
+        self.retry_policy.as_ref()
+    }
+
     /// Sets the listener capable of listening what happens during query execution.
     pub fn set_history_listener(&mut self, history_listener: Arc<dyn HistoryListener>) {
         self.config.history_listener = Some(history_listener);
@@ -134,6 +150,7 @@ impl Default for Batch {
     fn default() -> Self {
         Self {
             statements: Vec::new(),
+            retry_policy: None,
             batch_type: BatchType::Logged,
             config: Default::default(),
         }

scylla/src/statement/prepared_statement.rs

@@ -11,6 +11,7 @@ use crate::frame::response::result::PreparedMetadata;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::frame::value::SerializedValues;
 use crate::history::HistoryListener;
+use crate::retry_policy::RetryPolicy;
 use crate::transport::execution_profile::ExecutionProfileHandle;
 use crate::transport::partitioner::PartitionerName;
 
@@ -20,6 +21,9 @@ pub struct PreparedStatement {
     pub(crate) config: StatementConfig,
     pub prepare_tracing_ids: Vec<Uuid>,
 
+    // TODO: Move this after #701 is fixed
+    retry_policy: Option<Arc<dyn RetryPolicy>>,
+
     id: Bytes,
     shared: Arc<PreparedStatementSharedData>,
     page_size: Option<i32>,
@@ -37,6 +41,7 @@ impl Clone for PreparedStatement {
     fn clone(&self) -> Self {
         Self {
             config: self.config.clone(),
+            retry_policy: self.retry_policy.clone(),
             prepare_tracing_ids: Vec::new(),
             id: self.id.clone(),
             shared: self.shared.clone(),
@@ -53,6 +58,7 @@ impl PreparedStatement {
         is_lwt: bool,
         metadata: PreparedMetadata,
         statement: String,
+        retry_policy: Option<Arc<dyn RetryPolicy>>,
         page_size: Option<i32>,
         config: StatementConfig,
     ) -> Self {
@@ -62,6 +68,7 @@ impl PreparedStatement {
                 metadata,
                 statement,
             }),
+            retry_policy,
             prepare_tracing_ids: Vec::new(),
             page_size,
             config,
@@ -289,6 +296,18 @@ impl PreparedStatement {
         &self.partitioner_name
     }
 
+    /// Set the retry policy for this statement, overriding the one from execution profile if not None.
+    #[inline]
+    pub fn set_retry_policy(&mut self, retry_policy: Option<Arc<dyn RetryPolicy>>) {
+        self.retry_policy = retry_policy;
+    }
+
+    /// Get the retry policy set for the statement.
+    #[inline]
+    pub fn get_retry_policy(&self) -> Option<&Arc<dyn RetryPolicy>> {
+        self.retry_policy.as_ref()
+    }
+
     /// Sets the listener capable of listening what happens during query execution.
     pub fn set_history_listener(&mut self, history_listener: Arc<dyn HistoryListener>) {
         self.config.history_listener = Some(history_listener);

scylla/src/statement/query.rs

@@ -1,6 +1,7 @@
 use super::StatementConfig;
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::history::HistoryListener;
+use crate::retry_policy::RetryPolicy;
 use crate::transport::execution_profile::ExecutionProfileHandle;
 use std::sync::Arc;
 use std::time::Duration;
@@ -12,6 +13,9 @@ use std::time::Duration;
 pub struct Query {
     pub(crate) config: StatementConfig,
 
+    // TODO: Move this after #701 is fixed
+    retry_policy: Option<Arc<dyn RetryPolicy>>,
+
     pub contents: String,
     page_size: Option<i32>,
 }
@@ -21,6 +25,7 @@ impl Query {
     pub fn new(query_text: impl Into<String>) -> Self {
         Self {
             contents: query_text.into(),
+            retry_policy: None,
             page_size: None,
             config: Default::default(),
         }
@@ -123,6 +128,18 @@ impl Query {
         self.config.request_timeout
     }
 
+    /// Set the retry policy for this statement, overriding the one from execution profile if not None.
+    #[inline]
+    pub fn set_retry_policy(&mut self, retry_policy: Option<Arc<dyn RetryPolicy>>) {
+        self.retry_policy = retry_policy;
+    }
+
+    /// Get the retry policy set for the statement.
+    #[inline]
+    pub fn get_retry_policy(&self) -> Option<&Arc<dyn RetryPolicy>> {
+        self.retry_policy.as_ref()
+    }
+
     /// Sets the listener capable of listening what happens during query execution.
     pub fn set_history_listener(&mut self, history_listener: Arc<dyn HistoryListener>) {
         self.config.history_listener = Some(history_listener);

scylla/src/transport/caching_session.rs

@@ -166,11 +166,13 @@ where
 
         if let Some(raw) = self.cache.get(&query.contents) {
             let page_size = query.get_page_size();
+            let retry_policy = query.get_retry_policy().cloned();
             let mut stmt = PreparedStatement::new(
                 raw.id.clone(),
                 raw.is_confirmed_lwt,
                 raw.metadata.clone(),
                 query.contents,
+                retry_policy,
                 page_size,
                 query.config,
             );

scylla/src/transport/connection.rs

@@ -423,6 +423,7 @@ impl Connection {
                     .prepared_flags_contain_lwt_mark(p.prepared_metadata.flags as u32),
                 p.prepared_metadata,
                 query.contents.clone(),
+                query.get_retry_policy().cloned(),
                 query.get_page_size(),
                 query.config.clone(),
             ),

scylla/src/transport/iterator.rs

@@ -151,7 +151,11 @@ impl RowIterator {
             .serial_consistency
             .unwrap_or(execution_profile.serial_consistency);
 
-        let retry_session = execution_profile.retry_policy.new_session();
+        let retry_session = query
+            .get_retry_policy()
+            .map(|rp| &**rp)
+            .unwrap_or(&*execution_profile.retry_policy)
+            .new_session();
 
         let parent_span = tracing::Span::current();
         let worker_task = async move {
@@ -222,7 +226,12 @@ impl RowIterator {
             .config
             .serial_consistency
             .unwrap_or(config.execution_profile.serial_consistency);
-        let retry_session = config.execution_profile.retry_policy.new_session();
+        let retry_session = config
+            .prepared
+            .get_retry_policy()
+            .map(|rp| &**rp)
+            .unwrap_or(&*config.execution_profile.retry_policy)
+            .new_session();
 
         let parent_span = tracing::Span::current();
         let worker_task = async move {

scylla/src/transport/session.rs

@@ -7,6 +7,7 @@ use crate::cloud::CloudConfig;
 use crate::frame::types::LegacyConsistency;
 use crate::history;
 use crate::history::HistoryListener;
+use crate::retry_policy::RetryPolicy;
 use arc_swap::ArcSwapOption;
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -581,6 +582,7 @@ impl Session {
             .run_query(
                 RoutingInfo::default(),
                 &query.config,
+                query.get_retry_policy().map(|rp| &**rp),
                 |node: Arc<Node>| async move { node.random_connection().await },
                 |connection: Arc<Connection>,
                  consistency: Consistency,
@@ -920,6 +922,7 @@ impl Session {
             .run_query(
                 statement_info,
                 &prepared.config,
+                prepared.get_retry_policy().map(|rp| &**rp),
                 |node: Arc<Node>| async move {
                     match token {
                         Some(token) => node.connection_for_token(token).await,
@@ -1117,6 +1120,7 @@ impl Session {
             .run_query(
                 statement_info,
                 &batch.config,
+                batch.get_retry_policy().map(|rp| &**rp),
                 |node: Arc<Node>| async move {
                     match first_value_token {
                         Some(first_value_token) => {
@@ -1430,6 +1434,7 @@ impl Session {
         &'a self,
         statement_info: RoutingInfo<'a>,
         statement_config: &'a StatementConfig,
+        statement_retry_policy: Option<&dyn RetryPolicy>,
         choose_connection: impl Fn(Arc<Node>) -> ConnFut,
         do_query: impl Fn(Arc<Connection>, Consistency, &ExecutionProfileInner) -> QueryFut,
         request_span: &'a RequestSpan,
@@ -1479,7 +1484,7 @@ impl Session {
                 }
             }
 
-            let retry_policy = &execution_profile.retry_policy;
+            let retry_policy = statement_retry_policy.unwrap_or(&*execution_profile.retry_policy);
 
             let speculative_policy = execution_profile.speculative_execution_policy.as_ref();
 
@@ -1754,6 +1759,7 @@ impl Session {
             .run_query(
                 info,
                 &config,
+                None, // No specific retry policy needed for schema agreement
                 |node: Arc<Node>| async move { node.random_connection().await },
                 do_query,
                 &span,