adapt after master merge

Author: Ten0

scylla-cql/src/types/serialize/row.rs

@@ -80,6 +80,17 @@ pub trait SerializeRow {
     /// the bind marker types and names so that the values can be properly
     /// type checked and serialized.
     fn is_empty(&self) -> bool;
+
+    /// Specialization that allows the driver to not re-serialize the row if it's already
+    /// a `SerializedValues`
+    ///
+    /// Note that if using this, it's the user's responsibility to ensure that this
+    /// `SerializedValues` has been generated with the same prepared statement as the query
+    /// is going to be made with.
+    #[inline]
+    fn already_serialized(&self) -> Option<&SerializedValues> {
+        None
+    }
 }
 
 macro_rules! fallback_impl_contents {
@@ -255,12 +266,35 @@ impl<T: SerializeRow + ?Sized> SerializeRow for &T {
         ctx: &RowSerializationContext<'_>,
         writer: &mut RowWriter,
     ) -> Result<(), SerializationError> {
-        <T as SerializeRow>::serialize(self, ctx, writer)
+        <T as SerializeRow>::serialize(*self, ctx, writer)
     }
 
     #[inline]
     fn is_empty(&self) -> bool {
-        <T as SerializeRow>::is_empty(self)
+        <T as SerializeRow>::is_empty(*self)
+    }
+
+    #[inline]
+    fn already_serialized(&self) -> Option<&SerializedValues> {
+        <T as SerializeRow>::already_serialized(*self)
+    }
+}
+
+impl SerializeRow for SerializedValues {
+    fn serialize(
+        &self,
+        _ctx: &RowSerializationContext<'_>,
+        writer: &mut RowWriter,
+    ) -> Result<(), SerializationError> {
+        Ok(writer.append_serialize_row(self))
+    }
+
+    fn is_empty(&self) -> bool {
+        self.is_empty()
+    }
+
+    fn already_serialized(&self) -> Option<&SerializedValues> {
+        Some(self)
     }
 }
 

scylla/src/statement/batch.rs

@@ -4,8 +4,10 @@ use std::sync::Arc;
 use crate::history::HistoryListener;
 use crate::load_balancing;
 use crate::retry_policy::RetryPolicy;
+use crate::routing::Shard;
 use crate::statement::{prepared_statement::PreparedStatement, query::Query};
-use crate::transport::{execution_profile::ExecutionProfileHandle, Node};
+use crate::transport::execution_profile::ExecutionProfileHandle;
+use crate::transport::NodeRef;
 use crate::Session;
 
 use super::StatementConfig;
@@ -145,31 +147,31 @@ impl Batch {
         self.config.execution_profile_handle.as_ref()
     }
 
-    /// Associates the batch with a new execution profile that will have a load balancing policy
-    /// that will enforce the use of the provided [`Node`] to the extent possible.
+    /// Associates the batch with a new execution profile that will have a load
+    /// balancing policy that will enforce the use of the provided [`Node`]
+    /// to the extent possible.
     ///
-    /// This should typically be used in conjunction with [`Session::shard_for_statement`], where
-    /// you would constitute a batch by assigning to the same batch all the statements that would be executed in
-    /// the same shard.
+    /// This should typically be used in conjunction with
+    /// [`Session::shard_for_statement`], where you would constitute a batch
+    /// by assigning to the same batch all the statements that would be executed
+    /// in the same shard.
     ///
-    /// Since it is not guaranteed that subsequent calls to the load balancer would re-assign the statement
-    /// to the same node, you should use this method to enforce the use of the original node that was envisioned by
+    /// Since it is not guaranteed that subsequent calls to the load balancer
+    /// would re-assign the statement to the same node, you should use this
+    /// method to enforce the use of the original node that was envisioned by
     /// `shard_for_statement` for the batch:
     ///
     /// ```rust
     /// # use scylla::Session;
     /// # use std::error::Error;
     /// # async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
-    /// use scylla::{
-    ///     batch::Batch,
-    ///     frame::value::{SerializedValues, ValueList},
-    /// };
+    /// use scylla::{batch::Batch, serialize::row::SerializedValues};
     ///
     /// let prepared_statement = session
     ///     .prepare("INSERT INTO ks.tab(a, b) VALUES(?, ?)")
     ///     .await?;
     ///
-    /// let serialized_values: SerializedValues = (1, 2).serialized()?.into_owned();
+    /// let serialized_values: SerializedValues = prepared_statement.serialize_values(&(1, 2))?;
     /// let shard = session.shard_for_statement(&prepared_statement, &serialized_values)?;
     ///
     /// // Send that to a task that will handle statements targeted to the same shard
@@ -178,8 +180,8 @@ impl Batch {
     /// // Constitute a batch with all the statements that would be executed in the same shard
     ///
     /// let mut batch: Batch = Default::default();
-    /// if let Some((node, _shard_idx)) = shard {
-    ///     batch.enforce_target_node(&node, &session);
+    /// if let Some((node, shard_idx)) = shard {
+    ///     batch.enforce_target_node(&node, shard_idx, &session);
     /// }
     /// let mut batch_values = Vec::new();
     ///
@@ -195,13 +197,14 @@ impl Batch {
     /// ```
     ///
     ///
-    /// If the target node is not available anymore at the time of executing the statement, it will fallback to the
-    /// original load balancing policy:
+    /// If the target node is not available anymore at the time of executing the
+    /// statement, it will fallback to the original load balancing policy:
     /// - Either that currently set on the [`Batch`], if any
     /// - Or that of the [`Session`] if there isn't one on the `Batch`
     pub fn enforce_target_node(
         &mut self,
-        node: &Arc<Node>,
+        node: NodeRef<'_>,
+        shard: Shard,
         base_execution_profile_from_session: &Session,
     ) {
         let execution_profile_handle = self.get_execution_profile_handle().unwrap_or_else(|| {
@@ -210,8 +213,9 @@ impl Batch {
         self.set_execution_profile_handle(Some(
             execution_profile_handle
                 .pointee_to_builder()
-                .load_balancing_policy(Arc::new(load_balancing::EnforceTargetNodePolicy::new(
+                .load_balancing_policy(Arc::new(load_balancing::EnforceTargetShardPolicy::new(
                     node,
+                    shard,
                     execution_profile_handle.load_balancing_policy(),
                 )))
                 .build()

scylla/src/statement/prepared_statement.rs

@@ -459,7 +459,7 @@ impl PreparedStatement {
         self.config.execution_profile_handle.as_ref()
     }
 
-    pub(crate) fn serialize_values(
+    pub fn serialize_values(
         &self,
         values: &impl SerializeRow,
     ) -> Result<SerializedValues, SerializationError> {

scylla/src/transport/load_balancing/default.rs

@@ -703,7 +703,7 @@ impl DefaultPolicy {
         vec.into_iter()
     }
 
-    fn is_alive(node: NodeRef, _shard: Option<Shard>) -> bool {
+    pub(crate) fn is_alive(node: NodeRef, _shard: Option<Shard>) -> bool {
         // For now, we leave this as stub, until we have time to improve node events.
         // node.is_enabled() && !node.is_down()
         node.is_enabled()

scylla/src/transport/load_balancing/enforce_node.rs

@@ -1,5 +1,8 @@
 use super::{DefaultPolicy, FallbackPlan, LoadBalancingPolicy, NodeRef, RoutingInfo};
-use crate::transport::{cluster::ClusterData, Node};
+use crate::{
+    routing::Shard,
+    transport::{cluster::ClusterData, Node},
+};
 use std::sync::Arc;
 use uuid::Uuid;
 
@@ -8,25 +11,36 @@ use uuid::Uuid;
 ///
 /// This is meant to be used for shard-aware batching.
 #[derive(Debug)]
-pub struct EnforceTargetNodePolicy {
+pub struct EnforceTargetShardPolicy {
     target_node: Uuid,
+    shard: Shard,
     fallback: Arc<dyn LoadBalancingPolicy>,
 }
 
-impl EnforceTargetNodePolicy {
-    pub fn new(target_node: &Arc<Node>, fallback: Arc<dyn LoadBalancingPolicy>) -> Self {
+impl EnforceTargetShardPolicy {
+    pub fn new(
+        target_node: &Arc<Node>,
+        shard: Shard,
+        fallback: Arc<dyn LoadBalancingPolicy>,
+    ) -> Self {
         Self {
             target_node: target_node.host_id,
+            shard,
             fallback,
         }
     }
 }
-impl LoadBalancingPolicy for EnforceTargetNodePolicy {
-    fn pick<'a>(&'a self, query: &'a RoutingInfo, cluster: &'a ClusterData) -> Option<NodeRef<'a>> {
+impl LoadBalancingPolicy for EnforceTargetShardPolicy {
+    fn pick<'a>(
+        &'a self,
+        query: &'a RoutingInfo,
+        cluster: &'a ClusterData,
+    ) -> Option<(NodeRef<'a>, Option<Shard>)> {
         cluster
             .known_peers
             .get(&self.target_node)
-            .filter(DefaultPolicy::is_alive)
+            .map(|node| (node, Some(self.shard)))
+            .filter(|&(node, shard)| DefaultPolicy::is_alive(node, shard))
             .or_else(|| self.fallback.pick(query, cluster))
     }
 
@@ -40,7 +54,7 @@ impl LoadBalancingPolicy for EnforceTargetNodePolicy {
 
     fn name(&self) -> String {
         format!(
-            "Enforce target node Load balancing policy - Node: {} - fallback: {}",
+            "Enforce target shard Load balancing policy - Node: {} - fallback: {}",
             self.target_node,
             self.fallback.name()
         )

scylla/src/transport/load_balancing/mod.rs

@@ -17,7 +17,7 @@ mod plan;
 pub use plan::Plan;
 pub use {
     default::{DefaultPolicy, DefaultPolicyBuilder, LatencyAwarenessBuilder},
-    enforce_node::EnforceTargetNodePolicy,
+    enforce_node::EnforceTargetShardPolicy,
 };
 
 /// Represents info about statement that can be used by load balancing policies.