Merge pull request #1446 from Lorak-mmk/refactor-prepared-sharing

Centralize prepared statement data sharing

Author: Lorak-mmk

scylla/src/client/caching_session.rs

@@ -4,14 +4,11 @@
 use crate::errors::{ExecutionError, PagerExecutionError, PrepareError};
 use crate::response::query_result::QueryResult;
 use crate::response::{PagingState, PagingStateResponse};
-use crate::routing::partitioner::PartitionerName;
 use crate::statement::batch::{Batch, BatchStatement};
-use crate::statement::prepared::PreparedStatement;
+use crate::statement::prepared::{PreparedStatement, UnconfiguredPreparedStatement};
 use crate::statement::unprepared::Statement;
-use bytes::Bytes;
 use dashmap::DashMap;
 use futures::future::try_join_all;
-use scylla_cql::frame::response::result::{PreparedMetadata, ResultMetadata};
 use scylla_cql::serialize::batch::BatchValues;
 use scylla_cql::serialize::row::SerializeRow;
 use std::collections::hash_map::RandomState;
@@ -22,19 +19,6 @@ use std::sync::Arc;
 use crate::client::pager::QueryPager;
 use crate::client::session::Session;
 
-/// Contains just the parts of a prepared statement that were returned
-/// from the database. All remaining parts (query string, page size,
-/// consistency, etc.) are taken from the Query passed
-/// to the `CachingSession::execute` family of methods.
-#[derive(Debug)]
-struct RawPreparedStatementData {
-    id: Bytes,
-    is_confirmed_lwt: bool,
-    metadata: PreparedMetadata,
-    result_metadata: Arc<ResultMetadata<'static>>,
-    partitioner_name: PartitionerName,
-}
-
 /// Provides auto caching while executing queries
 pub struct CachingSession<S = RandomState>
 where
@@ -45,7 +29,7 @@ where
     /// If a prepared statement is added while the limit is reached, the oldest prepared statement
     /// is removed from the cache
     max_capacity: usize,
-    cache: DashMap<String, RawPreparedStatementData, S>,
+    cache: DashMap<String, UnconfiguredPreparedStatement, S>,
     use_cached_metadata: bool,
 }
 
@@ -204,16 +188,7 @@ where
 
         if let Some(raw) = self.cache.get(&query.contents) {
             let page_size = query.get_validated_page_size();
-            let mut stmt = PreparedStatement::new(
-                raw.id.clone(),
-                raw.is_confirmed_lwt,
-                raw.metadata.clone(),
-                raw.result_metadata.clone(),
-                query.contents,
-                page_size,
-                query.config,
-            );
-            stmt.set_partitioner_name(raw.partitioner_name.clone());
+            let mut stmt = raw.make_configured_handle(query.config, page_size);
             stmt.set_use_cached_result_metadata(self.use_cached_metadata);
             Ok(stmt)
         } else {
@@ -245,13 +220,7 @@ where
                 }
             }
 
-            let raw = RawPreparedStatementData {
-                id: prepared.get_id().clone(),
-                is_confirmed_lwt: prepared.is_confirmed_lwt(),
-                metadata: prepared.get_prepared_metadata().clone(),
-                result_metadata: prepared.get_result_metadata().clone(),
-                partitioner_name: prepared.get_partitioner_name().clone(),
-            };
+            let raw = prepared.make_unconfigured_handle();
             self.cache.insert(query_contents, raw);
 
             Ok(prepared)

scylla/src/client/session.rs

@@ -1341,20 +1341,20 @@ impl Session {
 
         // Validate prepared ids equality.
         for another_raw_prepared in raw_prepared_statements_results_iter.flatten() {
-            if prepared.get_id() != &another_raw_prepared.prepared_response.id {
+            if prepared.get_id() != another_raw_prepared.get_id() {
                 tracing::error!(
                     "Got differing ids upon statement preparation: statement \"{}\", id1: {:?}, id2: {:?}",
                     prepared.get_statement(),
                     prepared.get_id(),
-                    another_raw_prepared.prepared_response.id
+                    another_raw_prepared.get_id()
                 );
                 return Err(PrepareError::PreparedStatementIdsMismatch);
             }
 
             // Collect all tracing ids from prepare() queries in the final result
             prepared
                 .prepare_tracing_ids
-                .extend(another_raw_prepared.tracing_id);
+                .extend(another_raw_prepared.tracing_id());
         }
 
         // This is the first preparation that succeeded.

scylla/src/network/connection.rs

@@ -21,13 +21,11 @@ use crate::frame::{
 use crate::policies::address_translator::{AddressTranslator, UntranslatedPeer};
 use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::response::query_result::QueryResult;
-use crate::response::{
-    NonErrorAuthResponse, NonErrorStartupResponse, PagingState, QueryResponse, RawPreparedStatement,
-};
+use crate::response::{NonErrorAuthResponse, NonErrorStartupResponse, PagingState, QueryResponse};
 use crate::routing::locator::tablets::{RawTablet, TabletParsingError};
 use crate::routing::{Shard, ShardAwarePortRange, ShardInfo, Sharder, ShardingError};
 use crate::statement::batch::{Batch, BatchStatement};
-use crate::statement::prepared::PreparedStatement;
+use crate::statement::prepared::{PreparedStatement, RawPreparedStatement};
 use crate::statement::unprepared::Statement;
 use crate::statement::{Consistency, PageSize};
 use bytes::Bytes;
@@ -688,13 +686,13 @@ impl Connection {
         previous_prepared: &PreparedStatement,
     ) -> Result<(), RequestAttemptError> {
         let reprepare_query: Statement = query.into();
-        let prepared_response = self.prepare_raw(&reprepare_query).await?.prepared_response;
+        let raw_prepared = self.prepare_raw(&reprepare_query).await?;
 
         // Reprepared statement should keep its id - it's the md5 sum
         // of statement contents
-        if prepared_response.id != previous_prepared.get_id() {
+        if raw_prepared.get_id() != previous_prepared.get_id() {
             Err(RequestAttemptError::RepreparedIdChanged {
-                reprepared_id: prepared_response.id.into(),
+                reprepared_id: raw_prepared.get_id().clone().into(),
                 statement: reprepare_query.contents,
                 expected_id: previous_prepared.get_id().clone().into(),
             })

scylla/src/response/mod.rs

@@ -14,6 +14,5 @@ mod request_response;
 pub use coordinator::Coordinator;
 pub(crate) use request_response::{
     NonErrorAuthResponse, NonErrorQueryResponse, NonErrorStartupResponse, QueryResponse,
-    RawPreparedStatement,
 };
 pub use scylla_cql::frame::request::query::{PagingState, PagingStateResponse};

scylla/src/response/request_response.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::sync::Arc;
 
 use bytes::Bytes;
 use scylla_cql::frame::request::query::PagingStateResponse;
@@ -11,8 +10,6 @@ use crate::errors::RequestAttemptError;
 use crate::frame::response::{self, result};
 use crate::response::Coordinator;
 use crate::response::query_result::QueryResult;
-use crate::statement::Statement;
-use crate::statement::prepared::PreparedStatement;
 
 pub(crate) struct QueryResponse {
     pub(crate) response: Response,
@@ -145,59 +142,3 @@ pub(crate) enum NonErrorAuthResponse {
     AuthChallenge(response::authenticate::AuthChallenge),
     AuthSuccess(response::authenticate::AuthSuccess),
 }
-
-/// Parts which are needed to construct [PreparedStatement].
-///
-/// Kept separate for performance reasons, because constructing
-/// [PreparedStatement] involves allocations.
-pub(crate) struct RawPreparedStatement<'statement> {
-    pub(crate) statement: &'statement Statement,
-    pub(crate) prepared_response: result::Prepared,
-    pub(crate) is_lwt: bool,
-    pub(crate) tracing_id: Option<Uuid>,
-}
-
-impl<'statement> RawPreparedStatement<'statement> {
-    pub(crate) fn new(
-        statement: &'statement Statement,
-        prepared_response: result::Prepared,
-        is_lwt: bool,
-        tracing_id: Option<Uuid>,
-    ) -> Self {
-        Self {
-            statement,
-            prepared_response,
-            is_lwt,
-            tracing_id,
-        }
-    }
-}
-
-/// Constructs the fully-fledged [PreparedStatement].
-///
-/// This involves allocations.
-impl RawPreparedStatement<'_> {
-    pub(crate) fn into_prepared_statement(self) -> PreparedStatement {
-        let Self {
-            statement,
-            prepared_response,
-            is_lwt,
-            tracing_id,
-        } = self;
-        let mut prepared_statement = PreparedStatement::new(
-            prepared_response.id,
-            is_lwt,
-            prepared_response.prepared_metadata,
-            Arc::new(prepared_response.result_metadata),
-            statement.contents.clone(),
-            statement.get_validated_page_size(),
-            statement.config.clone(),
-        );
-
-        if let Some(tracing_id) = tracing_id {
-            prepared_statement.prepare_tracing_ids.push(tracing_id);
-        }
-
-        prepared_statement
-    }
-}

scylla/src/statement/prepared.rs

@@ -18,14 +18,92 @@ use uuid::Uuid;
 use super::{PageSize, StatementConfig};
 use crate::client::execution_profile::ExecutionProfileHandle;
 use crate::errors::{BadQuery, ExecutionError};
-use crate::frame::response::result::PreparedMetadata;
+use crate::frame::response::result::{self, PreparedMetadata};
 use crate::frame::types::{Consistency, SerialConsistency};
 use crate::observability::history::HistoryListener;
 use crate::policies::load_balancing::LoadBalancingPolicy;
 use crate::policies::retry::RetryPolicy;
 use crate::response::query_result::ColumnSpecs;
 use crate::routing::Token;
 use crate::routing::partitioner::{Partitioner, PartitionerHasher, PartitionerName};
+use crate::statement::Statement;
+
+/// Parts which are needed to construct [PreparedStatement].
+///
+/// Kept separate for performance reasons, because constructing
+/// [PreparedStatement] involves allocations.
+///
+/// # Lifecycle of prepared statement
+///
+/// When PREPARE request is issued, RawPreparedStatement is returned,
+/// and later converted to PreparedStatement.
+/// PreparedStatement can be cloned. Clone is a new handle to the same
+/// underlying shared data - if the result metadata is updated, it is
+/// updated for all clones.
+/// PreparedStatement can be turned into UnconfiguredPreparedStatement.
+/// Similarly to clone, unconfigured statement is also a handle to the
+/// same shared data.
+/// UnconfiguredPreparedStatement can be used to create new PreparedStatement
+/// objects. Those are also handles to the same shared data.
+pub(crate) struct RawPreparedStatement<'statement> {
+    statement: &'statement Statement,
+    prepared_response: result::Prepared,
+    is_lwt: bool,
+    tracing_id: Option<Uuid>,
+}
+
+impl<'statement> RawPreparedStatement<'statement> {
+    pub(crate) fn new(
+        statement: &'statement Statement,
+        prepared_response: result::Prepared,
+        is_lwt: bool,
+        tracing_id: Option<Uuid>,
+    ) -> Self {
+        Self {
+            statement,
+            prepared_response,
+            is_lwt,
+            tracing_id,
+        }
+    }
+
+    pub(crate) fn get_id(&self) -> &Bytes {
+        &self.prepared_response.id
+    }
+
+    pub(crate) fn tracing_id(&self) -> Option<Uuid> {
+        self.tracing_id
+    }
+}
+
+/// Constructs the fully-fledged [PreparedStatement].
+///
+/// This involves allocations.
+impl RawPreparedStatement<'_> {
+    pub(crate) fn into_prepared_statement(self) -> PreparedStatement {
+        let Self {
+            statement,
+            prepared_response,
+            is_lwt,
+            tracing_id,
+        } = self;
+        let mut prepared_statement = PreparedStatement::new(
+            prepared_response.id,
+            is_lwt,
+            prepared_response.prepared_metadata,
+            Arc::new(prepared_response.result_metadata),
+            statement.contents.clone(),
+            statement.get_validated_page_size(),
+            statement.config.clone(),
+        );
+
+        if let Some(tracing_id) = tracing_id {
+            prepared_statement.prepare_tracing_ids.push(tracing_id);
+        }
+
+        prepared_statement
+    }
+}
 
 /// Represents a statement prepared on the server.
 ///
@@ -108,7 +186,7 @@ pub struct PreparedStatement {
 struct PreparedStatementSharedData {
     id: Bytes,
     metadata: PreparedMetadata,
-    result_metadata: Arc<ResultMetadata<'static>>,
+    initial_result_metadata: Arc<ResultMetadata<'static>>,
     statement: String,
     is_confirmed_lwt: bool,
 }
@@ -126,7 +204,7 @@ impl Clone for PreparedStatement {
 }
 
 impl PreparedStatement {
-    pub(crate) fn new(
+    fn new(
         id: Bytes,
         is_lwt: bool,
         metadata: PreparedMetadata,
@@ -139,7 +217,7 @@ impl PreparedStatement {
             shared: Arc::new(PreparedStatementSharedData {
                 id,
                 metadata,
-                result_metadata,
+                initial_result_metadata: result_metadata,
                 statement,
                 is_confirmed_lwt: is_lwt,
             }),
@@ -430,12 +508,12 @@ impl PreparedStatement {
 
     /// Access metadata about the result of prepared statement returned by the database
     pub(crate) fn get_result_metadata(&self) -> &Arc<ResultMetadata<'static>> {
-        &self.shared.result_metadata
+        &self.shared.initial_result_metadata
     }
 
     /// Access column specifications of the result set returned after the execution of this statement
     pub fn get_result_set_col_specs(&self) -> ColumnSpecs<'_, 'static> {
-        ColumnSpecs::new(self.shared.result_metadata.col_specs())
+        ColumnSpecs::new(self.shared.initial_result_metadata.col_specs())
     }
 
     /// Get the name of the partitioner used for this statement.
@@ -506,6 +584,41 @@ impl PreparedStatement {
         let ctx = RowSerializationContext::from_prepared(self.get_prepared_metadata());
         SerializedValues::from_serializable(&ctx, values)
     }
+
+    pub(crate) fn make_unconfigured_handle(&self) -> UnconfiguredPreparedStatement {
+        UnconfiguredPreparedStatement {
+            shared: Arc::clone(&self.shared),
+            partitioner_name: self.get_partitioner_name().clone(),
+        }
+    }
+}
+
+/// This is a [PreparedStatement] without parts that are available to be configured
+/// on an unprepared [Statement]. It is intended to be used in [CachingSession] cache,
+/// as a type safety measue: it first needs to be configured with config taken from
+/// statement provided by the user.
+/// It contains partitioner_name field. It is configurable on prepared statement,
+/// but not on unprepared statement, so we need to keep it.
+#[derive(Debug)]
+pub(crate) struct UnconfiguredPreparedStatement {
+    shared: Arc<PreparedStatementSharedData>,
+    partitioner_name: PartitionerName,
+}
+
+impl UnconfiguredPreparedStatement {
+    pub(crate) fn make_configured_handle(
+        &self,
+        config: StatementConfig,
+        page_size: PageSize,
+    ) -> PreparedStatement {
+        PreparedStatement {
+            shared: Arc::clone(&self.shared),
+            prepare_tracing_ids: Vec::new(),
+            page_size,
+            partitioner_name: self.partitioner_name.clone(),
+            config,
+        }
+    }
 }
 
 /// Error when extracting partition key from bound values.