Repository: scylladb/scylla-rust-driver
Pull request #1325: Fix shard out of range panic
Author: login: muzarski
Labels: name: bug
description: Something isn't working
Created: 2025-04-18T10:16:09Z | Merged: 2025-04-19T14:47:35Z

In this PR implements verification logic for the picked shard. The shard should not be out of range for the chosen node. I also added a regression test where we would previously panic in such scenario.

## 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.
- [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.~



Diff:
diff --git a/scylla/src/client/pager.rs b/scylla/src/client/pager.rs
index dfc855b6fb..fe9fdc2209 100644
--- a/scylla/src/client/pager.rs
+++ b/scylla/src/client/pager.rs
@@ -173,8 +173,7 @@ where
         self.log_request_start();
 
         'nodes_in_plan: for (node, shard) in query_plan {
-            let span =
-                trace_span!(parent: &self.parent_span, "Executing query", node = %node.address);
+            let span = trace_span!(parent: &self.parent_span, "Executing query", node = %node.address, shard = %shard);
             // For each node in the plan choose a connection to use
             // This connection will be reused for same node retries to preserve paging cache on the shard
             let connection: Arc<Connection> = match node
diff --git a/scylla/src/client/session.rs b/scylla/src/client/session.rs
index 44de7862fb..ff3197f86f 100644
--- a/scylla/src/client/session.rs
+++ b/scylla/src/client/session.rs
@@ -1966,7 +1966,7 @@ impl Session {
             .unwrap_or(execution_profile.consistency);
 
         'nodes_in_plan: for (node, shard) in request_plan {
-            let span = trace_span!("Executing request", node = %node.address);
+            let span = trace_span!("Executing request", node = %node.address, shard = %shard);
             'same_node_retries: loop {
                 trace!(parent: &span, "Execution started");
                 let connection = match node.connection_for_shard(shard).await {
diff --git a/scylla/src/network/connection_pool.rs b/scylla/src/network/connection_pool.rs
index 85004dd214..3b5bda0edd 100644
--- a/scylla/src/network/connection_pool.rs
+++ b/scylla/src/network/connection_pool.rs
@@ -299,7 +299,18 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!\
+                    Targeting a random/arbitrary shard."
+                );
+                None
+            })
+            .and_then(|shard_conns| Self::choose_random_connection_from_slice(shard_conns))
         {
             trace!(shard = shard, "Found connection for the target shard");
             return conn;
diff --git a/scylla/src/policies/load_balancing/mod.rs b/scylla/src/policies/load_balancing/mod.rs
index aa8aa070eb..192b02ccc9 100644
--- a/scylla/src/policies/load_balancing/mod.rs
+++ b/scylla/src/policies/load_balancing/mod.rs
@@ -56,6 +56,11 @@ pub type FallbackPlan<'a> =
 /// the driver will try to send the request. The first elements of the plan are the targets which are
 /// the best to contact (e.g. they might have the lowest latency).
 ///
+/// Note: The chosen shard serves as a hint for the driver. The driver may ignore the chosen
+/// shard and fallback to some random shard for the node in two cases:
+/// 1. Picked shard is out of bounds (i.e. greater or equal to `nr_shards` specified by the node).
+/// 2. The connection pool to the picked shard is empty.
+///
 /// Most requests are sent on the first try, so the request execution layer rarely needs to know more
 /// than one target from plan. To better optimize that case, `LoadBalancingPolicy` has two methods:
 /// `pick` and `fallback`. `pick` returns the first target to contact for a given request, `fallback`
diff --git a/scylla/tests/integration/default_policy.rs b/scylla/tests/integration/default_policy.rs
deleted file mode 100644
index 272c10b21f..0000000000
--- a/scylla/tests/integration/default_policy.rs
+++ /dev/null
@@ -1,26 +0,0 @@
-use scylla::client::execution_profile::ExecutionProfile;
-use scylla::policies::load_balancing::{DefaultPolicy, LatencyAwarenessBuilder};
-
-use crate::utils::{create_new_session_builder, setup_tracing};
-
-// This is a regression test for #696.
-#[tokio::test]
-#[ntest::timeout(1000)]
-async fn latency_aware_query_completes() {
-    setup_tracing();
-    let policy = DefaultPolicy::builder()
-        .latency_awareness(LatencyAwarenessBuilder::default())
-        .build();
-    let handle = ExecutionProfile::builder()
-        .load_balancing_policy(policy)
-        .build()
-        .into_handle();
-
-    let session = create_new_session_builder()
-        .default_execution_profile_handle(handle)
-        .build()
-        .await
-        .unwrap();
-
-    session.query_unpaged("whatever", ()).await.unwrap_err();
-}
diff --git a/scylla/tests/integration/load_balancing.rs b/scylla/tests/integration/load_balancing.rs
new file mode 100644
index 0000000000..042ca1445b
--- /dev/null
+++ b/scylla/tests/integration/load_balancing.rs
@@ -0,0 +1,83 @@
+use std::sync::Arc;
+
+use scylla::client::execution_profile::ExecutionProfile;
+use scylla::cluster::{ClusterState, NodeRef};
+use scylla::policies::load_balancing::{
+    DefaultPolicy, FallbackPlan, LatencyAwarenessBuilder, LoadBalancingPolicy, RoutingInfo,
+};
+use scylla::routing::Shard;
+
+use crate::utils::{create_new_session_builder, setup_tracing};
+
+// This is a regression test for #696.
+#[tokio::test]
+#[ntest::timeout(1000)]
+async fn latency_aware_query_completes() {
+    setup_tracing();
+    let policy = DefaultPolicy::builder()
+        .latency_awareness(LatencyAwarenessBuilder::default())
+        .build();
+    let handle = ExecutionProfile::builder()
+        .load_balancing_policy(policy)
+        .build()
+        .into_handle();
+
+    let session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
+    session.query_unpaged("whatever", ()).await.unwrap_err();
+}
+
+/// Regression test for panic that appeared if LBP picked a
+/// shard that is out of range for the node.
+#[tokio::test]
+async fn test_shard_out_of_range() {
+    setup_tracing();
+
+    #[derive(Debug)]
+    struct ShardOutOfRangeLBP;
+    impl LoadBalancingPolicy for ShardOutOfRangeLBP {
+        fn pick<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            cluster: &'a ClusterState,
+        ) -> Option<(NodeRef<'a>, Option<Shard>)> {
+            let node = &cluster.get_nodes_info()[0];
+            match node.sharder() {
+                Some(sharder) => Some((node, Some(sharder.nr_shards.get() as u32))),
+                // For Cassandra let's pick some crazy shard number - it should be ignored anyway.
+                None => Some((node, Some(u16::MAX as u32))),
+            }
+        }
+
+        fn fallback<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            _cluster: &'a ClusterState,
+        ) -> FallbackPlan<'a> {
+            Box::new(std::iter::empty())
+        }
+
+        fn name(&self) -> String {
+            "ShardOutOfRangeLBP".into()
+        }
+    }
+
+    let handle = ExecutionProfile::builder()
+        .load_balancing_policy(Arc::new(ShardOutOfRangeLBP))
+        .build()
+        .into_handle();
+    let session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
+    let _ = session
+        .query_unpaged("SELECT * FROM system.local WHERE key='local'", ())
+        .await
+        .unwrap();
+}
diff --git a/scylla/tests/integration/main.rs b/scylla/tests/integration/main.rs
index 35fa0974c4..4b4191c6ea 100644
--- a/scylla/tests/integration/main.rs
+++ b/scylla/tests/integration/main.rs
@@ -4,11 +4,11 @@ mod consistency;
 mod cql_collections;
 mod cql_types;
 mod cql_value;
-mod default_policy;
 mod execution_profiles;
 mod history;
 mod hygiene;
 mod large_batch_statements;
+mod load_balancing;
 mod lwt_optimisation;
 mod metadata_custom_timeouts;
 mod new_session;


Reviews:
user: login: Lorak-mmk
body: I'm okay with merging as-is, but I'd prefer the warning to be added.
As you mentioned, there is currently a log. The problem is that it uses `trace` level, so it won't get noticed.

While both situations (out of range shard, no connections to shard) allow us to continue with the request attempt on this target, the first one suggests a bug in LBP, so warning about it with higher level seems like a good idea.
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-18T15:52:15Z
user: login: Lorak-mmk
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-18T16:25:38Z
user: login: wprzytula
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-19T08:43:01Z
user: login: wprzytula
body: LGTM.
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-19T08:43:27Z
user: login: wprzytula
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-19T10:51:50Z
user: login: muzarski
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T11:08:51Z
user: login: wprzytula
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-19T11:12:57Z
user: login: Lorak-mmk
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-19T12:45:58Z

Review comments:
diff_hunk: @@ -299,7 +299,17 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!"
+                );
path: scylla/src/network/connection_pool.rs
user: login: wprzytula
body: :wrench:
created_at: 2025-04-19T08:40:49Z
author_association: COLLABORATOR
diff_hunk: @@ -299,7 +299,17 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!"
+                );
path: scylla/src/network/connection_pool.rs
user: login: wprzytula
body: ```suggestion
                warn!(
                    shard = shard,
                    "Requested shard is out of bounds.\
                    This is most probably a bug in custom LoadBalancingPolicy implementation!\
                    Targetting a random/arbitrary shard."
                );
```
created_at: 2025-04-19T08:41:02Z
author_association: COLLABORATOR
diff_hunk: @@ -299,7 +299,18 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!\
+                    Targetting a random/arbitrary shard."
path: scylla/src/network/connection_pool.rs
user: login: wprzytula
body: Sorry, I made a typo myself.
created_at: 2025-04-19T10:51:23Z
author_association: COLLABORATOR
diff_hunk: @@ -299,7 +299,18 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!\
+                    Targetting a random/arbitrary shard."
path: scylla/src/network/connection_pool.rs
user: login: wprzytula
body: ```suggestion
                    Targeting a random/arbitrary shard."
```
created_at: 2025-04-19T10:51:35Z
author_association: COLLABORATOR
diff_hunk: @@ -299,7 +299,18 @@ impl NodeConnectionPool {
         shard_conns: &[Vec<Arc<Connection>>],
     ) -> Arc<Connection> {
         // Try getting the desired connection
-        if let Some(conn) = Self::choose_random_connection_from_slice(&shard_conns[shard as usize])
+        if let Some(conn) = shard_conns
+            .get(shard as usize)
+            .or_else(|| {
+                warn!(
+                    shard = shard,
+                    "Requested shard is out of bounds.\
+                    This is most probably a bug in custom LoadBalancingPolicy implementation!\
+                    Targetting a random/arbitrary shard."
path: scylla/src/network/connection_pool.rs
user: login: muzarski
body: > Sorry, I made a typo myself.

And I didn't notice:)
created_at: 2025-04-19T11:08:51Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-18T10:18:17Z
body: `cargo semver-checks` found no API-breaking changes in this PR.
Checked commit: 9325e7644591ca87ac15d5d7854e2c79b9db71c9
author_association: NONE
user: login: Lorak-mmk
created_at: 2025-04-18T11:23:35Z
body: The way I understand it the root cause of the problem is the bracket indexing here: https://github.com/scylladb/scylla-rust-driver/blob/v1.1.0/scylla/src/network/connection_pool.rs#L302
This PR does not change this code, and instead introduces more code to verify the passed shard number. Could you explain why you chose this approach?

The approach that feels simpler to me is to change the linked bracket indexing to get, and print a warning if `None` is returned.
author_association: COLLABORATOR
user: login: muzarski
created_at: 2025-04-18T11:32:41Z
body: > The way I understand it the root cause of the problem is the bracket indexing here: https://github.com/scylladb/scylla-rust-driver/blob/v1.1.0/scylla/src/network/connection_pool.rs#L302 This PR does not change this code, and instead introduces more code to verify the passed shard number. Could you explain why you chose this approach?
> 
> The approach that feels simpler to me is to change the linked bracket indexing to get, and print a warning if `None` is returned.

Yeah, I wanted to mention that in the cover letter - I forgot about it. Sorry about that.

Firstly, we need to answer ourselves the question: if LBP picks shard that is out of range, should we just print a warning and try on some random shard? If so, the code becomes much simpler.

If not, I assume that we should move to the next target from the plan. So if we wanted to handle this check in `NodeConnectionPool`, methods that pick the connection should somehow "inform" the caller about the problem.
I see two possibilities:
1. Change return type from `Result<_, ConnectionPoolError>` to `Option<Result<_, ConnectionPoolError>>` (`None` indicating shard out of range).
2. Add the `ShardOutOfRange` variant to `ConnectionPoolError`.

I don't like `2.`, because picking shard out of range by LBP is not a connection pool error. We could do `1.`. If you don't like `1.`, then the only thing we can do is to check the `shard` before calling `Node::connection_for_shard` - this is currently done in this PR.

I listed possible alternative solutions, and I'm open to suggestions. The only thing I'm against is adding `ShardOutOfRange` variant to `ConnectionPoolError`.
author_association: CONTRIBUTOR
user: login: Lorak-mmk
created_at: 2025-04-18T11:45:18Z
body: The way I understand it the shard chosen by LBP is not a strict requirement, but a preference. In the linked connection pool code we see that if there are no connections to selected shard, then some other shard is chosen instead.
Shard out of range is another type of situation where we can't get a connection to preferred shard, so it should also not be a critical error warranting skipping entire plan element. Printing a warning and choosing another shard is imo perfectly reasonable thing to do here.
author_association: COLLABORATOR
user: login: muzarski
created_at: 2025-04-18T11:47:25Z
body: > The way I understand it the shard chosen by LBP is not a strict requirement, but a preference. In the linked connection pool code we see that if there are no connections to selected shard, then some other shard is chosen instead. Shard out of range is another type of situation where we can't get a connection to preferred shard, so it should also not be a critical error warranting skipping entire plan element. Printing a warning and choosing another shard is imo perfectly reasonable thing to do here.

Makes sense. I'll update this (this will significantly reduce the size of this PR).
author_association: CONTRIBUTOR
user: login: muzarski
created_at: 2025-04-18T14:12:53Z
body: v2: After the discussion with @Lorak-mmk , I adjusted the logic accordingly - now, when shard is out of range, we fallback to picking the connection to some other random shard.
author_association: CONTRIBUTOR
user: login: muzarski
created_at: 2025-04-18T16:11:03Z
body: > I'm okay with merging as-is, but I'd prefer the warning to be added. As you mentioned, there is currently a log. The problem is that it uses `trace` level, so it won't get noticed.
> 
> While both situations (out of range shard, no connections to shard) allow us to continue with the request attempt on this target, the first one suggests a bug in LBP, so warning about it with higher level seems like a good idea.

Done
author_association: CONTRIBUTOR
user: login: wprzytula
created_at: 2025-04-19T08:47:30Z
body: I took a look into `LoadBalancingPolicy` docstring and there's nothing about the shard being considered just a preference / hint (not a strict requirement). What I see there is:

```
That plan is a list of
targets (target is a node + an optional shard) to which
the driver will try to send the request. The first elements of the plan are the targets which are
the best to contact (e.g. they might have the lowest latency).
```

I believe we must clarify our understanding of the shard selection in the LBP docstring. Namely, let's inform that we ignore the shard number (and choose another randomly) if it's out of bounds for the given node.
author_association: COLLABORATOR
user: login: muzarski
created_at: 2025-04-19T10:27:49Z
body: v2.1: Addressed @wprzytula comments. Mentioned that picked shard serves as a hint in the LBP docstring.
author_association: CONTRIBUTOR