feat(search): use stable node IDs for split affinity to keep disk cache warm when searcher pods restart

docs/configuration/node-config.md

@@ -21,7 +21,7 @@ A commented example is available here: [quickwit.yaml](https://github.com/quickw
 | --- | --- | --- | --- |
 | `version` | Config file version. `0.7` is the only available value with a retro compatibility on `0.5` and `0.4`. | | |
 | `cluster_id` | Unique identifier of the cluster the node will be joining. Clusters sharing the same network should use distinct cluster IDs.| `QW_CLUSTER_ID` | `quickwit-default-cluster` |
-| `node_id` | Unique identifier of the node. It must be distinct from the node IDs of its cluster peers. Defaults to the instance's short hostname if not set. | `QW_NODE_ID` | short hostname |
+| `node_id` | Unique identifier of the node. It must be distinct from the node IDs of its cluster peers. Searchers hash this ID for split affinity, so keep it stable across restarts (for example a StatefulSet pod name) if the same splits should keep landing on the same node. Defaults to the instance's short hostname if not set. | `QW_NODE_ID` | short hostname |
 | `enabled_services` | Enabled services (control_plane, indexer, janitor, metastore, metastore_read_replica, searcher) | `QW_ENABLED_SERVICES` | all services except metastore_read_replica |
 | `listen_address` | The IP address or hostname that Quickwit service binds to for starting REST and GRPC server and connecting this node to other nodes. By default, Quickwit binds itself to 127.0.0.1 (localhost). This default is not valid when trying to form a cluster. | `QW_LISTEN_ADDRESS` | `127.0.0.1` |
 | `advertise_address` | IP address advertised by the node, i.e. the IP address that peer nodes should use to connect to the node for RPCs. | `QW_ADVERTISE_ADDRESS` | `listen_address` |

quickwit/quickwit-datafusion/tests/distributed.rs

@@ -34,7 +34,8 @@ use quickwit_datafusion::test_utils::make_batch;
 use quickwit_datafusion::{
     DataFusionSessionBuilder, QuickwitObjectStoreRegistry, QuickwitWorkerResolver,
 };
-use quickwit_search::{SearcherPool, create_search_client_from_grpc_addr};
+use quickwit_proto::types::NodeId;
+use quickwit_search::{SearcherNode, SearcherPool, create_search_client_from_grpc_addr};
 
 mod common;
 mod metrics_splits;
@@ -104,12 +105,15 @@ async fn test_distributed_tasks_not_shuffles() {
     let worker_b = spawn_df_worker(make_builder()).await;
 
     // Populate the pool with the real worker gRPC addresses.
-    // Pool value is a SearchServiceClient — only the key (addr) matters for
+    // Pool value is a SearcherNode — only the key (addr) matters for
     // QuickwitWorkerResolver, which calls pool.keys() to get URLs.
     for addr in [worker_a.addr, worker_b.addr] {
         pool.insert(
             addr,
-            create_search_client_from_grpc_addr(addr, bytesize::ByteSize::mib(20)),
+            SearcherNode {
+                node_id: NodeId::from_str(&format!("worker-{}", addr.port())),
+                client: create_search_client_from_grpc_addr(addr, bytesize::ByteSize::mib(20)),
+            },
         );
     }
 

quickwit/quickwit-search/src/cluster_client.rs

@@ -305,15 +305,42 @@ mod tests {
     use std::net::SocketAddr;
     use std::sync::atomic::{AtomicU64, Ordering};
 
+    use quickwit_common::rendezvous_hasher::node_affinity;
     use quickwit_proto::search::{
         LeafRequestRef, PartialHit, SearchRequest, SortValue, SplitIdAndFooterOffsets,
         SplitSearchError,
     };
+    use quickwit_proto::types::NodeId;
     use quickwit_query::query_ast::qast_json_helper;
 
     use super::*;
     use crate::root::SearchJob;
-    use crate::{MockSearchService, SearchError, searcher_pool_for_test};
+    use crate::{
+        MockSearchService, SearchError, SearchServiceClient, SearcherNode, SearcherPool,
+        searcher_pool_for_test,
+    };
+
+    fn affinity_ordered_ports(key: &[u8], ports: &[u16]) -> Vec<u16> {
+        let mut nodes: Vec<(NodeId, u16)> = ports
+            .iter()
+            .copied()
+            .map(|port| (NodeId::from_str(&format!("node-{port}")), port))
+            .collect();
+        nodes
+            .sort_by_cached_key(|(node_id, _port)| std::cmp::Reverse(node_affinity(node_id, &key)));
+        nodes.into_iter().map(|(_node_id, port)| port).collect()
+    }
+
+    fn searcher_pool_from_port_mocks(
+        mocks: impl IntoIterator<Item = (u16, MockSearchService)>,
+    ) -> SearcherPool {
+        SearcherPool::from_iter(mocks.into_iter().map(|(port, mock_search_service)| {
+            let grpc_addr = SocketAddr::from(([127, 0, 0, 1], port));
+            let client =
+                SearchServiceClient::from_service(Arc::new(mock_search_service), grpc_addr);
+            (grpc_addr, SearcherNode::for_test(client))
+        }))
+    }
 
     fn mock_partial_hit(split_id: &str, sort_value: u64, doc_id: u32) -> PartialHit {
         PartialHit {
@@ -420,7 +447,7 @@ mod tests {
             ("127.0.0.1:1002", mock_search_service_2),
         ]);
         let first_client_addr: SocketAddr = "127.0.0.1:1001".parse().unwrap();
-        let first_client = searcher_pool.get(&first_client_addr).unwrap();
+        let first_client = searcher_pool.get(&first_client_addr).unwrap().client;
         let search_job_placer = SearchJobPlacer::new(searcher_pool);
         let cluster_client = ClusterClient::new(search_job_placer);
         let fetch_docs_response = cluster_client
@@ -441,7 +468,7 @@ mod tests {
         );
         let searcher_pool = searcher_pool_for_test([("127.0.0.1:1001", mock_search_service)]);
         let first_client_addr: SocketAddr = "127.0.0.1:1001".parse().unwrap();
-        let first_client = searcher_pool.get(&first_client_addr).unwrap();
+        let first_client = searcher_pool.get(&first_client_addr).unwrap().client;
         let search_job_placer = SearchJobPlacer::new(searcher_pool);
         let cluster_client = ClusterClient::new(search_job_placer);
         let search_error = cluster_client
@@ -602,35 +629,32 @@ mod tests {
 
     #[tokio::test]
     async fn test_put_kv_happy_path() {
-        // 3 servers 1, 2, 3
-        // Targeted key has affinity [2, 3, 1].
-        //
-        // Put on 2 and 3 is successful
-        // Get succeeds on 2.
-        let mock_search_service_1 = MockSearchService::new();
-        let mut mock_search_service_2 = MockSearchService::new();
-        mock_search_service_2.expect_put_kv().once().returning(
+        // Put on the two highest-affinity nodes; get succeeds on the first.
+        let ordered_ports = affinity_ordered_ports(b"my_key", &[1001, 1002, 1003]);
+        let mut mock_first = MockSearchService::new();
+        mock_first.expect_put_kv().once().returning(
             |put_req: quickwit_proto::search::PutKvRequest| {
                 assert_eq!(put_req.key, b"my_key");
                 assert_eq!(put_req.payload, b"my_payload");
             },
         );
-        mock_search_service_2.expect_get_kv().once().returning(
+        mock_first.expect_get_kv().once().returning(
             |get_req: quickwit_proto::search::GetKvRequest| {
                 assert_eq!(get_req.key, b"my_key");
                 Some(b"my_payload".to_vec())
             },
         );
-        let mut mock_search_service_3 = MockSearchService::new();
+        let mut mock_second = MockSearchService::new();
         // Due to the buffered call it is possible for the
-        // put request to 3 to be emitted too.
-        mock_search_service_3
+        // put request to the second node to be emitted too.
+        mock_second
             .expect_put_kv()
             .returning(|_put_req: quickwit_proto::search::PutKvRequest| {});
-        let searcher_pool = searcher_pool_for_test([
-            ("127.0.0.1:1001", mock_search_service_1),
-            ("127.0.0.1:1002", mock_search_service_2),
-            ("127.0.0.1:1003", mock_search_service_3),
+        let mock_third = MockSearchService::new();
+        let searcher_pool = searcher_pool_from_port_mocks([
+            (ordered_ports[0], mock_first),
+            (ordered_ports[1], mock_second),
+            (ordered_ports[2], mock_third),
         ]);
         let search_job_placer = SearchJobPlacer::new(searcher_pool);
         let cluster_client = ClusterClient::new(search_job_placer);
@@ -643,43 +667,40 @@ mod tests {
 
     #[tokio::test]
     async fn test_put_kv_failing_get() {
-        // 3 servers 1, 2, 3
-        // Targeted key has affinity [2, 3, 1].
-        //
-        // Put on 2 and 3 is successful
-        // Get fails on 2.
-        // Get succeeds on 3.
-        let mock_search_service_1 = MockSearchService::new();
-        let mut mock_search_service_2 = MockSearchService::new();
-        mock_search_service_2.expect_put_kv().once().returning(
+        // Put on the two highest-affinity nodes.
+        // Get fails on the first and succeeds on the second.
+        let ordered_ports = affinity_ordered_ports(b"my_key", &[1001, 1002, 1003]);
+        let mock_third = MockSearchService::new();
+        let mut mock_first = MockSearchService::new();
+        mock_first.expect_put_kv().once().returning(
             |put_req: quickwit_proto::search::PutKvRequest| {
                 assert_eq!(put_req.key, b"my_key");
                 assert_eq!(put_req.payload, b"my_payload");
             },
         );
-        mock_search_service_2.expect_get_kv().once().returning(
+        mock_first.expect_get_kv().once().returning(
             |get_req: quickwit_proto::search::GetKvRequest| {
                 assert_eq!(get_req.key, b"my_key");
                 None
             },
         );
-        let mut mock_search_service_3 = MockSearchService::new();
-        mock_search_service_3.expect_put_kv().once().returning(
+        let mut mock_second = MockSearchService::new();
+        mock_second.expect_put_kv().once().returning(
             |put_req: quickwit_proto::search::PutKvRequest| {
                 assert_eq!(put_req.key, b"my_key");
                 assert_eq!(put_req.payload, b"my_payload");
             },
         );
-        mock_search_service_3.expect_get_kv().once().returning(
+        mock_second.expect_get_kv().once().returning(
             |get_req: quickwit_proto::search::GetKvRequest| {
                 assert_eq!(get_req.key, b"my_key");
                 Some(b"my_payload".to_vec())
             },
         );
-        let searcher_pool = searcher_pool_for_test([
-            ("127.0.0.1:1001", mock_search_service_1),
-            ("127.0.0.1:1002", mock_search_service_2),
-            ("127.0.0.1:1003", mock_search_service_3),
+        let searcher_pool = searcher_pool_from_port_mocks([
+            (ordered_ports[0], mock_first),
+            (ordered_ports[1], mock_second),
+            (ordered_ports[2], mock_third),
         ]);
         let search_job_placer = SearchJobPlacer::new(searcher_pool);
         let cluster_client = ClusterClient::new(search_job_placer);

quickwit/quickwit-search/src/lib.rs

@@ -57,6 +57,7 @@ use tantivy::schema::NamedFieldDocument;
 /// Refer to this as `crate::Result<T>`.
 pub type Result<T> = std::result::Result<T, SearchError>;
 
+use std::hash::{Hash, Hasher};
 use std::net::{Ipv4Addr, SocketAddr};
 use std::sync::{Arc, LazyLock};
 
@@ -72,7 +73,7 @@ use quickwit_proto::search::{
     LeafResourceStats, PartialHit, SearchRequest, SearchResponse, SplitIdAndFooterOffsets,
     SplitResourceStats,
 };
-use quickwit_proto::types::IndexUid;
+use quickwit_proto::types::{IndexUid, NodeId};
 use quickwit_storage::StorageResolver;
 pub use service::SearcherContext;
 use tantivy::DocAddress;
@@ -94,8 +95,46 @@ pub use crate::search_response_rest::{
 };
 pub use crate::service::{MockSearchService, SearchService, SearchServiceImpl};
 
-/// A pool of searcher clients identified by their gRPC socket address.
-pub type SearcherPool = Pool<SocketAddr, SearchServiceClient>;
+/// A searcher pool entry: stable cluster identity plus the client used to dial it.
+///
+/// The pool stays keyed by gRPC address so RPC routing is unchanged. Rendezvous
+/// hashing uses [`Self::node_id`] (`node_id` / `QW_NODE_ID`) so a restart that
+/// keeps the same node ID keeps the same split affinity even if the listen
+/// address changes.
+#[derive(Clone, Debug)]
+pub struct SearcherNode {
+    /// Cluster node ID used for split affinity. Must stay unique and stable
+    /// across restarts.
+    pub node_id: NodeId,
+    /// Client used to send search RPCs to this node.
+    pub client: SearchServiceClient,
+}
+
+impl Hash for SearcherNode {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.node_id.hash(state);
+    }
+}
+
+#[cfg(any(test, feature = "testsuite"))]
+impl SearcherNode {
+    /// Wraps a client with a deterministic test node ID derived from its gRPC
+    /// port (`node-{port}`).
+    pub fn for_test(client: SearchServiceClient) -> Self {
+        let grpc_addr = client.grpc_addr();
+        Self {
+            node_id: NodeId::from_str(&format!("node-{}", grpc_addr.port())),
+            client,
+        }
+    }
+}
+
+/// A pool of searchers identified by their gRPC socket address.
+///
+/// Affinity hashing uses the value's [`SearcherNode::node_id`], not the address
+/// key, so placement survives listen-address changes when `QW_NODE_ID` is
+/// stable.
+pub type SearcherPool = Pool<SocketAddr, SearcherNode>;
 
 fn search_thread_pool() -> &'static ThreadPoolWithPriority {
     static SEARCH_THREAD_POOL: LazyLock<ThreadPoolWithPriority> =
@@ -298,7 +337,13 @@ pub async fn single_node_search(
     ));
     let search_service_client =
         SearchServiceClient::from_service(search_service.clone(), socket_addr);
-    searcher_pool.insert(socket_addr, search_service_client);
+    searcher_pool.insert(
+        socket_addr,
+        SearcherNode {
+            node_id: NodeId::from_str("single-node"),
+            client: search_service_client,
+        },
+    );
     root_search(
         &searcher_context,
         search_request,
@@ -342,7 +387,7 @@ pub fn searcher_pool_for_test(
                     .expect("The gRPC address should be valid socket address.");
                 let client =
                     SearchServiceClient::from_service(Arc::new(mock_search_service), grpc_addr);
-                (grpc_addr, client)
+                (grpc_addr, SearcherNode::for_test(client))
             }),
     )
 }

quickwit/quickwit-search/src/retry/mod.rs

@@ -116,8 +116,14 @@ mod tests {
             searcher_grpc_addr_2,
         );
         let searcher_pool = SearcherPool::from_iter([
-            (searcher_grpc_addr_1, searcher_client_1),
-            (searcher_grpc_addr_2, searcher_client_2),
+            (
+                searcher_grpc_addr_1,
+                crate::SearcherNode::for_test(searcher_client_1),
+            ),
+            (
+                searcher_grpc_addr_2,
+                crate::SearcherNode::for_test(searcher_client_2),
+            ),
         ]);
         let search_job_placer = SearchJobPlacer::new(searcher_pool);
         let _first_grpc_addr: SocketAddr = "127.0.0.1:1000".parse()?;