From 9407afbe06b04ffcdc1fe6f07175e904f687558a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Przytu=C5=82a?= <wojciech.przytula@scylladb.com>
Date: Tue, 9 Jul 2024 16:35:28 +0200
Subject: [PATCH] default_lb: LWT optimisation regression test

The test runs against ClusterData with node F being disabled by
HostFilter. In such arrangement, F should be never returned. As F is
the primary replica for executed statement and no DC is preferred,
the second replica (`A`) cannot be computed cheaply, so `pick` should
return None. Before the bug was fixed, `pick` would return an arbitrary
robinned node, e.g. `B` (not even a replica).
---
 .../src/transport/load_balancing/default.rs   | 76 ++++++++++++++++++-
 1 file changed, 75 insertions(+), 1 deletion(-)

diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
index 148887f566..46aa282992 100644
--- a/scylla/src/transport/load_balancing/default.rs
+++ b/scylla/src/transport/load_balancing/default.rs
@@ -993,6 +993,8 @@ impl<'a> TokenWithStrategy<'a> {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use scylla_cql::{frame::types::SerialConsistency, Consistency};
     use tracing::info;
 
@@ -1000,7 +1002,12 @@ mod tests {
         get_plan_and_collect_node_identifiers, mock_cluster_data_for_token_unaware_tests,
         ExpectedGroups, ExpectedGroupsBuilder,
     };
-    use crate::transport::locator::test::{TABLE_NTS_RF_2, TABLE_NTS_RF_3, TABLE_SS_RF_2};
+    use crate::host_filter::HostFilter;
+    use crate::transport::locator::tablets::TabletsInfo;
+    use crate::transport::locator::test::{
+        id_to_invalid_addr, mock_metadata_for_token_aware_tests, TABLE_NTS_RF_2, TABLE_NTS_RF_3,
+        TABLE_SS_RF_2,
+    };
     use crate::{
         load_balancing::{
             default::tests::framework::mock_cluster_data_for_token_aware_tests, Plan, RoutingInfo,
@@ -2306,6 +2313,73 @@ mod tests {
             )
             .await;
         }
+
+        let cluster_with_disabled_node_f = ClusterData::new(
+            mock_metadata_for_token_aware_tests(),
+            &Default::default(),
+            &HashMap::new(),
+            &None,
+            {
+                struct FHostFilter;
+                impl HostFilter for FHostFilter {
+                    fn accept(&self, peer: &crate::transport::topology::Peer) -> bool {
+                        peer.address != id_to_invalid_addr(F)
+                    }
+                }
+
+                Some(&FHostFilter)
+            },
+            TabletsInfo::new(),
+        )
+        .await;
+
+        let tests_with_disabled_node_f = [
+            // Keyspace NTS with RF=3 without preferred DC.
+            // The primary replica does not satisfy the predicate (being disabled by HostFilter),
+            // so pick() should return None and fallback should return A first.
+            //
+            // This is a regression test after a bug was fixed.
+            Test {
+                policy: DefaultPolicy {
+                    preferences: NodeLocationPreference::Any,
+                    is_token_aware: true,
+                    permit_dc_failover: true,
+                    pick_predicate: Box::new(|node, _shard| node.address != id_to_invalid_addr(F)),
+                    ..Default::default()
+                },
+                routing_info: RoutingInfo {
+                    token: Some(Token::new(160)),
+                    table: Some(TABLE_NTS_RF_3),
+                    consistency: Consistency::One,
+                    is_confirmed_lwt: true,
+                    ..Default::default()
+                },
+                // going through the ring, we get order: F , A , C , D , G , B , E
+                //                                      us  eu  eu  us  eu  eu  us
+                //                                      r2  r1  r1  r1  r2  r1  r1
+                expected_groups: ExpectedGroupsBuilder::new()
+                    // pick is empty, because the primary replica does not satisfy pick predicate,
+                    // and with LWT we cannot compute other replicas for NTS without allocations.
+                    .ordered([A, C, D, G, E]) // replicas
+                    .group([B]) // nodes
+                    .build(),
+            },
+        ];
+
+        for Test {
+            policy,
+            routing_info,
+            expected_groups,
+        } in tests_with_disabled_node_f
+        {
+            test_default_policy_with_given_cluster_and_routing_info(
+                &policy,
+                &cluster_with_disabled_node_f,
+                &routing_info,
+                &expected_groups,
+            )
+            .await;
+        }
     }
 }