Merge pull request #1340 from Lorak-mmk/caching_session_skip_metadata

CachingSession: Allow skipping result metadata

Author: Lorak-mmk

.gitignore

@@ -1,4 +1,5 @@
 .vscode
+.zed
 target/
 Cargo.lock
 /book/book

scylla/src/client/caching_session.rs

@@ -37,12 +37,13 @@ pub struct CachingSession<S = RandomState>
 where
     S: Clone + BuildHasher,
 {
-    session: Session,
+    session: Arc<Session>,
     /// The prepared statement cache size
     /// 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>,
+    use_cached_metadata: bool,
 }
 
 impl<S> fmt::Debug for CachingSession<S>
@@ -64,9 +65,10 @@ where
 {
     pub fn from(session: Session, cache_size: usize) -> Self {
         Self {
-            session,
+            session: Arc::new(session),
             max_capacity: cache_size,
             cache: Default::default(),
+            use_cached_metadata: false,
         }
     }
 }
@@ -79,9 +81,10 @@ where
     /// and a [`BuildHasher`], using a customer hasher.
     pub fn with_hasher(session: Session, cache_size: usize, hasher: S) -> Self {
         Self {
-            session,
+            session: Arc::new(session),
             max_capacity: cache_size,
             cache: DashMap::with_hasher(hasher),
+            use_cached_metadata: false,
         }
     }
 }
@@ -207,10 +210,15 @@ where
                 query.config,
             );
             stmt.set_partitioner_name(raw.partitioner_name.clone());
+            stmt.set_use_cached_result_metadata(self.use_cached_metadata);
             Ok(stmt)
         } else {
             let query_contents = query.contents.clone();
-            let prepared = self.session.prepare(query).await?;
+            let prepared = {
+                let mut stmt = self.session.prepare(query).await?;
+                stmt.set_use_cached_result_metadata(self.use_cached_metadata);
+                stmt
+            };
 
             if self.max_capacity == self.cache.len() {
                 // Cache is full, remove the first entry
@@ -280,20 +288,26 @@ pub struct CachingSessionBuilder<S = RandomState>
 where
     S: Clone + BuildHasher,
 {
-    session: Session,
+    session: Arc<Session>,
     max_capacity: usize,
     hasher: S,
+    use_cached_metadata: bool,
 }
 
 impl CachingSessionBuilder<RandomState> {
     /// Wraps a [Session] and creates a new [CachingSessionBuilder] instance,
     /// which can be used to create a new [CachingSession].
     ///
     pub fn new(session: Session) -> Self {
+        Self::new_shared(Arc::new(session))
+    }
+
+    pub fn new_shared(session: Arc<Session>) -> Self {
         Self {
             session,
             max_capacity: DEFAULT_MAX_CAPACITY,
             hasher: RandomState::default(),
+            use_cached_metadata: false,
         }
     }
 }
@@ -308,9 +322,33 @@ where
         self
     }
 
+    /// Make use of cached metadata to decode results
+    /// of the statement's execution.
+    ///
+    /// If true, the driver will request the server not to
+    /// attach the result metadata in response to the statement execution.
+    ///
+    /// The driver will cache the result metadata received from the server
+    /// after statement preparation and will use it
+    /// to deserialize the results of statement execution.
+    ///
+    /// See documentation of [`PreparedStatement`] for more details on limitations
+    /// of this functionality.
+    ///
+    /// This option is false by default.
+    pub fn use_cached_result_metadata(mut self, use_cached_metadata: bool) -> Self {
+        self.use_cached_metadata = use_cached_metadata;
+        self
+    }
+
     /// Finishes configuration of [CachingSession].
     pub fn build(self) -> CachingSession<S> {
-        CachingSession::with_hasher(self.session, self.max_capacity, self.hasher)
+        CachingSession {
+            session: self.session,
+            max_capacity: self.max_capacity,
+            cache: DashMap::with_hasher(self.hasher),
+            use_cached_metadata: self.use_cached_metadata,
+        }
     }
 }
 
@@ -364,11 +402,13 @@ where
             session,
             max_capacity,
             hasher: _,
+            use_cached_metadata,
         } = self;
         CachingSessionBuilder {
             session,
             max_capacity,
             hasher,
+            use_cached_metadata,
         }
     }
 }

scylla/tests/integration/session/caching_session.rs

@@ -0,0 +1,105 @@
+use std::sync::Arc;
+
+use scylla::client::caching_session::{CachingSession, CachingSessionBuilder};
+use scylla::client::session_builder::SessionBuilder;
+use scylla_cql::frame::request::Execute;
+use scylla_cql::frame::request::Request;
+use scylla_proxy::Condition;
+use scylla_proxy::ProxyError;
+use scylla_proxy::Reaction;
+use scylla_proxy::RequestFrame;
+use scylla_proxy::RequestOpcode;
+use scylla_proxy::RequestReaction;
+use scylla_proxy::RequestRule;
+use scylla_proxy::WorkerError;
+use tokio::sync::mpsc;
+
+use crate::utils::test_with_3_node_cluster;
+
+#[tokio::test]
+#[cfg_attr(scylla_cloud_tests, ignore)]
+async fn test_caching_session_metadata_cache() {
+    let res = test_with_3_node_cluster(
+        scylla_proxy::ShardAwareness::QueryNode,
+        |proxy_uris, translation_map, mut running_proxy| async move {
+            let (feedback_tx, mut feedback_rx) = mpsc::unbounded_channel();
+            let prepared_request_feedback_rule = RequestRule(
+                Condition::and(
+                    Condition::not(Condition::ConnectionRegisteredAnyEvent),
+                    Condition::RequestOpcode(RequestOpcode::Execute),
+                ),
+                RequestReaction::noop().with_feedback_when_performed(feedback_tx),
+            );
+            for node in running_proxy.running_nodes.iter_mut() {
+                node.change_request_rules(Some(vec![prepared_request_feedback_rule.clone()]));
+            }
+
+            async fn verify_statement_metadata(
+                session: &CachingSession,
+                statement: &str,
+                should_have_metadata: bool,
+                feedback: &mut mpsc::UnboundedReceiver<(RequestFrame, Option<u16>)>,
+            ) {
+                let _result = session.execute_unpaged(statement, ()).await.unwrap();
+                let (req_frame, _) = feedback.recv().await.unwrap();
+                let _ = feedback.try_recv().unwrap_err(); // There should be only one frame.
+                let request = req_frame.deserialize().unwrap();
+                let Request::Execute(Execute { parameters, .. }) = request else {
+                    panic!("Unexpected request type");
+                };
+                let has_metadata = !parameters.skip_metadata;
+                assert_eq!(has_metadata, should_have_metadata);
+            }
+
+            const REQUEST: &str = "SELECT * FROM system.local WHERE key = 'local'";
+
+            let session = Arc::new(
+                SessionBuilder::new()
+                    .known_node(proxy_uris[0].as_str())
+                    .address_translator(Arc::new(translation_map.clone()))
+                    .build()
+                    .await
+                    .unwrap(),
+            );
+            let caching_session: CachingSession =
+                CachingSessionBuilder::new_shared(Arc::clone(&session))
+                    .use_cached_result_metadata(false) // Default, set just to be more explicit
+                    .build();
+
+            // Skipping metadata was not set, so metadata should be present
+            verify_statement_metadata(&caching_session, REQUEST, true, &mut feedback_rx).await;
+
+            // It should also be present when executing statement already in cache
+            verify_statement_metadata(&caching_session, REQUEST, true, &mut feedback_rx).await;
+
+            let caching_session: CachingSession =
+                CachingSessionBuilder::new_shared(Arc::clone(&session))
+                    .use_cached_result_metadata(true)
+                    .build();
+
+            // Now we set skip_metadata to true, so metadata should not be present for a new query
+            verify_statement_metadata(&caching_session, REQUEST, false, &mut feedback_rx).await;
+
+            // It should also not be present when executing statement already in cache
+            verify_statement_metadata(&caching_session, REQUEST, false, &mut feedback_rx).await;
+
+            // Test also without setting it explicitly, to verify that it is false by default.
+            let caching_session: CachingSession =
+                CachingSessionBuilder::new_shared(Arc::clone(&session)).build();
+
+            // Skipping metadata was not set, so metadata should be present
+            verify_statement_metadata(&caching_session, REQUEST, true, &mut feedback_rx).await;
+
+            // It should also be present when executing statement already in cache
+            verify_statement_metadata(&caching_session, REQUEST, true, &mut feedback_rx).await;
+
+            running_proxy
+        },
+    )
+    .await;
+    match res {
+        Ok(()) => (),
+        Err(ProxyError::Worker(WorkerError::DriverDisconnected(_))) => (),
+        Err(err) => panic!("{}", err),
+    }
+}

scylla/tests/integration/session/mod.rs

@@ -1,4 +1,5 @@
 mod authenticate;
+mod caching_session;
 mod history;
 mod new_session;
 mod retries;