Repository: scylladb/scylla-rust-driver
Pull request #508 (2022-08-14 22:15:17)

Pick connections based on batch first statement's shard
Fixes: #448
Helps #468

To achieve this, it was useful to rework the BatchValues trait.
With this rework, it's also possible to pretty easily provide cloneable `Iterator`s over `ValueList` as `BatchValues` (through `BatchValuesFromIterator`) - so #499 may get closed if this gets merged (or this PR will require update if #499 gets merged before).

## Pre-review checklist
- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


cvybhu: Alright I had a go at the patch. My initial reaction was that it's way too complicated, but after thinking this through I think I agree that this is the way. Here is a summary to gather thoughts and make reviewing easier for others.

We want a few things from the new `BatchValues` trait:
1) We want to be able to write all of them to the request `BufMut`
2) We want to be able to get `SerializedValues` for one of the statements inside the batch in order to calculate the token. Being able to get them only for the first `ValueList` would work for now, but maybe in the future we will improve the heuristic - it would be good to be able to get them for any chosen statement.
3) Writing to request and getting the values must be done at separate times, because calculating token happens in `Session::batch`, while request serialization is much deeper. That's why any iterators of `ValueList` must be `Clone`.
4) We want to get the number of `BatchValues` in `Session::batch` to validate that the user passed values for each statement inside the batch. For all common implementations (tuples, slices, vecs, etc) this is easy to get. For iterators they have to be an `ExactSizeIterator` or we could do `.clone().count()`
5) `BatchValues` should be creatable from an iterator of `ValueLists` to avoid collecting it and needlessly allocating

Seeing these requirements one could attempt to write:
```rust
trait BatchValues {
    type BatchValuesIter: BatchValuesIterator;

    fn iter(&self) -> Self::BatchValuesIter;
    fn len(&self) -> usize;
}

trait BatchValuesIterator: Clone {
    fn write_next_to_request (&mut self,buf:  &mut impl BufMut) -> Option<()>;
    fn next_serialized(&mut self) -> Option
}
```

But there is a problem when the type `BatchValuesIter` has a refernce and requires a lifetime.
With GATs we could write:
```rust
trait BatchValues {
    type BatchValuesIter<'a>: BatchValuesIterator<'a>;

    fn iter<'a>(&'a self) -> Self::BatchValuesIter<'a>;
}
```
But because there are no GATs we end up with the workaround (which is really cool btw).

Looks very nice, I will read the rest and submit a review.
Ten0: > Alright I had a go at the patch. My initial reaction was that it's way too complicated, but after thinking this through I think I agree that this is the way.

That was exactly my thought process, thanks a lot for the detailed write-up! 😃

> which is really cool btw

Yeah at this point it has saved us more than once within [Diesel](https://diesel.rs/) 😊
cvybhu: I'm worried that our tests might not test all code paths here.
Everywhere I look the first statement in a batch is shard aware.

In cases where the first `PreparedStatement` is not shard aware `calculate_token` will fail, This failure seems to be handled gracefully but still it would be nice to actually test it.
Ten0: > In cases where the first `PreparedStatement` is not shard aware `calculate_token` will fail

Returns `Ok(None)`, I'm not sure whether to even call that a fail ^^
https://github.com/scylladb/scylla-rust-driver/blob/e484adc61a1f2ff087594863ac37ca6ae2d49117/scylla/src/transport/session.rs#L1346-L1348

Anyway, I'm just handling it the exact same way as is done in `execute_paged`:
https://github.com/scylladb/scylla-rust-driver/blob/dea816668d1496edd1c35ea5bc6944b0b7de81a7/scylla/src/transport/session.rs#L887
https://github.com/scylladb/scylla-rust-driver/blob/e484adc61a1f2ff087594863ac37ca6ae2d49117/scylla/src/transport/session.rs#L740-L743
cvybhu: There is a merge conflict with `scylla/src/transport/session.rs `, please rebase.
cvybhu: Also could you squash these two commits together, both of them contain changes related to review comments but only one is named like it does.
cvybhu: I'm almost ready to merge this, but it would be good for someone else to take a look as well, @piodul @psarna .
cvybhu: Apart from rebasing, one more thing I would love to see before merging is https://github.com/scylladb/scylla-rust-driver/pull/508#discussion_r952750126
The new interface gives us an iterator, but we don't test that after running out of items it returns `None`.
It isn't strictly required, but it would make the test complete and should be quick to add.

I will be on vacation next week, so I won't be able to merge it then.

Patch:
diff --git a/docs/source/load-balancing/dc-robin.md b/docs/source/load-balancing/dc-robin.md
index 825adc52..7a428a1d 100644
--- a/docs/source/load-balancing/dc-robin.md
+++ b/docs/source/load-balancing/dc-robin.md
@@ -10,6 +10,11 @@ For example if there are two datacenters:
 
 this policy when set to `us_east` will only use `A`, `B`, `C`, `A`, `B`, ...
 
+The policy may still return nodes in remote datacenters in the created plan as a fallback.
+This behaviour is configurable by the flag with a setter. By default the flag is set.
+In the example below, the flag is manually unset, which means that remote nodes
+are not present as a fallback at the end of the plan.
+
 ### Example
 To use this policy in `Session`:
 ```rust
@@ -23,11 +28,14 @@ use std::sync::Arc;
 
 let local_dc_name: String = "us_east".to_string();
 
+let mut policy = DcAwareRoundRobinPolicy::new(local_dc_name);
+policy.set_include_remote_nodes(false);
+
 let session: Session = SessionBuilder::new()
     .known_node("127.0.0.1:9042")
-    .load_balancing(Arc::new(DcAwareRoundRobinPolicy::new(local_dc_name)))
+    .load_balancing(Arc::new(policy))
     .build()
     .await?;
 # Ok(())
 # }
-```
\ No newline at end of file
+```
diff --git a/docs/source/load-balancing/token-dc-robin.md b/docs/source/load-balancing/token-dc-robin.md
index cbef84dc..f313f2f9 100644
--- a/docs/source/load-balancing/token-dc-robin.md
+++ b/docs/source/load-balancing/token-dc-robin.md
@@ -25,4 +25,4 @@ let session: Session = SessionBuilder::new()
     .await?;
 # Ok(())
 # }
-```
\ No newline at end of file
+```
diff --git a/scylla-cql/src/frame/request/batch.rs b/scylla-cql/src/frame/request/batch.rs
index cee7c66e..ac408852 100644
--- a/scylla-cql/src/frame/request/batch.rs
+++ b/scylla-cql/src/frame/request/batch.rs
@@ -1,4 +1,4 @@
-use crate::frame::frame_errors::ParseError;
+use crate::frame::{frame_errors::ParseError, value::BatchValuesIterator};
 use bytes::{BufMut, Bytes};
 use std::convert::TryInto;
 
@@ -54,9 +54,21 @@ where
         // Serializing queries
         types::write_short(self.statements_count.try_into()?, buf);
 
-        for (statement_num, statement) in self.statements.clone().enumerate() {
+        let mut n_serialized_statements = 0usize;
+        let mut value_lists = self.values.batch_values_iter();
+        for statement in self.statements.clone() {
             statement.serialize(buf)?;
-            self.values.write_nth_to_request(statement_num, buf)?;
+            value_lists.write_next_to_request(buf).ok_or_else(|| {
+                ParseError::BadDataToSerialize(
+                    "Mismatch between statement counts for batch query".to_owned(),
+                )
+            })??;
+            n_serialized_statements += 1;
+        }
+        if n_serialized_statements != self.statements_count {
+            return Err(ParseError::BadDataToSerialize(
+                "Mismatch between statement counts for batch query".to_owned(),
+            ));
         }
 
         // Serializing consistency
diff --git a/scylla-cql/src/frame/value.rs b/scylla-cql/src/frame/value.rs
index 2a85c450..7c59713f 100644
--- a/scylla-cql/src/frame/value.rs
+++ b/scylla-cql/src/frame/value.rs
@@ -99,21 +99,6 @@ pub trait ValueList {
     }
 }
 
-/// Represents List of ValueList for Batch statement
-pub trait BatchValues {
-    fn len(&self) -> usize;
-
-    fn write_nth_to_request(
-        &self,
-        n: usize,
-        buf: &mut impl BufMut,
-    ) -> Result<(), SerializeValuesError>;
-
-    fn is_empty(&self) -> bool {
-        self.len() == 0
-    }
-}
-
 impl SerializedValues {
     /// Creates empty value list
     pub const fn new() -> Self {
@@ -230,6 +215,85 @@ impl<'a> Iterator for SerializedValuesIterator<'a> {
     }
 }
 
+/// Represents List of ValueList for Batch statement
+///
+/// This trait is not implemented directly, but rather implemented through `BatchValuesGatWorkaround`
+/// (until GATs are made available in Rust)
+pub trait BatchValues: for<'r> BatchValuesGatWorkaround<'r> {}
+impl<T: for<'r> BatchValuesGatWorkaround<'r> + ?Sized> BatchValues for T {}
+
+pub trait BatchValuesGatWorkaround<'r> {
+    type BatchValuesIter: BatchValuesIterator;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter;
+    fn len(&'r self) -> usize;
+    fn is_empty(&'r self) -> bool {
+        self.len() == 0
+    }
+}
+
+/// An iterator-like for `ValueList`
+///
+/// An instance of this can be easily obtained from `IT: Iterator<Item: ValueList>`: that would be
+/// `BatchValuesIteratorFromIterator<IT>`
+///
+/// It's just essentially making methods from `ValueList` accessible instead of being an actual iterator because of
+/// several compiler limitations that would otherwise be very complex to overcome.\
+/// (specifically, types being different would require yielding enums for tuple impls, and the trait
+/// bound of `for<'r> <BatchValuesGatWorkaround<'r>::BatchValuesIter as Iterator>::Item: ValueList` is very
+/// hard to express considering several compiler limitations)
+pub trait BatchValuesIterator {
+    fn next_serialized(&mut self) -> Option<SerializedResult<'_>>;
+    fn write_next_to_request(
+        &mut self,
+        buf: &mut impl BufMut,
+    ) -> Option<Result<(), SerializeValuesError>>;
+    fn skip_next(&mut self) -> Option<()>;
+}
+
+/// Implements `BatchValuesIterator` from an `Iterator` over things that implement `ValueList`
+///
+/// Essentially used internally by this lib to provide implementors of `BatchValuesIterator` for cases
+/// that always serialize the same concrete `ValueList` type
+pub struct BatchValuesIteratorFromIterator<IT: Iterator> {
+    it: IT,
+    item_container_for_serialized: Option<IT::Item>,
+}
+
+impl<IT> BatchValuesIterator for BatchValuesIteratorFromIterator<IT>
+where
+    IT: Iterator,
+    IT::Item: ValueList,
+{
+    fn next_serialized(&mut self) -> Option<SerializedResult<'_>> {
+        self.item_container_for_serialized = self.it.next();
+        self.item_container_for_serialized
+            .as_ref()
+            .map(|vl| vl.serialized())
+    }
+    fn write_next_to_request(
+        &mut self,
+        buf: &mut impl BufMut,
+    ) -> Option<Result<(), SerializeValuesError>> {
+        self.it.next().map(|vl| vl.write_to_request(buf))
+    }
+    fn skip_next(&mut self) -> Option<()> {
+        self.it.next().map(|_| ())
+    }
+}
+
+impl<IT> From<IT> for BatchValuesIteratorFromIterator<IT>
+where
+    IT: Iterator,
+    IT::Item: ValueList,
+{
+    fn from(it: IT) -> Self {
+        BatchValuesIteratorFromIterator {
+            it,
+            item_container_for_serialized: None,
+        }
+    }
+}
+
 //
 //  Value impls
 //
@@ -805,78 +869,149 @@ impl<'b> ValueList for Cow<'b, SerializedValues> {
 // BatchValues impls
 //
 
-// Implement BatchValues for slices of ValueList types
-impl<T: ValueList> BatchValues for &[T] {
-    fn len(&self) -> usize {
-        <[T]>::len(*self)
+/// Implements `BatchValues` from an `Iterator` over things that implement `ValueList`
+///
+/// This is to avoid requiring allocating a new `Vec` containing all the `ValueList`s directly:
+/// with this, one can write:
+/// `session.batch(&batch, BatchValuesFromIterator::from(lines_to_insert.iter().map(|l| &l.value_list)))`
+/// where `lines_to_insert` may also contain e.g. data to pick the statement...
+///
+/// The underlying iterator will always be cloned at least once, once to compute the length if it can't be known
+/// in advance, and be re-cloned at every retry.
+/// It is consequently expected that the provided iterator is cheap to clone (e.g. `slice.iter().map(...)`).
+pub struct BatchValuesFromIter<IT> {
+    it: IT,
+}
+
+impl<IT> BatchValuesFromIter<IT>
+where
+    IT: Iterator + Clone,
+    IT::Item: ValueList,
+{
+    pub fn new(into_iter: impl IntoIterator<IntoIter = IT>) -> Self {
+        Self {
+            it: into_iter.into_iter(),
+        }
     }
+}
 
-    fn write_nth_to_request(
-        &self,
-        n: usize,
-        buf: &mut impl BufMut,
-    ) -> Result<(), SerializeValuesError> {
-        self[n].write_to_request(buf)?;
-        Ok(())
+impl<IT> From<IT> for BatchValuesFromIter<IT>
+where
+    IT: Iterator + Clone,
+    IT::Item: ValueList,
+{
+    fn from(it: IT) -> Self {
+        Self { it }
     }
 }
 
-// Implement BatchValues for Vec<ValueList>
-impl<T: ValueList> BatchValues for Vec<T> {
-    fn len(&self) -> usize {
-        Vec::<T>::len(self)
+impl<'r, IT> BatchValuesGatWorkaround<'r> for BatchValuesFromIter<IT>
+where
+    IT: Iterator + Clone,
+    IT::Item: ValueList,
+{
+    type BatchValuesIter = BatchValuesIteratorFromIterator<IT>;
+    fn batch_values_iter(&'r self) -> <Self as BatchValuesGatWorkaround>::BatchValuesIter {
+        self.it.clone().into()
+    }
+    fn len(&'r self) -> usize {
+        match self.it.size_hint() {
+            (l, Some(h)) if l == h => l,
+            _ => self.it.clone().count(),
+        }
     }
+}
 
-    fn write_nth_to_request(
-        &self,
-        n: usize,
-        buf: &mut impl BufMut,
-    ) -> Result<(), SerializeValuesError> {
-        self[n].write_to_request(buf)?;
-        Ok(())
+// Implement BatchValues for slices of ValueList types
+impl<'r, T: ValueList + 'r> BatchValuesGatWorkaround<'r> for [T] {
+    type BatchValuesIter = BatchValuesIteratorFromIterator<std::slice::Iter<'r, T>>;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+        self.iter().into()
+    }
+    fn len(&'r self) -> usize {
+        <[_]>::len(self)
+    }
+}
+
+// Implement BatchValues for Vec<ValueList>
+impl<'r, T: ValueList + 'r> BatchValuesGatWorkaround<'r> for Vec<T> {
+    type BatchValuesIter = BatchValuesIteratorFromIterator<std::slice::Iter<'r, T>>;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+        BatchValuesGatWorkaround::<'r>::batch_values_iter(self.as_slice())
+    }
+    fn len(&'r self) -> usize {
+        Vec::len(self)
     }
 }
 
 // Here is an example implementation for (T0, )
 // Further variants are done using a macro
-impl<T0: ValueList> BatchValues for (T0,) {
-    fn len(&self) -> usize {
+impl<'r, T0: ValueList + 'r> BatchValuesGatWorkaround<'r> for (T0,) {
+    type BatchValuesIter = BatchValuesIteratorFromIterator<std::iter::Once<&'r T0>>;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+        std::iter::once(&self.0).into()
+    }
+    fn len(&'r self) -> usize {
         1
     }
+}
 
-    fn write_nth_to_request(
-        &self,
-        n: usize,
-        buf: &mut impl BufMut,
-    ) -> Result<(), SerializeValuesError> {
-        match n {
-            0 => self.0.write_to_request(buf)?,
-            _ => panic!("Tried to serialize ValueList with an out of range index! index: {}, ValueList len: {}", n, 1),
-        };
-
-        Ok(())
-    }
+pub struct TupleValuesIter<'a, T> {
+    tuple: &'a T,
+    idx: usize,
 }
 
 macro_rules! impl_batch_values_for_tuple {
-    ( $($Ti:ident),* ; $($FieldI:tt),* ; $TupleSize:tt ) => {
-        impl<$($Ti),+> BatchValues for ($($Ti,)+)
+    ( $($Ti:ident),* ; $($FieldI:tt),* ; $TupleSize:tt) => {
+        impl<'r, $($Ti),+> BatchValuesGatWorkaround<'r> for ($($Ti,)+)
         where
-        $($Ti: ValueList),+
+            $($Ti: ValueList + 'r),+
         {
+            type BatchValuesIter = TupleValuesIter<'r, ($($Ti,)+)>;
+            fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+                TupleValuesIter {
+                    tuple: self,
+                    idx: 0,
+                }
+            }
             fn len(&self) -> usize{
                 $TupleSize
             }
-
-            fn write_nth_to_request(&self, n: usize, buf: &mut impl BufMut) -> Result<(), SerializeValuesError> {
-                match n {
+        }
+        impl<'r, $($Ti),+> BatchValuesIterator for TupleValuesIter<'r, ($($Ti,)+)>
+        where
+            $($Ti: ValueList + 'r),+
+        {
+            fn next_serialized(&mut self) -> Option<SerializedResult<'_>> {
+                let serialized_value_res = match self.idx {
                     $(
-                        $FieldI => self.$FieldI.write_to_request(buf) ?,
+                        $FieldI => self.tuple.$FieldI.serialized(),
                     )*
-                    _ => panic!("Tried to serialize ValueList with an out of range index! index: {}, ValueList len: {}", n, $TupleSize),
+                    _ => return None,
+                };
+                self.idx += 1;
+                Some(serialized_value_res)
+            }
+            fn write_next_to_request(
+                &mut self,
+                buf: &mut impl BufMut,
+            ) -> Option<Result<(), SerializeValuesError>> {
+                let ret = match self.idx {
+                    $(
+                        $FieldI => self.tuple.$FieldI.write_to_request(buf),
+                    )*
+                    _ => return None,
+                };
+                self.idx += 1;
+                Some(ret)
+            }
+            fn skip_next(&mut self) -> Option<()> {
+                if self.idx < $TupleSize {
+                    self.idx += 1;
+                    Some(())
+                } else {
+                    None
                 }
-
-                Ok(())
             }
         }
     }
@@ -906,17 +1041,75 @@ impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15; 16);
 
 // Every &impl BatchValues should also implement BatchValues
-impl<T: BatchValues> BatchValues for &T {
-    fn len(&self) -> usize {
-        <T as BatchValues>::len(*self)
+impl<'a, 'r, T: BatchValues + ?Sized> BatchValuesGatWorkaround<'r> for &'a T {
+    type BatchValuesIter = <T as BatchValuesGatWorkaround<'a>>::BatchValuesIter;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+        <T as BatchValuesGatWorkaround<'a>>::batch_values_iter(*self)
     }
+    fn len(&'r self) -> usize {
+        <T as BatchValuesGatWorkaround<'a>>::len(*self)
+    }
+}
+
+/// Allows reusing already-serialized first value
+///
+/// We'll need to build a `SerializedValues` for the first ~`ValueList` of a batch to figure out the shard (#448).
+/// Once that is done, we can use that instead of re-serializing.
+///
+/// This struct implements both `BatchValues` and `BatchValuesIterator` for that purpose
+pub struct BatchValuesFirstSerialized<'f, T> {
+    first: Option<&'f SerializedValues>,
+    rest: T,
+}
+
+impl<'f, T: BatchValues> BatchValuesFirstSerialized<'f, T> {
+    pub fn new(batch_values: T, already_serialized_first: Option<&'f SerializedValues>) -> Self {
+        Self {
+            first: already_serialized_first,
+            rest: batch_values,
+        }
+    }
+}
+
+impl<'r, 'f, BV: BatchValues> BatchValuesGatWorkaround<'r> for BatchValuesFirstSerialized<'f, BV> {
+    type BatchValuesIter =
+        BatchValuesFirstSerialized<'f, <BV as BatchValuesGatWorkaround<'r>>::BatchValuesIter>;
+    fn batch_values_iter(&'r self) -> Self::BatchValuesIter {
+        BatchValuesFirstSerialized {
+            first: self.first,
+            rest: self.rest.batch_values_iter(),
+        }
+    }
+    fn len(&'r self) -> usize {
+        self.rest.len()
+    }
+}
 
-    fn write_nth_to_request(
-        &self,
-        n: usize,
+impl<'f, IT: BatchValuesIterator> BatchValuesIterator for BatchValuesFirstSerialized<'f, IT> {
+    fn next_serialized(&mut self) -> Option<SerializedResult<'_>> {
+        match self.first.take() {
+            Some(first) => {
+                self.rest.skip_next();
+                Some(Ok(Cow::Borrowed(first)))
+            }
+            None => self.rest.next_serialized(),
+        }
+    }
+    fn write_next_to_request(
+        &mut self,
         buf: &mut impl BufMut,
-    ) -> Result<(), SerializeValuesError> {
-        <T as BatchValues>::write_nth_to_request(*self, n, buf)?;
-        Ok(())
+    ) -> Option<Result<(), SerializeValuesError>> {
+        match self.first.take() {
+            Some(first) => {
+                self.rest.skip_next();
+                first.write_to_request(buf);
+                Some(Ok(()))
+            }
+            None => self.rest.write_next_to_request(buf),
+        }
+    }
+    fn skip_next(&mut self) -> Option<()> {
+        self.rest.skip_next();
+        self.first.take().map(|_| ())
     }
 }
diff --git a/scylla/src/transport/load_balancing/dc_aware_round_robin.rs b/scylla/src/transport/load_balancing/dc_aware_round_robin.rs
index bd78aa6f..12fe489d 100644
--- a/scylla/src/transport/load_balancing/dc_aware_round_robin.rs
+++ b/scylla/src/transport/load_balancing/dc_aware_round_robin.rs
@@ -11,6 +11,7 @@ use tracing::trace;
 pub struct DcAwareRoundRobinPolicy {
     index: AtomicUsize,
     local_dc: String,
+    include_remote_nodes: bool,
 }
 
 impl DcAwareRoundRobinPolicy {
@@ -18,9 +19,14 @@ impl DcAwareRoundRobinPolicy {
         Self {
             index: AtomicUsize::new(0),
             local_dc,
+            include_remote_nodes: true,
         }
     }
 
+    pub fn set_include_remote_nodes(&mut self, val: bool) {
+        self.include_remote_nodes = val;
+    }
+
     fn is_local_node(node: &Node, local_dc: &str) -> bool {
         node.datacenter.as_deref() == Some(local_dc)
     }
@@ -60,28 +66,40 @@ impl LoadBalancingPolicy for DcAwareRoundRobinPolicy {
         let rotated_local_nodes =
             super::slice_rotated_left(local_nodes, local_nodes_rotation).cloned();
 
-        let remote_nodes = self.retrieve_remote_nodes(cluster);
-        let remote_nodes_count = cluster.all_nodes.len() - local_nodes.len();
-        let remote_nodes_rotation = super::compute_rotation(index, remote_nodes_count);
-        let rotated_remote_nodes = super::iter_rotated_left(remote_nodes, remote_nodes_rotation);
-        trace!(
-            local_nodes = rotated_local_nodes
-                .clone()
-                .map(|node| node.address.to_string())
-                .collect::<Vec<String>>()
-                .join(",")
-                .as_str(),
-            remote_nodes = rotated_remote_nodes
-                .clone()
-                .map(|node| node.address.to_string())
-                .collect::<Vec<String>>()
-                .join(",")
-                .as_str(),
-            "DC Aware"
-        );
-
-        let plan = rotated_local_nodes.chain(rotated_remote_nodes);
-        Box::new(plan)
+        if self.include_remote_nodes {
+            let remote_nodes = self.retrieve_remote_nodes(cluster);
+            let remote_nodes_count = cluster.all_nodes.len() - local_nodes.len();
+            let remote_nodes_rotation = super::compute_rotation(index, remote_nodes_count);
+            let rotated_remote_nodes =
+                super::iter_rotated_left(remote_nodes, remote_nodes_rotation);
+            trace!(
+                local_nodes = rotated_local_nodes
+                    .clone()
+                    .map(|node| node.address.to_string())
+                    .collect::<Vec<String>>()
+                    .join(",")
+                    .as_str(),
+                remote_nodes = rotated_remote_nodes
+                    .clone()
+                    .map(|node| node.address.to_string())
+                    .collect::<Vec<String>>()
+                    .join(",")
+                    .as_str(),
+                "DC Aware"
+            );
+            Box::new(rotated_local_nodes.chain(rotated_remote_nodes))
+        } else {
+            trace!(
+                local_nodes = rotated_local_nodes
+                    .clone()
+                    .map(|node| node.address.to_string())
+                    .collect::<Vec<String>>()
+                    .join(",")
+                    .as_str(),
+                "DC Aware"
+            );
+            Box::new(rotated_local_nodes)
+        }
     }
 
     fn name(&self) -> String {
@@ -103,15 +121,19 @@ impl ChildLoadBalancingPolicy for DcAwareRoundRobinPolicy {
         let local_nodes_rotation = super::compute_rotation(index, local_nodes.len());
         let rotated_local_nodes = super::slice_rotated_left(&local_nodes, local_nodes_rotation);
 
-        let remote_nodes_rotation = super::compute_rotation(index, remote_nodes.len());
-        let rotated_remote_nodes = super::slice_rotated_left(&remote_nodes, remote_nodes_rotation);
-
-        let plan = rotated_local_nodes
-            .chain(rotated_remote_nodes)
-            .cloned()
-            .collect::<Vec<_>>()
-            .into_iter();
-        Box::new(plan)
+        let plan = if self.include_remote_nodes {
+            let remote_nodes_rotation = super::compute_rotation(index, remote_nodes.len());
+            let rotated_remote_nodes =
+                super::slice_rotated_left(&remote_nodes, remote_nodes_rotation);
+
+            rotated_local_nodes
+                .chain(rotated_remote_nodes)
+                .cloned()
+                .collect::<Vec<_>>()
+        } else {
+            rotated_local_nodes.cloned().collect()
+        };
+        Box::new(plan.into_iter())
     }
 }
 
@@ -122,13 +144,12 @@ mod tests {
     use crate::transport::load_balancing::tests;
     use std::collections::HashSet;
 
-    #[tokio::test]
-    async fn test_dc_aware_round_robin_policy() {
+    async fn test_dc_aware_round_robin_policy(
+        policy: DcAwareRoundRobinPolicy,
+        expected_plans: HashSet<Vec<u16>>,
+    ) {
         let cluster = tests::mock_cluster_data_for_round_robin_tests();
 
-        let local_dc = "eu".to_string();
-        let policy = DcAwareRoundRobinPolicy::new(local_dc);
-
         let plans = (0..32)
             .map(|_| {
                 tests::get_plan_and_collect_node_identifiers(
@@ -139,6 +160,14 @@ mod tests {
             })
             .collect::<HashSet<_>>();
 
+        assert_eq!(expected_plans, plans);
+    }
+
+    #[tokio::test]
+    async fn test_dc_aware_round_robin_policy_with_remote_nodes() {
+        let local_dc = "eu".to_string();
+        let policy = DcAwareRoundRobinPolicy::new(local_dc);
+
         let expected_plans = vec![
             vec![1, 2, 3, 4, 5],
             vec![1, 2, 3, 5, 4],
@@ -150,6 +179,19 @@ mod tests {
         .into_iter()
         .collect::<HashSet<_>>();
 
-        assert_eq!(expected_plans, plans);
+        test_dc_aware_round_robin_policy(policy, expected_plans).await;
+    }
+
+    #[tokio::test]
+    async fn test_dc_aware_round_robin_policy_without_remote_nodes() {
+        let local_dc = "eu".to_string();
+        let mut policy = DcAwareRoundRobinPolicy::new(local_dc);
+        policy.set_include_remote_nodes(false);
+
+        let expected_plans = vec![vec![1, 2, 3], vec![2, 3, 1], vec![3, 1, 2]]
+            .into_iter()
+            .collect::<HashSet<_>>();
+
+        test_dc_aware_round_robin_policy(policy, expected_plans).await;
     }
 }
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 26a2e8db..002e4a7e 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -27,7 +27,9 @@ use super::topology::UntranslatedPeer;
 use crate::cql_to_rust::FromRow;
 use crate::frame::response::cql_to_rust::FromRowError;
 use crate::frame::response::result;
-use crate::frame::value::{BatchValues, SerializedValues, ValueList};
+use crate::frame::value::{
+    BatchValues, BatchValuesFirstSerialized, BatchValuesIterator, SerializedValues, ValueList,
+};
 use crate::prepared_statement::{PartitionKeyError, PreparedStatement};
 use crate::query::Query;
 use crate::routing::Token;
@@ -1006,13 +1008,43 @@ impl Session {
         batch: &Batch,
         values: impl BatchValues,
     ) -> Result<QueryResult, QueryError> {
+        // Shard-awareness behavior for batch will be to pick shard based on first batch statement's shard
+        // If users batch statements by shard, they will be rewarded with full shard awareness
+
+        // Extract first serialized_value
+        let mut batch_values_iter_for_first_serialized_value = values.batch_values_iter();
+        let first_serialized_value = batch_values_iter_for_first_serialized_value
+            .next_serialized()
+            .transpose()?;
+        let first_serialized_value = first_serialized_value.as_deref();
+        let statement_info = match (first_serialized_value, batch.statements.first()) {
+            (Some(first_serialized_value), Some(BatchStatement::PreparedStatement(ps))) => {
+                Statement {
+                    token: self.calculate_token(ps, first_serialized_value)?,
+                    keyspace: ps.get_keyspace_name(),
+                }
+            }
+            _ => Statement::default(),
+        };
+        let first_value_token = statement_info.token;
+
+        // Reuse first serialized value when serializing query, and delegate to `BatchValues::write_next_to_request`
+        // directly for others (if they weren't already serialized, possibly don't even allocate the `SerializedValues`)
+        let values = BatchValuesFirstSerialized::new(&values, first_serialized_value);
         let values_ref = &values;
 
         let run_query_result = self
             .run_query(
-                Statement::default(),
+                statement_info,
                 &batch.config,
-                |node: Arc<Node>| async move { node.random_connection().await },
+                |node: Arc<Node>| async move {
+                    match first_value_token {
+                        Some(first_value_token) => {
+                            node.connection_for_token(first_value_token).await
+                        }
+                        None => node.random_connection().await,
+                    }
+                },
                 |connection: Arc<Connection>, consistency: Consistency| async move {
                     connection
                         .batch_with_consistency(batch, values_ref, consistency)


Test patch:
diff --git a/scylla-cql/src/frame/value_tests.rs b/scylla-cql/src/frame/value_tests.rs
index 7d5a8deb..fbbf4f87 100644
--- a/scylla-cql/src/frame/value_tests.rs
+++ b/scylla-cql/src/frame/value_tests.rs
@@ -1,3 +1,5 @@
+use crate::frame::value::{BatchValuesGatWorkaround, BatchValuesIterator};
+
 use super::value::{
     BatchValues, Date, MaybeUnset, SerializeValuesError, SerializedValues, Time, Timestamp, Unset,
     Value, ValueList, ValueTooBig,
@@ -391,17 +393,21 @@ fn cow_serialized_values_value_list() {
 fn slice_batch_values() {
     let batch_values: &[&[i8]] = &[&[1, 2], &[2, 3, 4, 5], &[6]];
 
-    assert_eq!(<&[&[i8]] as BatchValues>::len(&batch_values), 3);
+    assert_eq!(
+        <&[&[i8]] as BatchValuesGatWorkaround<'_>>::len(&batch_values),
+        3
+    );
 
+    let mut it = batch_values.batch_values_iter();
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(0, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(request, vec![0, 2, 0, 0, 0, 1, 1, 0, 0, 0, 1, 2]);
     }
 
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(1, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(
             request,
             vec![0, 4, 0, 0, 0, 1, 2, 0, 0, 0, 1, 3, 0, 0, 0, 1, 4, 0, 0, 0, 1, 5]
@@ -410,26 +416,32 @@ fn slice_batch_values() {
 
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(2, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(request, vec![0, 1, 0, 0, 0, 1, 6]);
     }
+
+    assert_eq!(it.write_next_to_request(&mut Vec::new()), None);
 }
 
 #[test]
 fn vec_batch_values() {
     let batch_values: Vec<Vec<i8>> = vec![vec![1, 2], vec![2, 3, 4, 5], vec![6]];
 
-    assert_eq!(<Vec<Vec<i8>> as BatchValues>::len(&batch_values), 3);
+    assert_eq!(
+        <Vec<Vec<i8>> as BatchValuesGatWorkaround<'_>>::len(&batch_values),
+        3
+    );
 
+    let mut it = batch_values.batch_values_iter();
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(0, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(request, vec![0, 2, 0, 0, 0, 1, 1, 0, 0, 0, 1, 2]);
     }
 
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(1, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(
             request,
             vec![0, 4, 0, 0, 0, 1, 2, 0, 0, 0, 1, 3, 0, 0, 0, 1, 4, 0, 0, 0, 1, 5]
@@ -438,7 +450,7 @@ fn vec_batch_values() {
 
     {
         let mut request: Vec<u8> = Vec::new();
-        batch_values.write_nth_to_request(2, &mut request).unwrap();
+        it.write_next_to_request(&mut request).unwrap().unwrap();
         assert_eq!(request, vec![0, 1, 0, 0, 0, 1, 6]);
     }
 }
@@ -448,9 +460,10 @@ fn tuple_batch_values() {
     fn check_twoi32_tuple(tuple: impl BatchValues, size: usize) {
         assert_eq!(tuple.len(), size);
 
+        let mut it = tuple.batch_values_iter();
         for i in 0..size {
             let mut request: Vec<u8> = Vec::new();
-            tuple.write_nth_to_request(i, &mut request).unwrap();
+            it.write_next_to_request(&mut request).unwrap().unwrap();
 
             let mut expected: Vec<u8> = Vec::new();
             let i: i32 = i.try_into().unwrap();
@@ -637,12 +650,13 @@ fn tuple_batch_values() {
 fn ref_batch_values() {
     let batch_values: &[&[i8]] = &[&[1, 2], &[2, 3, 4, 5], &[6]];
 
-    assert_eq!(<&&&&[&[i8]] as BatchValues>::len(&&&&batch_values), 3);
+    assert_eq!(
+        <&&&&[&[i8]] as BatchValuesGatWorkaround<'_>>::len(&&&&batch_values),
+        3
+    );
+    let mut it = <&&&&[&[i8]] as BatchValuesGatWorkaround<'_>>::batch_values_iter(&&&&batch_values);
 
-    {
-        let mut request: Vec<u8> = Vec::new();
-        <&&&&[&[i8]] as BatchValues>::write_nth_to_request(&&&&batch_values, 0, &mut request)
-            .unwrap();
-        assert_eq!(request, vec![0, 2, 0, 0, 0, 1, 1, 0, 0, 0, 1, 2]);
-    }
+    let mut request: Vec<u8> = Vec::new();
+    it.write_next_to_request(&mut request).unwrap().unwrap();
+    assert_eq!(request, vec![0, 2, 0, 0, 0, 1, 1, 0, 0, 0, 1, 2]);
 }