Repository: scylladb/scylla-rust-driver
Pull request #788 (2023-08-09 08:01:50)

proxy: send shard number as feedback along with the frame 
## Motivation
It appears (e.g. in #738) that our proxy lacks ability to inform which shard was queried. This is because it is represented in a per-node, not per-shard way.

## What's done
- proxy workers are made aware of the shard they are connected to. This, of course, requires that a shards count is known, either by `ShardAwareness::FixedNum` or by querying a node;
- proxy logs are enriched with info about target shard;
- the optional feedback (triggered by a frame matching the specified criteria) is extended: previously it sent a frame through a channel, now it sends a pair (frame, shard number).

## Outcome
- while using proxy manually, it is easier to trace what shard is queried;
- it is possible now to write automated tests asserting that a particular shard is queried. 

## 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.
- ~~[ ] I have provided docstrings for the public items that I want to introduce.~~
- ~~[ ] I have adjusted the documentation in `./docs/source/`.~~
- ~~[ ] I added appropriate `Fixes:` annotations to PR description.~~


avelanarius: Would it make sense to add a new test "showcasing" the new functionality? For example some shard awareness test that validates we always send a request to partition X to the same shard Y?
wprzytula: > Would it make sense to add a new test "showcasing" the new functionality? For example some shard awareness test that validates we always send a request to partition X to the same shard Y?

It's a good idea both overall, and your particular example seems to be a well-suited showcase.
wprzytula: > Would it make sense to add a new test "showcasing" the new functionality? For example some shard awareness test that validates we always send a request to partition X to the same shard Y?

I've added the showcase test.
wprzytula: Fixed flakiness caused by some replicas never being queried (due to randomness in default policy), which resulted in 0 shards queried (as opposed to expected 1). Now 0 or 1 are accepted.

Patch:
diff --git a/scylla-proxy/src/actions.rs b/scylla-proxy/src/actions.rs
index ef0945ba..d571afd4 100644
--- a/scylla-proxy/src/actions.rs
+++ b/scylla-proxy/src/actions.rs
@@ -4,8 +4,9 @@ use bytes::Bytes;
 use rand::{Rng, RngCore};
 use tokio::sync::mpsc;
 
-use crate::frame::{
-    FrameOpcode, FrameParams, RequestFrame, RequestOpcode, ResponseFrame, ResponseOpcode,
+use crate::{
+    frame::{FrameOpcode, FrameParams, RequestFrame, RequestOpcode, ResponseFrame, ResponseOpcode},
+    TargetShard,
 };
 use scylla_cql::{
     errors::{DbError, WriteType},
@@ -171,7 +172,10 @@ pub trait Reaction: Sized {
 
     /// Adds sending the matching frame as feedback using the provided channel.
     /// Modifies the existing `Reaction`.
-    fn with_feedback_when_performed(self, tx: mpsc::UnboundedSender<Self::Incoming>) -> Self;
+    fn with_feedback_when_performed(
+        self,
+        tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
+    ) -> Self;
 }
 
 fn fmt_reaction(
@@ -202,7 +206,7 @@ pub struct RequestReaction {
     pub to_addressee: Option<Action<RequestFrame, RequestFrame>>,
     pub to_sender: Option<Action<RequestFrame, ResponseFrame>>,
     pub drop_connection: Option<Option<Duration>>,
-    pub feedback_channel: Option<mpsc::UnboundedSender<RequestFrame>>,
+    pub feedback_channel: Option<mpsc::UnboundedSender<(RequestFrame, Option<TargetShard>)>>,
 }
 
 impl fmt::Debug for RequestReaction {
@@ -223,7 +227,7 @@ pub struct ResponseReaction {
     pub to_addressee: Option<Action<ResponseFrame, ResponseFrame>>,
     pub to_sender: Option<Action<ResponseFrame, RequestFrame>>,
     pub drop_connection: Option<Option<Duration>>,
-    pub feedback_channel: Option<mpsc::UnboundedSender<ResponseFrame>>,
+    pub feedback_channel: Option<mpsc::UnboundedSender<(ResponseFrame, Option<TargetShard>)>>,
 }
 
 impl fmt::Debug for ResponseReaction {
@@ -321,7 +325,10 @@ impl Reaction for RequestReaction {
         }
     }
 
-    fn with_feedback_when_performed(self, tx: mpsc::UnboundedSender<Self::Incoming>) -> Self {
+    fn with_feedback_when_performed(
+        self,
+        tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
+    ) -> Self {
         Self {
             feedback_channel: Some(tx),
             ..self
@@ -662,7 +669,10 @@ impl Reaction for ResponseReaction {
         }
     }
 
-    fn with_feedback_when_performed(self, tx: mpsc::UnboundedSender<Self::Incoming>) -> Self {
+    fn with_feedback_when_performed(
+        self,
+        tx: mpsc::UnboundedSender<(Self::Incoming, Option<TargetShard>)>,
+    ) -> Self {
         Self {
             feedback_channel: Some(tx),
             ..self
diff --git a/scylla-proxy/src/lib.rs b/scylla-proxy/src/lib.rs
index 8e91ae97..5a4a12f5 100644
--- a/scylla-proxy/src/lib.rs
+++ b/scylla-proxy/src/lib.rs
@@ -2,6 +2,9 @@ mod actions;
 mod errors;
 mod frame;
 mod proxy;
+
+pub type TargetShard = u16;
+
 pub use actions::{
     Action, Condition, Reaction, RequestReaction, RequestRule, ResponseReaction, ResponseRule,
 };
diff --git a/scylla-proxy/src/proxy.rs b/scylla-proxy/src/proxy.rs
index 31a9c0cd..ae3eea3c 100644
--- a/scylla-proxy/src/proxy.rs
+++ b/scylla-proxy/src/proxy.rs
@@ -3,7 +3,7 @@ use crate::errors::{DoorkeeperError, ProxyError, WorkerError};
 use crate::frame::{
     self, read_response_frame, write_frame, FrameOpcode, FrameParams, RequestFrame, ResponseFrame,
 };
-use crate::RequestOpcode;
+use crate::{RequestOpcode, TargetShard};
 use bytes::Bytes;
 use scylla_cql::frame::types::read_string_multimap;
 use std::collections::HashMap;
@@ -206,6 +206,18 @@ impl Display for DisplayableRealAddrOption {
     }
 }
 
+#[derive(Clone, Copy)]
+struct DisplayableShard(Option<TargetShard>);
+impl Display for DisplayableShard {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Some(shard) = self.0 {
+            write!(f, "shard {}", shard)
+        } else {
+            write!(f, "unknown shard")
+        }
+    }
+}
+
 enum InternalNode {
     Real {
         real_addr: SocketAddr,
@@ -600,6 +612,7 @@ impl Doorkeeper {
         connection_no: usize,
         driver_stream: TcpStream,
         cluster_stream: Option<TcpStream>,
+        shard: Option<TargetShard>,
     ) {
         let (driver_read, driver_write) = driver_stream.into_split();
 
@@ -611,6 +624,7 @@ impl Doorkeeper {
             driver_addr,
             real_addr: self.node.real_addr(),
             proxy_addr: self.node.proxy_addr(),
+            shard,
         };
 
         let (tx_request, rx_request) = mpsc::unbounded_channel::<RequestFrame>();
@@ -667,11 +681,13 @@ impl Doorkeeper {
         connection_no: usize,
     ) -> Result<(), DoorkeeperError> {
         let (driver_stream, driver_addr) = self.make_driver_stream(connection_no).await?;
-        let cluster_stream = match self.node {
+        let (cluster_stream, shard) = match self.node {
             InternalNode::Real { real_addr, .. } => {
-                Some(self.make_cluster_stream(driver_addr, real_addr).await?)
+                let (cluster_stream, shard) =
+                    self.make_cluster_stream(driver_addr, real_addr).await?;
+                (Some(cluster_stream), shard)
             }
-            InternalNode::Simulated { .. } => None,
+            InternalNode::Simulated { .. } => (None, None),
         };
 
         self.spawn_workers(
@@ -680,6 +696,7 @@ impl Doorkeeper {
             connection_no,
             driver_stream,
             cluster_stream,
+            shard,
         )
         .await;
 
@@ -707,8 +724,8 @@ impl Doorkeeper {
         &mut self,
         driver_addr: SocketAddr,
         real_addr: SocketAddr,
-    ) -> Result<TcpStream, DoorkeeperError> {
-        let cluster_stream = if let Some(shards) = self.shards_count {
+    ) -> Result<(TcpStream, Option<TargetShard>), DoorkeeperError> {
+        let mut cluster_stream = if let Some(shards) = self.shards_count {
             let socket = match self.node.proxy_addr().ip() {
                 std::net::IpAddr::V4(_) => TcpSocket::new_v4(),
                 std::net::IpAddr::V6(_) => TcpSocket::new_v6(),
@@ -731,7 +748,7 @@ impl Doorkeeper {
 
             socket.connect(real_addr).await.map(|ok| {
                 info!(
-                    "Connected to the cluster from {} at {}, shard {}.",
+                    "Connected to the cluster from {} at {}, intended shard {}.",
                     ok.local_addr().unwrap(),
                     real_addr,
                     shard_preserving_addr.port() % shards
@@ -746,32 +763,53 @@ impl Doorkeeper {
         }
         .map_err(|err| DoorkeeperError::NodeConnectionAttempt(real_addr, err))?;
 
-        Ok(cluster_stream)
+        // If ShardAwareness is aware (QueryNode or FixedNum variants) and the
+        // proxy succeeded to know the shards count (in FixedNum we get it for
+        // free, in QueryNode the initial Options query succceeded and Supported
+        // contained SCYLLA_SHARDS_NUM), then upon opening each connection to the
+        // node, the proxy issues another Options requests and acknowledges the
+        // shard it got connected to.
+        let shard = if self.shards_count.is_some() {
+            self.obtain_shard_number(real_addr, &mut cluster_stream)
+                .await?
+        } else {
+            None
+        };
+
+        Ok((cluster_stream, shard))
     }
 
     fn next_port_to_same_shard(&self, port: u16) -> u16 {
         port.wrapping_add(self.shards_count.unwrap())
     }
 
-    async fn obtain_shards_count(&self, real_addr: SocketAddr) -> Result<u16, DoorkeeperError> {
-        let mut connection = TcpStream::connect(real_addr)
-            .await
-            .map_err(|err| DoorkeeperError::NodeConnectionAttempt(real_addr, err))?;
+    async fn get_supported_options(
+        connection: &mut TcpStream,
+    ) -> Result<HashMap<String, Vec<String>>, DoorkeeperError> {
         write_frame(
             HARDCODED_OPTIONS_PARAMS,
             FrameOpcode::Request(RequestOpcode::Options),
             &Bytes::new(),
-            &mut connection,
+            connection,
         )
         .await
         .map_err(DoorkeeperError::ObtainingShardNumber)?;
 
-        let supported_frame = read_response_frame(&mut connection)
+        let supported_frame = read_response_frame(connection)
             .await
             .map_err(DoorkeeperError::ObtainingShardNumberFrame)?;
 
         let options = read_string_multimap(&mut supported_frame.body.as_ref())
             .map_err(DoorkeeperError::ObtainingShardNumberParseOptions)?;
+
+        Ok(options)
+    }
+
+    async fn obtain_shards_count(&self, real_addr: SocketAddr) -> Result<u16, DoorkeeperError> {
+        let mut connection = TcpStream::connect(real_addr)
+            .await
+            .map_err(|err| DoorkeeperError::NodeConnectionAttempt(real_addr, err))?;
+        let options = Self::get_supported_options(&mut connection).await?;
         let nr_shards_entry = options.get("SCYLLA_NR_SHARDS");
         let shards = match nr_shards_entry
             .and_then(|vec| vec.first())
@@ -785,6 +823,24 @@ impl Doorkeeper {
         info!("Obtained shards number on node {}: {}", real_addr, shards);
         Ok(shards)
     }
+
+    async fn obtain_shard_number(
+        &self,
+        real_addr: SocketAddr,
+        connection: &mut TcpStream,
+    ) -> Result<Option<TargetShard>, DoorkeeperError> {
+        let options = Self::get_supported_options(connection).await?;
+        let shard_entry = options.get("SCYLLA_SHARD");
+        let shard = shard_entry
+            .and_then(|vec| vec.first())
+            .map(|s| {
+                s.parse::<u16>()
+                    .map_err(DoorkeeperError::ObtainingShardNumberParseShardNumber)
+            })
+            .transpose()?;
+        info!("Connected to node {}, shard {:?}", real_addr, shard);
+        Ok(shard)
+    }
 }
 
 struct ProxyWorker {
@@ -795,15 +851,17 @@ struct ProxyWorker {
     driver_addr: SocketAddr,
     real_addr: Option<SocketAddr>,
     proxy_addr: SocketAddr,
+    shard: Option<TargetShard>,
 }
 
 impl ProxyWorker {
     fn exit(self, duty: &'static str) {
         debug!(
-            "Worker exits: [driver: {}, proxy: {}, node: {}]::{}.",
+            "Worker exits: [driver: {}, proxy: {}, node: {}, {}]::{}.",
             self.driver_addr,
             self.proxy_addr,
             DisplayableRealAddrOption(self.real_addr),
+            DisplayableShard(self.shard),
             duty
         );
         std::mem::drop(self.finish_guard);
@@ -834,6 +892,7 @@ impl ProxyWorker {
         mut read_half: (impl AsyncRead + Unpin),
         request_processor_tx: mpsc::UnboundedSender<RequestFrame>,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted(
             "receiver_from_driver",
             |driver_addr, proxy_addr, _real_addr| async move {
@@ -846,8 +905,11 @@ impl ProxyWorker {
                         })?;
 
                     debug!(
-                        "Intercepted Driver ({}) -> Cluster ({}) frame. opcode: {:?}.",
-                        driver_addr, proxy_addr, &frame.opcode
+                        "Intercepted Driver ({}) -> Cluster ({}) ({}) frame. opcode: {:?}.",
+                        driver_addr,
+                        proxy_addr,
+                        DisplayableShard(shard),
+                        &frame.opcode
                     );
                     if request_processor_tx.send(frame).is_err() {
                         warn!("request_processor had exited.");
@@ -864,6 +926,7 @@ impl ProxyWorker {
         mut read_half: (impl AsyncRead + Unpin),
         response_processor_tx: mpsc::UnboundedSender<ResponseFrame>,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted(
             "receiver_from_cluster",
             |driver_addr, _proxy_addr, real_addr| async move {
@@ -878,8 +941,11 @@ impl ProxyWorker {
                             })?;
 
                     debug!(
-                        "Intercepted Cluster ({}) -> Driver ({}) frame. opcode: {:?}.",
-                        real_addr, driver_addr, &frame.opcode
+                        "Intercepted Cluster ({}) -> Driver ({}) ({}) frame. opcode: {:?}.",
+                        real_addr,
+                        driver_addr,
+                        DisplayableShard(shard),
+                        &frame.opcode
                     );
 
                     if response_processor_tx.send(frame).is_err() {
@@ -899,6 +965,7 @@ impl ProxyWorker {
         mut connection_close_notifier: ConnectionCloseNotifier,
         mut terminate_notifier: TerminateNotifier,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted(
             "sender_to_driver",
             |driver_addr, proxy_addr, _real_addr| async move {
@@ -916,8 +983,11 @@ impl ProxyWorker {
                     };
 
                     debug!(
-                        "Sending Proxy ({}) -> Driver ({}) frame. opcode: {:?}.",
-                        proxy_addr, driver_addr, &response.opcode
+                        "Sending Proxy ({}) -> Driver ({}) ({}) frame. opcode: {:?}.",
+                        proxy_addr,
+                        driver_addr,
+                        DisplayableShard(shard),
+                        &response.opcode
                     );
                     if response.write(&mut write_half).await.is_err() {
                         if terminate_notifier.try_recv().is_err()
@@ -941,6 +1011,7 @@ impl ProxyWorker {
         mut connection_close_notifier: ConnectionCloseNotifier,
         mut terminate_notifier: TerminateNotifier,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted(
             "sender_to_driver",
             |_driver_addr, proxy_addr, real_addr| async move {
@@ -959,8 +1030,11 @@ impl ProxyWorker {
                     };
 
                     debug!(
-                        "Sending Proxy ({}) -> Cluster ({}) frame. opcode: {:?}.",
-                        proxy_addr, real_addr, &request.opcode
+                        "Sending Proxy ({}) -> Cluster ({}) ({}) frame. opcode: {:?}.",
+                        proxy_addr,
+                        real_addr,
+                        DisplayableShard(shard),
+                        &request.opcode
                     );
 
                     if request.write(&mut write_half).await.is_err() {
@@ -987,6 +1061,7 @@ impl ProxyWorker {
         request_rules: Arc<Mutex<Vec<RequestRule>>>,
         connection_close_signaler: ConnectionCloseSignaler,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted("request_processor", |driver_addr, _, real_addr| async move {
             'mainloop: loop {
                 match requests_rx.recv().await {
@@ -999,13 +1074,13 @@ impl ProxyWorker {
                         let mut guard = request_rules.lock().unwrap();
                         '_ruleloop: for (i, request_rule) in guard.iter_mut().enumerate() {
                             if request_rule.0.eval(&ctx) {
-                                info!("Applying rule no={} to request ({} -> {}).", i, driver_addr, DisplayableRealAddrOption(real_addr));
+                                info!("Applying rule no={} to request ({} -> {} ({})).", i, driver_addr, DisplayableRealAddrOption(real_addr), DisplayableShard(shard));
                                 debug!("-> Applied rule: {:?}", request_rule);
                                 debug!("-> To request: {:?}", ctx.opcode);
                                 trace!("{:?}", request);
 
                                 if let Some(ref tx) = request_rule.1.feedback_channel {
-                                    tx.send(request.clone()).unwrap_or_else(|err|
+                                    tx.send((request.clone(), shard)).unwrap_or_else(|err|
                                         warn!("Could not send received request as feedback: {}", err)
                                     );
                                 }
@@ -1055,8 +1130,11 @@ impl ProxyWorker {
                                         }
                                         // close connection.
                                         info!(
-                                            "Dropping connection between {} and {} (as requested by a proxy rule)!"
-                                        , driver_addr, DisplayableRealAddrOption(real_addr));
+                                            "Dropping connection between {} and {} ({}) (as requested by a proxy rule)!",
+                                            driver_addr,
+                                            DisplayableRealAddrOption(real_addr),
+                                            DisplayableShard(shard),
+                                        );
                                         let _ = connection_close_signaler_clone.send(());
                                     }
                                 };
@@ -1086,6 +1164,7 @@ impl ProxyWorker {
         response_rules: Arc<Mutex<Vec<ResponseRule>>>,
         connection_close_signaler: ConnectionCloseSignaler,
     ) {
+        let shard = self.shard;
         self.run_until_interrupted("request_processor", |driver_addr, _, real_addr| async move {
             'mainloop: loop {
                 match responses_rx.recv().await {
@@ -1098,13 +1177,13 @@ impl ProxyWorker {
                         let mut guard = response_rules.lock().unwrap();
                         '_ruleloop: for (i, response_rule) in guard.iter_mut().enumerate() {
                             if response_rule.0.eval(&ctx) {
-                                info!("Applying rule no={} to request ({} -> {}).", i, DisplayableRealAddrOption(real_addr), driver_addr);
+                                info!("Applying rule no={} to request ({} -> {} ({})).", i, DisplayableRealAddrOption(real_addr), driver_addr, DisplayableShard(shard));
                                 debug!("-> Applied rule: {:?}", response_rule);
                                 debug!("-> To response: {:?}", ctx.opcode);
                                 trace!("{:?}", response);
 
                                 if let Some(ref tx) = response_rule.1.feedback_channel {
-                                    tx.send(response.clone()).unwrap_or_else(|err| warn!(
+                                    tx.send((response.clone(), shard)).unwrap_or_else(|err| warn!(
                                         "Could not send received response as feedback: {}", err
                                     ));
                                 }
@@ -1154,8 +1233,11 @@ impl ProxyWorker {
                                         }
                                         // close connection.
                                         info!(
-                                            "Dropping connection between {} and {} (as requested by a proxy rule)!"
-                                        , driver_addr, real_addr.expect("BUG: response rules are unavailable for dry-mode proxy!"));
+                                            "Dropping connection between {} and {} ({}) (as requested by a proxy rule)!",
+                                            driver_addr,
+                                            real_addr.expect("BUG: response rules are unavailable for dry-mode proxy!"),
+                                            DisplayableShard(shard)
+                                        );
                                         let _ = connection_close_signaler_clone.send(());
                                     }
                                 };
@@ -1203,7 +1285,7 @@ mod tests {
     use crate::{Condition, Reaction as _, RequestReaction, ResponseOpcode, ResponseReaction};
     use assert_matches::assert_matches;
     use bytes::{BufMut, BytesMut};
-    use futures::future::join;
+    use futures::future::{join, join3};
     use rand::RngCore;
     use scylla_cql::frame::frame_errors::FrameError;
     use scylla_cql::frame::types::write_string_multimap;
@@ -1230,23 +1312,21 @@ mod tests {
         body.freeze()
     }
 
-    async fn respond_with_shards_number(mut conn: TcpStream, shards_num: u16) {
+    async fn respond_with_supported(
+        conn: &mut TcpStream,
+        supported_options: &HashMap<String, Vec<String>>,
+    ) {
         let RequestFrame {
             params: recvd_params,
             opcode: recvd_opcode,
             body: recvd_body,
-        } = read_request_frame(&mut conn).await.unwrap();
+        } = read_request_frame(conn).await.unwrap();
         assert_eq!(recvd_params, HARDCODED_OPTIONS_PARAMS);
         assert_eq!(recvd_opcode, RequestOpcode::Options);
         assert_eq!(recvd_body, Bytes::new()); // body should be empty
 
         let mut body = BytesMut::new();
-        let mut sharded_info = HashMap::new();
-        sharded_info.insert(
-            String::from("SCYLLA_NR_SHARDS"),
-            vec![shards_num.to_string()],
-        );
-        write_string_multimap(&sharded_info, &mut body).unwrap();
+        write_string_multimap(supported_options, &mut body).unwrap();
 
         let body = body.freeze();
 
@@ -1254,12 +1334,35 @@ mod tests {
             HARDCODED_OPTIONS_PARAMS.for_response(),
             FrameOpcode::Response(ResponseOpcode::Supported),
             &body,
-            &mut conn,
+            conn,
         )
         .await
         .unwrap();
     }
 
+    fn supported_shards_count(shards_count: u16) -> HashMap<String, Vec<String>> {
+        let mut sharded_info = HashMap::new();
+        sharded_info.insert(
+            String::from("SCYLLA_NR_SHARDS"),
+            vec![shards_count.to_string()],
+        );
+        sharded_info
+    }
+
+    fn supported_shard_number(shard_num: TargetShard) -> HashMap<String, Vec<String>> {
+        let mut sharded_info = HashMap::new();
+        sharded_info.insert(String::from("SCYLLA_SHARD"), vec![shard_num.to_string()]);
+        sharded_info
+    }
+
+    async fn respond_with_shards_count(conn: &mut TcpStream, shards_count: u16) {
+        respond_with_supported(conn, &supported_shards_count(shards_count)).await;
+    }
+
+    async fn respond_with_shard_num(conn: &mut TcpStream, shard_num: TargetShard) {
+        respond_with_supported(conn, &supported_shard_number(shard_num)).await;
+    }
+
     fn next_local_address_with_port(port: u16) -> SocketAddr {
         SocketAddr::new(get_exclusive_local_address(), port)
     }
@@ -1287,7 +1390,7 @@ mod tests {
 
         let body = random_body();
 
-        let mock_driver_action = async {
+        let send_frame_to_shard = async {
             let mut conn = TcpStream::connect(node1_proxy_addr).await.unwrap();
 
             write_frame(params, opcode, &body, &mut conn).await.unwrap();
@@ -1296,9 +1399,13 @@ mod tests {
 
         let mock_node_action = async {
             if let ShardAwareness::QueryNode = shard_awareness {
-                respond_with_shards_number(mock_node_listener.accept().await.unwrap().0, 1).await;
+                respond_with_shards_count(&mut mock_node_listener.accept().await.unwrap().0, 1)
+                    .await;
             }
             let (mut conn, _) = mock_node_listener.accept().await.unwrap();
+            if shard_awareness.is_aware() {
+                respond_with_shard_num(&mut conn, 1).await;
+            }
             let RequestFrame {
                 params: recvd_params,
                 opcode: recvd_opcode,
@@ -1311,7 +1418,7 @@ mod tests {
         };
 
         // we keep the connections open until proxy finishes to let it perform clean exit with no disconnects
-        let (_node_conn, _driver_conn) = join(mock_node_action, mock_driver_action).await;
+        let (_node_conn, _driver_conn) = join(mock_node_action, send_frame_to_shard).await;
         running_proxy.finish().await.unwrap();
     }
 
@@ -1347,7 +1454,7 @@ mod tests {
 
             let (driver_addr_tx, driver_addr_rx) = oneshot::channel::<SocketAddr>();
 
-            let mock_driver_action = async {
+            let send_frame_to_shard = async {
                 let socket = TcpSocket::new_v4().unwrap();
                 socket
                     .bind(SocketAddr::from_str("0.0.0.0:0").unwrap())
@@ -1368,7 +1475,7 @@ mod tests {
             };
 
             // we keep the connections open until proxy finishes to let it perform clean exit with no disconnects
-            let (_node_conn, _driver_conn) = join(mock_node_action, mock_driver_action).await;
+            let (_node_conn, _driver_conn) = join(mock_node_action, send_frame_to_shard).await;
             running_proxy.finish().await.unwrap();
         }
 
@@ -1398,7 +1505,7 @@ mod tests {
                 let (driver_addr_tx, driver_addr_rx) = oneshot::channel::<SocketAddr>();
 
                 let mock_driver_addr = next_local_address_with_port(shards_num * 1234 + shard_num);
-                let mock_driver_action = async {
+                let send_frame_to_shard = async {
                     let socket = TcpSocket::new_v4().unwrap();
                     socket
                         .bind(mock_driver_addr)
@@ -1408,8 +1515,8 @@ mod tests {
                 };
 
                 let mock_node_action = async {
-                    respond_with_shards_number(
-                        mock_node_listener.accept().await.unwrap().0,
+                    respond_with_shards_count(
+                        &mut mock_node_listener.accept().await.unwrap().0,
                         shards_num,
                     )
                     .await;
@@ -1422,7 +1529,7 @@ mod tests {
                     conn
                 };
 
-                let (_node_conn, _driver_conn) = join(mock_node_action, mock_driver_action).await;
+                let (_node_conn, _driver_conn) = join(mock_node_action, send_frame_to_shard).await;
                 running_proxy.finish().await.unwrap();
             }
         }
@@ -1507,7 +1614,7 @@ mod tests {
             body.freeze()
         };
 
-        let mock_driver_action = async {
+        let send_frame_to_shard = async {
             let mut conn = TcpStream::connect(node1_proxy_addr).await.unwrap();
 
             write_frame(params1, opcode1, &body1, &mut conn)
@@ -1555,7 +1662,7 @@ mod tests {
             conn
         };
 
-        let (mut node_conn, mut driver_conn) = join(mock_node_action, mock_driver_action).await;
+        let (mut node_conn, mut driver_conn) = join(mock_node_action, send_frame_to_shard).await;
 
         running_proxy.finish().await.unwrap();
 
@@ -1793,7 +1900,7 @@ mod tests {
 
         let body = random_body();
 
-        let mock_driver_action = async {
+        let send_frame_to_shard = async {
             let mut conn = TcpStream::connect(node1_proxy_addr).await.unwrap();
             write_frame(params, request_opcode, &body, &mut conn)
                 .await
@@ -1810,16 +1917,16 @@ mod tests {
         };
 
         // we keep the connections open until proxy finishes to let it perform clean exit with no disconnects
-        let (_node_conn, _driver_conn) = join(mock_node_action, mock_driver_action).await;
+        let (_node_conn, _driver_conn) = join(mock_node_action, send_frame_to_shard).await;
 
-        let feedback_request = request_feedback_rx.recv().await.unwrap();
+        let (feedback_request, _shard) = request_feedback_rx.recv().await.unwrap();
         assert_eq!(feedback_request.params, params);
         assert_eq!(
             FrameOpcode::Request(feedback_request.opcode),
             request_opcode
         );
         assert_eq!(feedback_request.body, body);
-        let feedback_response = response_feedback_rx.recv().await.unwrap();
+        let (feedback_response, _shard) = response_feedback_rx.recv().await.unwrap();
         assert_eq!(feedback_response.params, params.for_response());
         assert_eq!(
             FrameOpcode::Response(feedback_response.opcode),
@@ -1846,7 +1953,7 @@ mod tests {
 
         let mock_node_listener = TcpListener::bind(node1_real_addr).await.unwrap();
 
-        let mock_driver_action = async {
+        let send_frame_to_shard = async {
             let mut conn = TcpStream::connect(node1_proxy_addr).await.unwrap();
 
             conn.write_all(b"uselessJunk").await.unwrap();
@@ -1858,7 +1965,7 @@ mod tests {
             conn
         };
 
-        let (node_conn, driver_conn) = join(mock_node_action, mock_driver_action).await;
+        let (node_conn, driver_conn) = join(mock_node_action, send_frame_to_shard).await;
 
         running_proxy.sanity_check().unwrap();
 
@@ -1921,7 +2028,7 @@ mod tests {
 
         let body2 = random_body();
 
-        let mock_driver_action = async {
+        let send_frame_to_shard = async {
             let mut conn = TcpStream::connect(node1_proxy_addr).await.unwrap();
 
             write_frame(params1, opcode1, &body1, &mut conn)
@@ -1948,7 +2055,7 @@ mod tests {
 
         // we keep the connections open until proxy finishes to let it perform clean exit with no disconnects
         let (_node_conn, _driver_conn) =
-            tokio::time::timeout(delay, join(mock_node_action, mock_driver_action))
+            tokio::time::timeout(delay, join(mock_node_action, send_frame_to_shard))
                 .await
                 .expect("Request processing was not concurrent");
         running_proxy.finish().await.unwrap();
@@ -2081,4 +2188,155 @@ mod tests {
 
         assert_matches!(conn.read(&mut [0u8; 1]).await, Ok(0));
     }
+
+    // The test asserts that once a (mock) driver connects to the proxy from some port,
+    // the proxy will connect to a shard corresponding to that port and that the target
+    // shard number will be sent through the feedback channel.
+    #[tokio::test]
+    #[ntest::timeout(1000)]
+    async fn proxy_reports_target_shard_as_feedback() {
+        init_logger();
+
+        let node_port = 10101;
+        let node_real_addr = next_local_address_with_port(node_port);
+        let mock_node_listener = TcpListener::bind(node_real_addr).await.unwrap();
+
+        let params = FrameParams {
+            flags: 0,
+            version: 0x04,
+            stream: 0,
+        };
+        let request_opcode = FrameOpcode::Request(RequestOpcode::Options);
+        let response_opcode = FrameOpcode::Response(ResponseOpcode::Ready);
+
+        let body = random_body();
+
+        for shards_count in 2..9 {
+            // Two driver connections are simulated, each to a different shard.
+            let driver1_shard = shards_count - 1;
+            let driver2_shard = shards_count - 2;
+            let node_proxy_addr = next_local_address_with_port(node_port);
+
+            let (request_feedback_tx, mut request_feedback_rx) = mpsc::unbounded_channel();
+            let (response_feedback_tx, mut response_feedback_rx) = mpsc::unbounded_channel();
+
+            let proxy = Proxy::new([Node::new(
+                node_real_addr,
+                node_proxy_addr,
+                ShardAwareness::FixedNum(shards_count),
+                Some(vec![RequestRule(
+                    Condition::True,
+                    RequestReaction::drop_frame().with_feedback_when_performed(request_feedback_tx),
+                )]),
+                Some(vec![ResponseRule(
+                    Condition::True,
+                    ResponseReaction::drop_frame()
+                        .with_feedback_when_performed(response_feedback_tx),
+                )]),
+            )]);
+            let running_proxy = proxy.run().await.unwrap();
+
+            /// Choose a source port `p` such that `shard == shard_of_source_port(p)`.
+            fn draw_source_port_for_shard(shards_count: u16, shard: u16) -> u16 {
+                assert!(shard < shards_count);
+                (49152 + shards_count - 1) / shards_count * shards_count + shard
+            }
+
+            async fn bind_socket_for_shard(shards_count: u16, shard: u16) -> TcpSocket {
+                let socket = TcpSocket::new_v4().unwrap();
+                let initial_port = draw_source_port_for_shard(shards_count, shard);
+
+                let mut desired_addr =
+                    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), initial_port);
+                while socket.bind(desired_addr).is_err() {
+                    // in search for a port that translates to the desired shard
+                    let next_port = desired_addr.port().wrapping_add(shards_count);
+                    if next_port == initial_port {
+                        panic!("No more ports left");
+                    }
+                    desired_addr.set_port(next_port);
+                }
+
+                socket
+            }
+
+            let body_ref = &body;
+            let send_frame_to_shard = |driver_shard: u16| async move {
+                let socket = bind_socket_for_shard(shards_count, driver_shard).await;
+                let mut conn = socket.connect(node_proxy_addr).await.unwrap();
+
+                write_frame(params, request_opcode, body_ref, &mut conn)
+                    .await
+                    .unwrap();
+                conn
+            };
+
+            let mock_driver1_action = send_frame_to_shard(driver1_shard);
+            let mock_driver2_action = send_frame_to_shard(driver2_shard);
+
+            // Accepts two connections and sends a response to each of them.
+            let mock_node_action = async {
+                let mut conns_futs = (0..2)
+                    .map(|_| async {
+                        let (mut conn, driver_addr) = mock_node_listener.accept().await.unwrap();
+                        respond_with_shard_num(&mut conn, driver_addr.port() % shards_count).await;
+                        write_frame(params.for_response(), response_opcode, body_ref, &mut conn)
+                            .await
+                            .unwrap();
+                        conn
+                    })
+                    .collect::<Vec<_>>();
+                let conn2 = conns_futs.pop().unwrap().await;
+                let conn1 = conns_futs.pop().unwrap().await;
+                (conn1, conn2)
+            };
+
+            // we keep the connections open until proxy finishes to let it perform clean exit with no disconnects
+            let (_node_conns, _driver1_conn, _driver2_conn) =
+                join3(mock_node_action, mock_driver1_action, mock_driver2_action).await;
+
+            let assert_feedback_request = |feedback_request: RequestFrame| {
+                assert_eq!(feedback_request.params, params);
+                assert_eq!(
+                    FrameOpcode::Request(feedback_request.opcode),
+                    request_opcode
+                );
+                assert_eq!(feedback_request.body, body);
+            };
+
+            let assert_feedback_response = |feedback_response: ResponseFrame| {
+                assert_eq!(feedback_response.params, params.for_response());
+                assert_eq!(
+                    FrameOpcode::Response(feedback_response.opcode),
+                    response_opcode
+                );
+                assert_eq!(feedback_response.body, body);
+            };
+
+            let (feedback_request, shard1) = request_feedback_rx.recv().await.unwrap();
+            assert_feedback_request(feedback_request);
+            let (feedback_request, shard2) = request_feedback_rx.recv().await.unwrap();
+            assert_feedback_request(feedback_request);
+            let (feedback_response, shard3) = response_feedback_rx.recv().await.unwrap();
+            assert_feedback_response(feedback_response);
+            let (feedback_response, shard4) = response_feedback_rx.recv().await.unwrap();
+            assert_feedback_response(feedback_response);
+
+            // expected: {driver1_shard request, driver1_shard response, driver2_shard request, driver2_shard response}
+            let mut expected_shards = [driver1_shard, driver1_shard, driver2_shard, driver2_shard];
+            expected_shards.sort_unstable();
+
+            let mut got_shards = [
+                shard1.unwrap(),
+                shard2.unwrap(),
+                shard3.unwrap(),
+                shard4.unwrap(),
+            ];
+            got_shards.sort_unstable();
+
+            assert_eq!(expected_shards, got_shards);
+
+            running_proxy.finish().await.unwrap();
+        }
+    }
 }
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 2eb66f7e..6012906f 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -2048,7 +2048,7 @@ mod tests {
             .unwrap();
 
         // We must interrupt the driver's full connection opening, because our proxy does not interact further after Startup.
-        let startup_without_lwt_optimisation = select! {
+        let (startup_without_lwt_optimisation, _shard) = select! {
             _ = open_connection(UntranslatedEndpoint::ContactPoint(ContactPoint{address: proxy_addr, datacenter: None}), None, config.clone()) => unreachable!(),
             startup = startup_rx.recv() => startup.unwrap(),
         };
@@ -2056,7 +2056,7 @@ mod tests {
         proxy.running_nodes[0]
             .change_request_rules(Some(make_rules(options_with_lwt_optimisation_support)));
 
-        let startup_with_lwt_optimisation = select! {
+        let (startup_with_lwt_optimisation, _shard) = select! {
             _ = open_connection(UntranslatedEndpoint::ContactPoint(ContactPoint{address: proxy_addr, datacenter: None}), None, config.clone()) => unreachable!(),
             startup = startup_rx.recv() => startup.unwrap(),
         };


Test patch:
diff --git a/scylla/tests/integration/consistency.rs b/scylla/tests/integration/consistency.rs
index 9eaabf03..93a0ad8f 100644
--- a/scylla/tests/integration/consistency.rs
+++ b/scylla/tests/integration/consistency.rs
@@ -20,7 +20,7 @@ use scylla_cql::Consistency;
 use scylla_proxy::ShardAwareness;
 use scylla_proxy::{
     Condition, ProxyError, Reaction, RequestFrame, RequestOpcode, RequestReaction, RequestRule,
-    WorkerError,
+    TargetShard, WorkerError,
 };
 use std::sync::Arc;
 use std::time::Duration;
@@ -296,9 +296,9 @@ async fn consistency_is_correctly_set_in_cql_requests() {
             async fn check_consistencies(
                 consistency: Consistency,
                 serial_consistency: Option<SerialConsistency>,
-                mut request_rx: UnboundedReceiver<RequestFrame>,
-            ) -> UnboundedReceiver<RequestFrame> {
-                let request_frame = request_rx.recv().await.unwrap();
+                mut request_rx: UnboundedReceiver<(RequestFrame, Option<TargetShard>)>,
+            ) -> UnboundedReceiver<(RequestFrame, Option<TargetShard>)> {
+                let (request_frame, _shard) = request_rx.recv().await.unwrap();
                 let deserialized_request = request_frame.deserialize().unwrap();
                 assert_eq!(deserialized_request.get_consistency().unwrap(), consistency);
                 assert_eq!(
diff --git a/scylla/tests/integration/lwt_optimisation.rs b/scylla/tests/integration/lwt_optimisation.rs
index c96b47ba..19ed45be 100644
--- a/scylla/tests/integration/lwt_optimisation.rs
+++ b/scylla/tests/integration/lwt_optimisation.rs
@@ -9,7 +9,7 @@ use tokio::sync::mpsc;
 
 use scylla_proxy::{
     Condition, ProxyError, Reaction, RequestFrame, RequestOpcode, RequestReaction, RequestRule,
-    ResponseOpcode, ResponseReaction, ResponseRule, ShardAwareness, WorkerError,
+    ResponseOpcode, ResponseReaction, ResponseRule, ShardAwareness, TargetShard, WorkerError,
 };
 
 #[tokio::test]
@@ -56,7 +56,7 @@ async fn if_lwt_optimisation_mark_offered_then_negotiatied_and_lwt_routed_optima
             .await
             .unwrap();
 
-        let supported_frame = supported_rx.recv().await.unwrap();
+        let (supported_frame, _shard) = supported_rx.recv().await.unwrap();
         let supported_options = types::read_string_multimap(&mut &*supported_frame.body).unwrap();
         let supported_features = ProtocolFeatures::parse_from_supported(&supported_options);
 
@@ -74,20 +74,20 @@ async fn if_lwt_optimisation_mark_offered_then_negotiatied_and_lwt_routed_optima
             .unwrap();
 
 
-        fn clear_rxs(rxs: &mut [mpsc::UnboundedReceiver<RequestFrame>; 3]) {
+        fn clear_rxs(rxs: &mut [mpsc::UnboundedReceiver<(RequestFrame, Option<TargetShard>)>; 3]) {
             for rx in rxs.iter_mut() {
                 while rx.try_recv().is_ok() {}
             }
         }
 
-        async fn assert_all_replicas_queried(rxs: &mut [mpsc::UnboundedReceiver<RequestFrame>; 3]) {
+        async fn assert_all_replicas_queried(rxs: &mut [mpsc::UnboundedReceiver<(RequestFrame, Option<TargetShard>)>; 3]) {
             for rx in rxs.iter_mut() {
                 rx.recv().await.unwrap();
             }
             clear_rxs(rxs);
         }
 
-        fn assert_one_replica_queried(rxs: &mut [mpsc::UnboundedReceiver<RequestFrame>; 3]) {
+        fn assert_one_replica_queried(rxs: &mut [mpsc::UnboundedReceiver<(RequestFrame, Option<TargetShard>)>; 3]) {
             let mut found_queried = false;
             for rx in rxs.iter_mut() {
                 if rx.try_recv().is_ok() {
@@ -100,7 +100,7 @@ async fn if_lwt_optimisation_mark_offered_then_negotiatied_and_lwt_routed_optima
         }
 
         #[allow(unused)]
-        fn who_was_queried(rxs: &mut [mpsc::UnboundedReceiver<RequestFrame>; 3]) {
+        fn who_was_queried(rxs: &mut [mpsc::UnboundedReceiver<(RequestFrame, Option<TargetShard>)>; 3]) {
             for (i, rx) in rxs.iter_mut().enumerate() {
                 if rx.try_recv().is_ok() {
                     println!("{} was queried.", i);
diff --git a/scylla/tests/integration/main.rs b/scylla/tests/integration/main.rs
index c2cb8fc6..ab64db0c 100644
--- a/scylla/tests/integration/main.rs
+++ b/scylla/tests/integration/main.rs
@@ -4,4 +4,5 @@ mod hygiene;
 mod lwt_optimisation;
 mod new_session;
 mod retries;
+mod shards;
 pub(crate) mod utils;
diff --git a/scylla/tests/integration/shards.rs b/scylla/tests/integration/shards.rs
new file mode 100644
index 00000000..7bee6a3a
--- /dev/null
+++ b/scylla/tests/integration/shards.rs
@@ -0,0 +1,82 @@
+use std::sync::Arc;
+
+use crate::utils::test_with_3_node_cluster;
+use scylla::{test_utils::unique_keyspace_name, SessionBuilder};
+use tokio::sync::mpsc;
+
+use scylla_proxy::TargetShard;
+use scylla_proxy::{
+    Condition, Reaction, RequestOpcode, RequestReaction, RequestRule, ShardAwareness,
+};
+use scylla_proxy::{ProxyError, RequestFrame, WorkerError};
+
+#[tokio::test]
+#[ntest::timeout(30000)]
+#[cfg(not(scylla_cloud_tests))]
+async fn test_consistent_shard_awareness() {
+    use std::collections::HashSet;
+
+    let res = test_with_3_node_cluster(ShardAwareness::QueryNode, |proxy_uris, translation_map, mut running_proxy| async move {
+
+        let (feedback_txs, mut feedback_rxs): (Vec<_>, Vec<_>) = (0..3).map(|_| {
+            mpsc::unbounded_channel::<(RequestFrame, Option<TargetShard>)>()
+        }).unzip();
+        for (i, tx) in feedback_txs.iter().cloned().enumerate() {
+            running_proxy.running_nodes[i].change_request_rules(Some(vec![
+                RequestRule(Condition::RequestOpcode(RequestOpcode::Execute), RequestReaction::noop().with_feedback_when_performed(tx))
+            ]));
+        }
+
+        let session = SessionBuilder::new()
+            .known_node(proxy_uris[0].as_str())
+            .address_translator(Arc::new(translation_map))
+            .build()
+            .await
+            .unwrap();
+        let ks = unique_keyspace_name();
+
+        /* Prepare schema */
+        session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 3}}", ks), &[]).await.unwrap();
+        session
+            .query(
+                format!(
+                    "CREATE TABLE IF NOT EXISTS {}.t (a int, b int, c text, primary key (a, b))",
+                    ks
+                ),
+                &[],
+            )
+            .await
+            .unwrap();
+
+        let prepared = session.prepare(format!("INSERT INTO {}.t (a, b, c) VALUES (?, ?, 'abc')", ks)).await.unwrap();
+
+        let value_lists = [
+            (4, 2),
+            (2, 1),
+            (3, 7),
+        ];
+
+        fn assert_one_shard_queried(rx: &mut mpsc::UnboundedReceiver<(RequestFrame, Option<TargetShard>)>) {
+            let shards = std::iter::from_fn(|| rx.try_recv().ok().map(|(_frame, shard)| shard)).collect::<HashSet<_>>();
+            if !shards.is_empty() {
+                assert_eq!(shards.len(), 1);
+            }
+        }
+
+        for values in value_lists {
+            for _ in 0..10 {
+                session.execute(&prepared, values).await.unwrap();
+            }
+            for rx in feedback_rxs.iter_mut() {
+                assert_one_shard_queried(rx);
+            }
+        }
+
+        running_proxy
+    }).await;
+    match res {
+        Ok(()) => (),
+        Err(ProxyError::Worker(WorkerError::DriverDisconnected(_))) => (),
+        Err(err) => panic!("{}", err),
+    }
+}