Repository: scylladb/scylla-rust-driver
Pull request #567 (2022-10-04 15:17:08)

Host filter
This PR adds host filters. With host filters, it is possible to prevent some nodes from establishing connections to the cluster at all. For example, the set of nodes for which connections are allowed can be restricted to the local DC.

Tested manually with cql-stress and a 3 node cluster - applied an allowlist of nodes and observed in metrics that forbidden nodes do not have additional connections and do not receive any requests.

Fixes: https://github.com/scylladb/scylla-rust-driver/issues/561

## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [ ] 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: Let's also make an issue that documentation for `HostFilter` is missing in the book.

Patch:
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/lib.rs b/scylla/src/lib.rs
index c9b3b33c..a6e22245 100644
--- a/scylla/src/lib.rs
+++ b/scylla/src/lib.rs
@@ -116,6 +116,7 @@ pub use transport::query_result::QueryResult;
 pub use transport::session::{IntoTypedRows, Session, SessionConfig};
 pub use transport::session_builder::SessionBuilder;
 
+pub use transport::host_filter;
 pub use transport::load_balancing;
 pub use transport::retry_policy;
 pub use transport::speculative_execution;
diff --git a/scylla/src/transport/cluster.rs b/scylla/src/transport/cluster.rs
index 6e300205..befd3fea 100644
--- a/scylla/src/transport/cluster.rs
+++ b/scylla/src/transport/cluster.rs
@@ -3,6 +3,7 @@ use crate::frame::response::event::{Event, StatusChangeEvent};
 use crate::frame::value::ValueList;
 use crate::load_balancing::TokenAwarePolicy;
 use crate::routing::Token;
+use crate::transport::host_filter::HostFilter;
 use crate::transport::{
     connection::{Connection, VerifiedKeyspaceName},
     connection_pool::PoolConfig,
@@ -110,6 +111,10 @@ struct ClusterWorker {
 
     // Keyspace send in "USE <keyspace name>" when opening each connection
     used_keyspace: Option<VerifiedKeyspaceName>,
+
+    // The host filter determines towards which nodes we should open
+    // connections
+    host_filter: Option<Arc<dyn HostFilter>>,
 }
 
 #[derive(Debug)]
@@ -129,6 +134,7 @@ impl Cluster {
         pool_config: PoolConfig,
         fetch_schema_metadata: bool,
         address_translator: &Option<Arc<dyn AddressTranslator>>,
+        host_filter: &Option<Arc<dyn HostFilter>>,
     ) -> Result<Cluster, QueryError> {
         let (refresh_sender, refresh_receiver) = tokio::sync::mpsc::channel(32);
         let (use_keyspace_sender, use_keyspace_receiver) = tokio::sync::mpsc::channel(32);
@@ -141,10 +147,17 @@ impl Cluster {
             server_events_sender,
             fetch_schema_metadata,
             address_translator,
+            host_filter,
         );
 
         let metadata = metadata_reader.read_metadata(true).await?;
-        let cluster_data = ClusterData::new(metadata, &pool_config, &HashMap::new(), &None);
+        let cluster_data = ClusterData::new(
+            metadata,
+            &pool_config,
+            &HashMap::new(),
+            &None,
+            host_filter.as_deref(),
+        );
         cluster_data.wait_until_all_pools_are_initialized().await;
         let cluster_data: Arc<ArcSwap<ClusterData>> =
             Arc::new(ArcSwap::from(Arc::new(cluster_data)));
@@ -160,6 +173,8 @@ impl Cluster {
 
             use_keyspace_channel: use_keyspace_receiver,
             used_keyspace: None,
+
+            host_filter: host_filter.clone(),
         };
 
         let (fut, worker_handle) = worker.work().remote_handle();
@@ -273,6 +288,7 @@ impl ClusterData {
         pool_config: &PoolConfig,
         known_peers: &HashMap<SocketAddr, Arc<Node>>,
         used_keyspace: &Option<VerifiedKeyspaceName>,
+        host_filter: Option<&dyn HostFilter>,
     ) -> Self {
         // Create new updated known_peers and ring
         let mut new_known_peers: HashMap<SocketAddr, Arc<Node>> =
@@ -289,13 +305,17 @@ impl ClusterData {
                 Some(node) if node.datacenter == peer.datacenter && node.rack == peer.rack => {
                     node.clone()
                 }
-                _ => Arc::new(Node::new(
-                    peer.address,
-                    pool_config.clone(),
-                    peer.datacenter,
-                    peer.rack,
-                    used_keyspace.clone(),
-                )),
+                _ => {
+                    let is_enabled = host_filter.map_or(true, |f| f.accept(&peer));
+                    Arc::new(Node::new(
+                        peer.address,
+                        pool_config.clone(),
+                        peer.datacenter,
+                        peer.rack,
+                        used_keyspace.clone(),
+                        is_enabled,
+                    ))
+                }
             };
 
             new_known_peers.insert(peer.address, node.clone());
@@ -567,6 +587,7 @@ impl ClusterWorker {
             &self.pool_config,
             &cluster_data.known_peers,
             &self.used_keyspace,
+            self.host_filter.as_deref(),
         ));
 
         new_cluster_data
diff --git a/scylla/src/transport/host_filter.rs b/scylla/src/transport/host_filter.rs
new file mode 100644
index 00000000..48a547ea
--- /dev/null
+++ b/scylla/src/transport/host_filter.rs
@@ -0,0 +1,78 @@
+//! Host filters.
+//!
+//! Host filters are essentially just a predicate over
+//! [`Peer`](crate::transport::topology::Peer)s. Currently, they are used
+//! by the [`Session`](crate::transport::session::Session) to determine whether
+//! connections should be opened to a given node or not.
+
+use std::collections::HashSet;
+use std::io::Error;
+use std::net::{SocketAddr, ToSocketAddrs};
+
+use crate::transport::topology::Peer;
+
+/// The `HostFilter` trait.
+pub trait HostFilter: Send + Sync {
+    /// Returns whether a peer should be accepted or not.
+    fn accept(&self, peer: &Peer) -> bool;
+}
+
+/// Unconditionally accepts all nodes.
+pub struct AcceptAllHostFilter;
+
+impl HostFilter for AcceptAllHostFilter {
+    fn accept(&self, _peer: &Peer) -> bool {
+        true
+    }
+}
+
+/// Accepts nodes whose addresses are present in the allow list provided
+/// during filter's construction.
+pub struct AllowListHostFilter {
+    allowed: HashSet<SocketAddr>,
+}
+
+impl AllowListHostFilter {
+    /// Creates a new `AllowListHostFilter` which only accepts nodes from the
+    /// list.
+    pub fn new<I, A>(allowed_iter: I) -> Result<Self, Error>
+    where
+        I: IntoIterator<Item = A>,
+        A: ToSocketAddrs,
+    {
+        // I couldn't get the iterator combinators to work
+        let mut allowed = HashSet::new();
+        for item in allowed_iter {
+            for addr in item.to_socket_addrs()? {
+                allowed.insert(addr);
+            }
+        }
+
+        Ok(Self { allowed })
+    }
+}
+
+impl HostFilter for AllowListHostFilter {
+    fn accept(&self, peer: &Peer) -> bool {
+        self.allowed.contains(&peer.address)
+    }
+}
+
+/// Accepts nodes from given DC.
+pub struct DcHostFilter {
+    local_dc: String,
+}
+
+impl DcHostFilter {
+    /// Creates a new `DcHostFilter` that accepts nodes only from the
+    /// `local_dc`.
+    pub fn new(local_dc: String) -> Self {
+        Self { local_dc }
+    }
+}
+
+impl HostFilter for DcHostFilter {
+    fn accept(&self, peer: &Peer) -> bool {
+        peer.datacenter.as_ref() == Some(&self.local_dc)
+    }
+}
diff --git a/scylla/src/transport/load_balancing/mod.rs b/scylla/src/transport/load_balancing/mod.rs
index d8edb73d..f893933e 100644
--- a/scylla/src/transport/load_balancing/mod.rs
+++ b/scylla/src/transport/load_balancing/mod.rs
@@ -150,7 +150,7 @@ mod tests {
             keyspaces: HashMap::new(),
         };
 
-        ClusterData::new(info, &Default::default(), &HashMap::new(), &None)
+        ClusterData::new(info, &Default::default(), &HashMap::new(), &None, None)
     }
 
     pub const EMPTY_STATEMENT: Statement = Statement {
diff --git a/scylla/src/transport/load_balancing/token_aware.rs b/scylla/src/transport/load_balancing/token_aware.rs
index 373959e7..140c01c9 100644
--- a/scylla/src/transport/load_balancing/token_aware.rs
+++ b/scylla/src/transport/load_balancing/token_aware.rs
@@ -345,7 +345,7 @@ mod tests {
             keyspaces,
         };
 
-        ClusterData::new(info, &Default::default(), &HashMap::new(), &None)
+        ClusterData::new(info, &Default::default(), &HashMap::new(), &None, None)
     }
 
     // creates ClusterData with info about 8 nodes living in two different datacenters
@@ -444,7 +444,7 @@ mod tests {
             keyspaces,
         };
 
-        ClusterData::new(info, &Default::default(), &HashMap::new(), &None)
+        ClusterData::new(info, &Default::default(), &HashMap::new(), &None, None)
     }
 
     // Used as child policy for TokenAwarePolicy tests
diff --git a/scylla/src/transport/mod.rs b/scylla/src/transport/mod.rs
index c8dc0d9b..69cede77 100644
--- a/scylla/src/transport/mod.rs
+++ b/scylla/src/transport/mod.rs
@@ -3,6 +3,7 @@ mod cluster;
 pub(crate) mod connection;
 mod connection_pool;
 pub mod downgrading_consistency_retry_policy;
+pub mod host_filter;
 pub mod iterator;
 pub mod load_balancing;
 pub(crate) mod metrics;
diff --git a/scylla/src/transport/node.rs b/scylla/src/transport/node.rs
index 88b587b0..683a0166 100644
--- a/scylla/src/transport/node.rs
+++ b/scylla/src/transport/node.rs
@@ -21,7 +21,8 @@ pub struct Node {
     pub datacenter: Option<String>,
     pub rack: Option<String>,
 
-    pool: NodeConnectionPool,
+    // If the node is filtered out by the host filter, this will be None
+    pool: Option<NodeConnectionPool>,
 
     down_marker: AtomicBool,
 }
@@ -40,9 +41,11 @@ impl Node {
         datacenter: Option<String>,
         rack: Option<String>,
         keyspace_name: Option<VerifiedKeyspaceName>,
+        enabled: bool,
     ) -> Self {
-        let pool =
-            NodeConnectionPool::new(address.ip(), address.port(), pool_config, keyspace_name);
+        let pool = enabled.then(|| {
+            NodeConnectionPool::new(address.ip(), address.port(), pool_config, keyspace_name)
+        });
 
         Node {
             address,
@@ -54,7 +57,7 @@ impl Node {
     }
 
     pub fn sharder(&self) -> Option<Sharder> {
-        self.pool.sharder()
+        self.pool.as_ref()?.sharder()
     }
 
     /// Get connection which should be used to connect using given token
@@ -63,18 +66,25 @@ impl Node {
         &self,
         token: Token,
     ) -> Result<Arc<Connection>, QueryError> {
-        self.pool.connection_for_token(token)
+        self.get_pool()?.connection_for_token(token)
     }
 
     /// Get random connection
     pub(crate) async fn random_connection(&self) -> Result<Arc<Connection>, QueryError> {
-        self.pool.random_connection()
+        self.get_pool()?.random_connection()
     }
 
     pub fn is_down(&self) -> bool {
         self.down_marker.load(Ordering::Relaxed)
     }
 
+    /// Returns a boolean which indicates whether this node was is enabled.
+    /// Only enabled nodes will have connections open. For disabled nodes,
+    /// no connections will be opened.
+    pub fn is_enabled(&self) -> bool {
+        self.pool.is_some()
+    }
+
     pub(crate) fn change_down_marker(&self, is_down: bool) {
         self.down_marker.store(is_down, Ordering::Relaxed);
     }
@@ -83,15 +93,30 @@ impl Node {
         &self,
         keyspace_name: VerifiedKeyspaceName,
     ) -> Result<(), QueryError> {
-        self.pool.use_keyspace(keyspace_name).await
+        if let Some(pool) = &self.pool {
+            pool.use_keyspace(keyspace_name).await?;
+        }
+        Ok(())
     }
 
     pub(crate) fn get_working_connections(&self) -> Result<Vec<Arc<Connection>>, QueryError> {
-        self.pool.get_working_connections()
+        self.get_pool()?.get_working_connections()
     }
 
     pub(crate) async fn wait_until_pool_initialized(&self) {
-        self.pool.wait_until_initialized().await
+        if let Some(pool) = &self.pool {
+            pool.wait_until_initialized().await;
+        }
+    }
+
+    fn get_pool(&self) -> Result<&NodeConnectionPool, QueryError> {
+        self.pool.as_ref().ok_or_else(|| {
+            QueryError::IoError(Arc::new(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "No connections in the pool: the node has been disabled \
+                by the host filter",
+            )))
+        })
     }
 }
 
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 26a2e8db..4fef63fb 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;
@@ -36,6 +38,7 @@ use crate::tracing::{GetTracingConfig, TracingEvent, TracingInfo};
 use crate::transport::cluster::{Cluster, ClusterData, ClusterNeatDebug};
 use crate::transport::connection::{Connection, ConnectionConfig, VerifiedKeyspaceName};
 use crate::transport::connection_pool::PoolConfig;
+use crate::transport::host_filter::HostFilter;
 use crate::transport::iterator::{PreparedIteratorConfig, RowIterator};
 use crate::transport::load_balancing::{
     LoadBalancingPolicy, RoundRobinPolicy, Statement, TokenAwarePolicy,
@@ -203,6 +206,11 @@ pub struct SessionConfig {
 
     pub address_translator: Option<Arc<dyn AddressTranslator>>,
 
+    /// The host filter decides whether any connections should be opened
+    /// to the node or not. The driver will also avoid filtered out nodes when
+    /// re-establishing the control connection.
+    pub host_filter: Option<Arc<dyn HostFilter>>,
+
     /// If true, full schema metadata is fetched after successfully reaching a schema agreement.
     /// It is true by default but can be disabled if successive schema-altering statements should be performed.
     pub refresh_metadata_on_auto_schema_agreement: bool,
@@ -250,6 +258,7 @@ impl SessionConfig {
             auto_await_schema_agreement_timeout: Some(std::time::Duration::from_secs(60)),
             request_timeout: Some(Duration::from_secs(30)),
             address_translator: None,
+            host_filter: None,
             refresh_metadata_on_auto_schema_agreement: true,
         }
     }
@@ -436,6 +445,7 @@ impl Session {
             config.get_pool_config(),
             config.fetch_schema_metadata,
             &config.address_translator,
+            &config.host_filter,
         )
         .await?;
 
@@ -1006,13 +1016,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)
diff --git a/scylla/src/transport/session_builder.rs b/scylla/src/transport/session_builder.rs
index 163b005f..09d5e039 100644
--- a/scylla/src/transport/session_builder.rs
+++ b/scylla/src/transport/session_builder.rs
@@ -5,6 +5,7 @@ use super::load_balancing::LoadBalancingPolicy;
 use super::session::{AddressTranslator, Session, SessionConfig};
 use super::speculative_execution::SpeculativeExecutionPolicy;
 use super::Compression;
+use crate::transport::host_filter::HostFilter;
 use crate::transport::{connection_pool::PoolSize, retry_policy::RetryPolicy};
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -595,7 +596,7 @@ impl SessionBuilder {
     /// # Ok(())
     /// # }
     /// ```
-    /// # Example
+    /// /// # Example
     /// ```
     /// # use std::net::SocketAddr;
     /// # use std::sync::Arc;
@@ -603,22 +604,54 @@ impl SessionBuilder {
     /// # use std::str::FromStr;
     /// # use scylla::{Session, SessionBuilder};
     /// # use scylla::transport::session::{AddressTranslator, TranslationError};
-    /// #
+    ///
     /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
-    /// let mut translation_rules = HashMap::new();
-    /// let addr_before_translation = SocketAddr::from_str("192.168.0.42:19042").unwrap();
-    /// let addr_after_translation = SocketAddr::from_str("157.123.12.42:23203").unwrap();
-    /// translation_rules.insert(addr_before_translation, addr_after_translation);
+    ///     let mut translation_rules = HashMap::new();
+    ///     let addr_before_translation = SocketAddr::from_str("192.168.0.42").unwrap();
+    ///     let addr_after_translation = SocketAddr::from_str("157.123.12.42").unwrap();
+    ///     translation_rules.insert(addr_before_translation, addr_after_translation);
+    ///     let session: Session = SessionBuilder::new()
+    ///         .known_node("127.0.0.1:9042")
+    ///         .address_translator(Arc::new(translation_rules))
+    ///         .build()
+    ///         .await?;
+    /// #    Ok(())
+    /// # }
+    /// ```
+    pub fn address_translator(mut self, translator: Arc<dyn AddressTranslator>) -> Self {
+        self.config.address_translator = Some(translator);
+        self
+    }
+
+    /// Sets the host filter. The host filter decides whether any connections
+    /// should be opened to the node or not. The driver will also avoid
+    /// those nodes when re-establishing the control connection.
+    ///
+    /// See the [host filter](crate::transport::host_filter) module for a list
+    /// of pre-defined filters. It is also possible to provide a custom filter
+    /// by implementing the HostFilter trait.
+    ///
+    /// # Example
+    /// ```
+    /// # use async_trait::async_trait;
+    /// # use std::net::SocketAddr;
+    /// # use std::sync::Arc;
+    /// # use scylla::{Session, SessionBuilder};
+    /// # use scylla::transport::session::{AddressTranslator, TranslationError};
+    /// # use scylla::transport::host_filter::DcHostFilter;
+    ///
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// // The session will only connect to nodes from "my-local-dc"
     /// let session: Session = SessionBuilder::new()
     ///     .known_node("127.0.0.1:9042")
-    ///     .address_translator(Arc::new(translation_rules))
+    ///     .host_filter(Arc::new(DcHostFilter::new("my-local-dc".to_string())))
     ///     .build()
     ///     .await?;
-    /// #    Ok(())
+    /// # Ok(())
     /// # }
     /// ```
-    pub fn address_translator(mut self, translator: Arc<dyn AddressTranslator>) -> Self {
-        self.config.address_translator = Some(translator);
+    pub fn host_filter(mut self, filter: Arc<dyn HostFilter>) -> Self {
+        self.config.host_filter = Some(filter);
         self
     }
 
diff --git a/scylla/src/transport/topology.rs b/scylla/src/transport/topology.rs
index d3008a5b..949514ad 100644
--- a/scylla/src/transport/topology.rs
+++ b/scylla/src/transport/topology.rs
@@ -4,6 +4,7 @@ use crate::statement::query::Query;
 use crate::transport::connection::{Connection, ConnectionConfig};
 use crate::transport::connection_pool::{NodeConnectionPool, PoolConfig, PoolSize};
 use crate::transport::errors::{DbError, QueryError};
+use crate::transport::host_filter::HostFilter;
 use crate::transport::session::{AddressTranslator, IntoTypedRows};
 use crate::utils::parse::{ParseErrorCause, ParseResult, ParserState};
 
@@ -36,6 +37,7 @@ pub(crate) struct MetadataReader {
     fetch_schema: bool,
 
     address_translator: Option<Arc<dyn AddressTranslator>>,
+    host_filter: Option<Arc<dyn HostFilter>>,
 }
 
 /// Describes all metadata retrieved from the cluster
@@ -219,6 +221,7 @@ impl MetadataReader {
         server_event_sender: mpsc::Sender<Event>,
         fetch_schema: bool,
         address_translator: &Option<Arc<dyn AddressTranslator>>,
+        host_filter: &Option<Arc<dyn HostFilter>>,
     ) -> Self {
         let control_connection_address = *known_peers
             .choose(&mut thread_rng())
@@ -243,6 +246,7 @@ impl MetadataReader {
             known_peers: known_peers.into(),
             fetch_schema,
             address_translator: address_translator.clone(),
+            host_filter: host_filter.clone(),
         }
     }
 
@@ -251,6 +255,9 @@ impl MetadataReader {
         let mut result = self.fetch_metadata(initial).await;
         if let Ok(metadata) = result {
             self.update_known_peers(&metadata);
+            if initial {
+                self.handle_unaccepted_host_in_control_connection(&metadata);
+            }
             return Ok(metadata);
         }
 
@@ -302,6 +309,7 @@ impl MetadataReader {
         match &result {
             Ok(metadata) => {
                 self.update_known_peers(metadata);
+                self.handle_unaccepted_host_in_control_connection(metadata);
                 debug!("Fetched new metadata");
             }
             Err(error) => error!(
@@ -343,7 +351,67 @@ impl MetadataReader {
     }
 
     fn update_known_peers(&mut self, metadata: &Metadata) {
-        self.known_peers = metadata.peers.iter().map(|peer| peer.address).collect();
+        let host_filter = self.host_filter.as_ref();
+        self.known_peers = metadata
+            .peers
+            .iter()
+            .filter(|peer| host_filter.map_or(true, |f| f.accept(peer)))
+            .map(|peer| peer.address)
+            .collect();
+
+        // Check if the host filter isn't accidentally too restrictive,
+        // and print an error message about this fact
+        if !metadata.peers.is_empty() && self.known_peers.is_empty() {
+            error!(
+                node_ips = ?metadata
+                    .peers
+                    .iter()
+                    .map(|peer| peer.address)
+                    .collect::<Vec<_>>(),
+                "The host filter rejected all nodes in the cluster, \
+                no connections that can serve user queries have been \
+                established. The session cannot serve any queries!"
+            )
+        }
+    }
+
+    fn handle_unaccepted_host_in_control_connection(&mut self, metadata: &Metadata) {
+        let control_connection_peer = metadata
+            .peers
+            .iter()
+            .find(|peer| peer.address == self.control_connection_address);
+        if let Some(peer) = control_connection_peer {
+            if !self.host_filter.as_ref().map_or(true, |f| f.accept(peer)) {
+                warn!(
+                    filtered_node_ips = ?metadata
+                        .peers
+                        .iter()
+                        .filter(|peer| self.host_filter.as_ref().map_or(true, |p| p.accept(peer)))
+                        .map(|peer| peer.address)
+                        .collect::<Vec<_>>(),
+                    control_connection_address = ?self.control_connection_address,
+                    "The node that the control connection is established to \
+                    is not accepted by the host filter. Please verify that \
+                    the nodes in your initial peers list are accepted by the \
+                    host filter. The driver will try to re-establish the \
+                    control connection to a different node."
+                );
+
+                // Assuming here that known_peers are up-to-date
+                if !self.known_peers.is_empty() {
+                    self.control_connection_address = *self
+                        .known_peers
+                        .choose(&mut thread_rng())
+                        .expect("known_peers is empty - should be impossible");
+
+                    self.control_connection = Self::make_control_connection_pool(
+                        self.control_connection_address,
+                        self.connection_config.clone(),
+                        self.keepalive_interval,
+                    );
+                }
+            }
+        }
     }
 
     fn make_control_connection_pool(


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]);
 }
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 872977ee..6214eb7a 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -388,8 +388,7 @@ async fn test_batch() {
     batch.append_statement(&format!("INSERT INTO {}.t_batch (a, b, c) VALUES (7, 11, '')", ks)[..]);
     batch.append_statement(prepared_statement.clone());
 
-    let four_value: i32 = 4;
-    let values = ((1_i32, 2_i32, "abc"), (), (1_i32, &four_value, "hello"));
+    let values = ((1_i32, 2_i32, "abc"), (), (1_i32, 4_i32, "hello"));
 
     session.batch(&batch, values).await.unwrap();