Merge pull request #1329 from wprzytula/construct-prepared-statement-only-where-needed

Construct `PreparedStatement` only where needed

Author: wprzytula

scylla/src/client/session.rs

@@ -30,7 +30,9 @@ use crate::policies::retry::{RequestInfo, RetryDecision, RetrySession};
 use crate::policies::speculative_execution;
 use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::response::query_result::{MaybeFirstRowError, QueryResult, RowsError};
-use crate::response::{NonErrorQueryResponse, PagingState, PagingStateResponse, QueryResponse};
+use crate::response::{
+    NonErrorQueryResponse, PagingState, PagingStateResponse, QueryResponse, RawPreparedStatement,
+};
 use crate::routing::partitioner::PartitionerName;
 use crate::routing::{Shard, ShardAwarePortRange};
 use crate::statement::batch::batch_values;
@@ -1064,15 +1066,14 @@ impl Session {
                         .serial_consistency
                         .unwrap_or(execution_profile.serial_consistency);
                     // Needed to avoid moving query and values into async move block
-                    let statement_ref = &statement;
                     let values_ref = &values;
                     let paging_state_ref = &paging_state;
                     async move {
                         if values_ref.is_empty() {
                             span_ref.record_request_size(0);
                             connection
                                 .query_raw_with_consistency(
-                                    statement_ref,
+                                    statement,
                                     consistency,
                                     serial_consistency,
                                     page_size,
@@ -1081,7 +1082,7 @@ impl Session {
                                 .await
                                 .and_then(QueryResponse::into_non_error_query_response)
                         } else {
-                            let prepared = connection.prepare(statement_ref).await?;
+                            let prepared = connection.prepare(statement).await?;
                             let serialized = prepared.serialize_values(values_ref)?;
                             span_ref.record_request_size(serialized.buffer_size());
                             connection
@@ -1242,29 +1243,28 @@ impl Session {
         let connections_iter = cluster_state.iter_working_connections()?;
 
         // Prepare statements on all connections concurrently
-        let handles = connections_iter.map(|c| async move { c.prepare(statement_ref).await });
+        let handles = connections_iter.map(|c| async move { c.prepare_raw(statement_ref).await });
         let mut results = join_all(handles).await.into_iter();
 
         // If at least one prepare was successful, `prepare()` returns Ok.
         // Find the first result that is Ok, or Err if all failed.
 
         // Safety: there is at least one node in the cluster, and `Cluster::iter_working_connections()`
         // returns either an error or an iterator with at least one connection, so there will be at least one result.
-        let first_ok: Result<PreparedStatement, RequestAttemptError> =
+        let first_ok: Result<RawPreparedStatement, RequestAttemptError> =
             results.by_ref().find_or_first(Result::is_ok).unwrap();
-        let mut prepared: PreparedStatement =
-            first_ok.map_err(|first_attempt| PrepareError::AllAttemptsFailed { first_attempt })?;
+        let mut prepared: PreparedStatement = first_ok
+            .map_err(|first_attempt| PrepareError::AllAttemptsFailed { first_attempt })?
+            .into_prepared_statement();
 
         // Validate prepared ids equality
         for statement in results.flatten() {
-            if prepared.get_id() != statement.get_id() {
+            if prepared.get_id() != &statement.prepared_response.id {
                 return Err(PrepareError::PreparedStatementIdsMismatch);
             }
 
             // Collect all tracing ids from prepare() queries in the final result
-            prepared
-                .prepare_tracing_ids
-                .extend(statement.prepare_tracing_ids);
+            prepared.prepare_tracing_ids.extend(statement.tracing_id);
         }
 
         prepared.set_partitioner_name(

scylla/src/network/connection.rs

@@ -24,6 +24,7 @@ use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::response::query_result::QueryResult;
 use crate::response::{
     NonErrorAuthResponse, NonErrorStartupResponse, PagingState, PagingStateResponse, QueryResponse,
+    RawPreparedStatement,
 };
 use crate::routing::locator::tablets::{RawTablet, TabletParsingError};
 use crate::routing::{Shard, ShardAwarePortRange, ShardInfo, Sharder, ShardingError};
@@ -632,46 +633,56 @@ impl Connection {
         Ok(supported)
     }
 
-    pub(crate) async fn prepare(
+    /// Prepares the given statement and returns raw parts that can be used to construct
+    /// a prepared statement.
+    ///
+    /// Extracted in order to avoid needless allocations upon [PreparedStatement] construction
+    /// in cases when the [PreparedStatement] is not used anyway.
+    pub(crate) async fn prepare_raw<'statement>(
         &self,
-        query: &Statement,
-    ) -> Result<PreparedStatement, RequestAttemptError> {
+        statement: &'statement Statement,
+    ) -> Result<RawPreparedStatement<'statement>, RequestAttemptError> {
         let query_response = self
             .send_request(
                 &request::Prepare {
-                    query: &query.contents,
+                    query: &statement.contents,
                 },
                 true,
-                query.config.tracing,
+                statement.config.tracing,
                 None,
             )
             .await?;
 
-        let mut prepared_statement = match query_response.response {
+        match query_response.response {
             Response::Error(error::Error { error, reason }) => {
-                return Err(RequestAttemptError::DbError(error, reason))
+                Err(RequestAttemptError::DbError(error, reason))
             }
-            Response::Result(result::Result::Prepared(p)) => PreparedStatement::new(
-                p.id,
-                self.features
+            Response::Result(result::Result::Prepared(p)) => {
+                let is_lwt = self
+                    .features
                     .protocol_features
-                    .prepared_flags_contain_lwt_mark(p.prepared_metadata.flags as u32),
-                p.prepared_metadata,
-                Arc::new(p.result_metadata),
-                query.contents.clone(),
-                query.get_validated_page_size(),
-                query.config.clone(),
-            ),
-            _ => {
-                return Err(RequestAttemptError::UnexpectedResponse(
-                    query_response.response.to_response_kind(),
+                    .prepared_flags_contain_lwt_mark(p.prepared_metadata.flags as u32);
+                Ok(RawPreparedStatement::new(
+                    statement,
+                    p,
+                    is_lwt,
+                    query_response.tracing_id,
                 ))
             }
-        };
-
-        if let Some(tracing_id) = query_response.tracing_id {
-            prepared_statement.prepare_tracing_ids.push(tracing_id);
+            _ => Err(RequestAttemptError::UnexpectedResponse(
+                query_response.response.to_response_kind(),
+            )),
         }
+    }
+
+    /// Prepares the given statement and returns [PreparedStatement].
+    pub(crate) async fn prepare(
+        &self,
+        statement: &Statement,
+    ) -> Result<PreparedStatement, RequestAttemptError> {
+        let raw_prepared_statement = self.prepare_raw(statement).await?;
+        let prepared_statement = raw_prepared_statement.into_prepared_statement();
+
         Ok(prepared_statement)
     }
 
@@ -681,14 +692,15 @@ impl Connection {
         previous_prepared: &PreparedStatement,
     ) -> Result<(), RequestAttemptError> {
         let reprepare_query: Statement = query.into();
-        let reprepared = self.prepare(&reprepare_query).await?;
+        let prepared_response = self.prepare_raw(&reprepare_query).await?.prepared_response;
+
         // Reprepared statement should keep its id - it's the md5 sum
         // of statement contents
-        if reprepared.get_id() != previous_prepared.get_id() {
+        if prepared_response.id != previous_prepared.get_id() {
             Err(RequestAttemptError::RepreparedIdChanged {
+                reprepared_id: prepared_response.id.into(),
                 statement: reprepare_query.contents,
                 expected_id: previous_prepared.get_id().clone().into(),
-                reprepared_id: reprepared.get_id().clone().into(),
             })
         } else {
             Ok(())

scylla/src/response/mod.rs

@@ -12,5 +12,6 @@ mod request_response;
 
 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,4 +1,5 @@
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use bytes::Bytes;
 use scylla_cql::frame::request::query::PagingStateResponse;
@@ -9,6 +10,8 @@ use uuid::Uuid;
 use crate::errors::RequestAttemptError;
 use crate::frame::response::{self, result};
 use crate::response::query_result::QueryResult;
+use crate::statement::prepared::PreparedStatement;
+use crate::statement::Statement;
 
 pub(crate) struct QueryResponse {
     pub(crate) response: Response,
@@ -108,3 +111,59 @@ 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
+    }
+}