From 057b3f7124333f886084e61cd3ceb865d3e11ad8 Mon Sep 17 00:00:00 2001 From: "cong.xie" Date: Wed, 19 Aug 2026 16:34:21 -0400 Subject: [PATCH 1/2] feat(search): use stable node IDs for split affinity Rendezvous hashing previously keyed off gRPC addresses, so a searcher restart that kept QW_NODE_ID but changed listen address remapped every split. Hash the cluster node ID instead, while keeping the searcher pool keyed by address for RPC. --- docs/configuration/node-config.md | 2 +- .../quickwit-datafusion/tests/distributed.rs | 10 +- .../quickwit-search/src/cluster_client.rs | 93 ++++++----- quickwit/quickwit-search/src/lib.rs | 55 ++++++- quickwit/quickwit-search/src/retry/mod.rs | 10 +- .../quickwit-search/src/search_job_placer.rs | 149 +++++++++++++++--- .../src/datafusion_api/setup.rs | 11 +- quickwit/quickwit-serve/src/lib.rs | 25 +-- 8 files changed, 267 insertions(+), 88 deletions(-) diff --git a/docs/configuration/node-config.md b/docs/configuration/node-config.md index cbe3f8b0022..8ce96ede62b 100644 --- a/docs/configuration/node-config.md +++ b/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` | diff --git a/quickwit/quickwit-datafusion/tests/distributed.rs b/quickwit/quickwit-datafusion/tests/distributed.rs index 7c8c498e35f..ca1cec9204f 100644 --- a/quickwit/quickwit-datafusion/tests/distributed.rs +++ b/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)), + }, ); } diff --git a/quickwit/quickwit-search/src/cluster_client.rs b/quickwit/quickwit-search/src/cluster_client.rs index 9cdadc6b413..0c295ef8ee0 100644 --- a/quickwit/quickwit-search/src/cluster_client.rs +++ b/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 { + 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, + ) -> 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); diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index 2d891dbfa65..90b545b9981 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -57,6 +57,7 @@ use tantivy::schema::NamedFieldDocument; /// Refer to this as `crate::Result`. pub type Result = std::result::Result; +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; +/// 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(&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; fn search_thread_pool() -> &'static ThreadPoolWithPriority { static SEARCH_THREAD_POOL: LazyLock = @@ -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)) }), ) } diff --git a/quickwit/quickwit-search/src/retry/mod.rs b/quickwit/quickwit-search/src/retry/mod.rs index 996665717cf..3420ac35acb 100644 --- a/quickwit/quickwit-search/src/retry/mod.rs +++ b/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()?; diff --git a/quickwit/quickwit-search/src/search_job_placer.rs b/quickwit/quickwit-search/src/search_job_placer.rs index dc1bd53c120..0d31b30e469 100644 --- a/quickwit/quickwit-search/src/search_job_placer.rs +++ b/quickwit/quickwit-search/src/search_job_placer.rs @@ -22,15 +22,16 @@ use std::time::Duration; use anyhow::bail; use async_trait::async_trait; use futures::future::join_all; +use quickwit_common::get_bool_from_env_cached; use quickwit_common::pubsub::EventSubscriber; use quickwit_common::rendezvous_hasher::{node_affinity, sort_by_rendez_vous_hash}; -use quickwit_common::{SocketAddrLegacyHash, get_bool_from_env_cached}; use quickwit_metrics::counter; use quickwit_proto::search::{ReportSplit, ReportSplitsRequest}; +use quickwit_proto::types::NodeId; use tracing::{info, warn}; use crate::metrics::JOB_ASSIGNED_TOTAL; -use crate::{SearchJob, SearchServiceClient, SearcherPool}; +use crate::{SearchJob, SearchServiceClient, SearcherNode, SearcherPool}; /// Job. /// The unit in which distributed search is performed. @@ -67,7 +68,7 @@ pub struct SearchJobPlacer { #[async_trait] impl EventSubscriber for SearchJobPlacer { async fn handle_event(&mut self, evt: ReportSplitsRequest) { - let mut nodes: HashMap = + let mut nodes: HashMap = self.searcher_pool.pairs().into_iter().collect(); if nodes.is_empty() { return; @@ -76,22 +77,23 @@ impl EventSubscriber for SearchJobPlacer { HashMap::with_capacity(nodes.len().min(evt.report_splits.len())); for report_split in evt.report_splits { let node_addr = nodes - .keys() - .max_by_key(|node_addr| { - node_affinity(SocketAddrLegacyHash(node_addr), &report_split.split_id) + .iter() + .max_by_key(|(_node_addr, node)| { + node_affinity(&node.node_id, &report_split.split_id) }) // This actually never happens thanks to the if-condition at the // top of this function. + .map(|(node_addr, _node)| *node_addr) .expect("`nodes` should not be empty"); splits_per_node - .entry(*node_addr) + .entry(node_addr) .or_default() .push(report_split); } for (node_addr, report_splits) in splits_per_node { - if let Some(search_client) = nodes.get_mut(&node_addr) { + if let Some(searcher_node) = nodes.get_mut(&node_addr) { let report_splits_req = ReportSplitsRequest { report_splits }; - let _ = search_client.report_splits(report_splits_req).await; + let _ = searcher_node.client.report_splits(report_splits_req).await; } } } @@ -115,13 +117,13 @@ impl SearchJobPlacer { } struct SocketAddrAndClient { - socket_addr: SocketAddr, + node_id: NodeId, client: SearchServiceClient, } impl Hash for SocketAddrAndClient { fn hash(&self, hasher: &mut H) { - SocketAddrLegacyHash(&self.socket_addr).hash(hasher); + self.node_id.hash(hasher); } } @@ -136,9 +138,9 @@ impl SearchJobPlacer { .searcher_pool .pairs() .into_iter() - .map(|(socket_addr, client)| SocketAddrAndClient { - socket_addr, - client, + .map(|(_socket_addr, searcher_node)| SocketAddrAndClient { + node_id: searcher_node.node_id, + client: searcher_node.client, }) .collect(); sort_by_rendez_vous_hash(&mut nodes[..], affinity_key); @@ -147,6 +149,22 @@ impl SearchJobPlacer { .map(|socket_addr_and_client| socket_addr_and_client.client) } + /// Returns searcher node IDs ordered by decreasing affinity with `affinity_key`. + #[cfg(test)] + async fn best_node_ids_per_affinity(&self, affinity_key: &[u8]) -> Vec { + let mut nodes: Vec = self + .searcher_pool + .pairs() + .into_iter() + .map(|(_grpc_addr, searcher_node)| searcher_node) + .collect(); + sort_by_rendez_vous_hash(&mut nodes[..], affinity_key); + nodes + .into_iter() + .map(|searcher_node| searcher_node.node_id) + .collect() + } + /// Assign the given job to the clients /// Returns a list of pair (SocketAddr, `Vec`) /// @@ -204,9 +222,10 @@ impl SearchJobPlacer { } let mut candidate_nodes: Vec = all_nodes .into_iter() - .map(|(grpc_addr, client)| CandidateNode { + .map(|(grpc_addr, searcher_node)| CandidateNode { + affinity_id: searcher_node.node_id, grpc_addr, - client, + client: searcher_node.client, load: None, }) .collect(); @@ -352,6 +371,7 @@ impl SearchJobPlacer { #[derive(Debug, Clone)] struct CandidateNode { + affinity_id: NodeId, pub grpc_addr: SocketAddr, pub client: SearchServiceClient, /// Current load of this node in job-cost units. `None` means the node @@ -361,13 +381,13 @@ struct CandidateNode { impl Hash for CandidateNode { fn hash(&self, state: &mut H) { - SocketAddrLegacyHash(&self.grpc_addr).hash(state); + self.affinity_id.hash(state); } } impl PartialEq for CandidateNode { fn eq(&self, other: &Self) -> bool { - self.grpc_addr == other.grpc_addr + self.affinity_id == other.affinity_id } } @@ -415,7 +435,9 @@ mod tests { use std::sync::Arc; use super::*; - use crate::{MockSearchService, SearchJob, SearchServiceClient, searcher_pool_for_test}; + use crate::{ + MockSearchService, SearchJob, SearchServiceClient, SearcherNode, searcher_pool_for_test, + }; fn searcher_pool_with_loads_for_test( iter: impl IntoIterator, @@ -427,7 +449,27 @@ mod tests { let client = SearchServiceClient::from_service(Arc::new(MockSearchService::new()), grpc_addr) .with_test_load(load); - (grpc_addr, client) + (grpc_addr, SearcherNode::for_test(client)) + })) + } + + fn searcher_pool_for_named_nodes( + iter: impl IntoIterator, + ) -> SearcherPool { + SearcherPool::from_iter(iter.into_iter().map(|(node_id, grpc_addr_str, load)| { + let grpc_addr: SocketAddr = grpc_addr_str + .parse() + .expect("the gRPC address should be a valid socket address"); + let client = + SearchServiceClient::from_service(Arc::new(MockSearchService::new()), grpc_addr) + .with_test_load(load); + ( + grpc_addr, + SearcherNode { + node_id: NodeId::from_str(node_id), + client, + }, + ) })) } @@ -552,17 +594,17 @@ mod tests { ( expected_searcher_addr_1, vec![ - SearchJob::for_test("split5", 5), SearchJob::for_test("split4", 4), SearchJob::for_test("split3", 3), + SearchJob::for_test("split2", 2), + SearchJob::for_test("split1", 1), ], ), ( expected_searcher_addr_2, vec![ SearchJob::for_test("split6", 6), - SearchJob::for_test("split2", 2), - SearchJob::for_test("split1", 1), + SearchJob::for_test("split5", 5), ], ), ]; @@ -629,9 +671,9 @@ mod tests { // With both nodes at equal load, each split should go to its highest-affinity // node as determined by rendezvous hashing. // - // Affinities for the (1001, 1002) pool (from test_search_job_placer): - // 1001 ← split3, split4, split5 - // 1002 ← split1, split2, split6 + // Affinities for the (node-1001, node-1002) pool: + // 1001 ← split3 + // 1002 ← split1 #[tokio::test] async fn test_equal_load_affinity_respected() { let searcher_pool = searcher_pool_for_test([ @@ -745,4 +787,59 @@ mod tests { split_ids.sort_unstable(); assert_eq!(split_ids, vec!["split1", "split3"]); } + + #[tokio::test] + async fn test_rendezvous_order_survives_address_change_for_same_node_ids() { + let first_pool = searcher_pool_for_named_nodes([ + ("searcher-0", "127.0.0.1:1001", 0), + ("searcher-1", "127.0.0.1:1002", 0), + ]); + let restarted_pool = searcher_pool_for_named_nodes([ + ("searcher-0", "127.0.0.1:2001", 0), + ("searcher-1", "127.0.0.1:2002", 0), + ]); + let before = SearchJobPlacer::new(first_pool) + .best_node_ids_per_affinity(b"split-a") + .await; + let after = SearchJobPlacer::new(restarted_pool) + .best_node_ids_per_affinity(b"split-a") + .await; + assert_eq!(before, after); + assert_eq!(before.len(), 2); + } + + // Load-aware fallback still skips the highest-affinity node when that node + // is overloaded, and picks the next node in the stable node-id order. + #[tokio::test] + async fn test_load_fallback_still_uses_next_stable_candidate() { + let overloaded_addr: SocketAddr = ([127, 0, 0, 1], 1001).into(); + let idle_addr: SocketAddr = ([127, 0, 0, 1], 1002).into(); + let searcher_pool = searcher_pool_for_named_nodes([ + ("searcher-0", "127.0.0.1:1001", 1_000_000), + ("searcher-1", "127.0.0.1:1002", 0), + ]); + let placer = SearchJobPlacer::new(searcher_pool); + + let mut split_id = "split-0".to_string(); + let mut ordered = Vec::new(); + for split_ord in 0..200 { + split_id = format!("split-{split_ord}"); + ordered = placer.best_node_ids_per_affinity(split_id.as_bytes()).await; + if ordered.first().map(|node_id| node_id.as_str()) == Some("searcher-0") { + break; + } + } + assert_eq!( + ordered.first().map(|node_id| node_id.as_str()), + Some("searcher-0"), + "could not find a split whose highest-affinity node is searcher-0" + ); + + let selected = placer + .assign_job(SearchJob::for_test(&split_id, 1), &HashSet::new()) + .await + .unwrap(); + assert_eq!(selected.grpc_addr(), idle_addr); + assert_ne!(selected.grpc_addr(), overloaded_addr); + } } diff --git a/quickwit/quickwit-serve/src/datafusion_api/setup.rs b/quickwit/quickwit-serve/src/datafusion_api/setup.rs index aad5e45a57b..7c1d61d5199 100644 --- a/quickwit/quickwit-serve/src/datafusion_api/setup.rs +++ b/quickwit/quickwit-serve/src/datafusion_api/setup.rs @@ -40,7 +40,7 @@ use quickwit_datafusion::{ QuickwitWorkerResolver, build_worker, }; use quickwit_proto::metastore::MetastoreServiceClient; -use quickwit_search::{SearchServiceClient, SearcherPool, create_search_client_from_grpc_addr}; +use quickwit_search::{SearcherNode, SearcherPool, create_search_client_from_grpc_addr}; use quickwit_storage::StorageResolver; use tokio::time::timeout; use tonic::transport::server::Router; @@ -112,7 +112,7 @@ fn setup_datafusion_worker_pool( async fn datafusion_worker_changes( cluster_change: ClusterChange, max_message_size: ByteSize, -) -> Vec> { +) -> Vec> { match cluster_change { ClusterChange::Add(node) if is_datafusion_worker_node(&node).await => { vec![insert_datafusion_worker(&node, max_message_size)] @@ -172,11 +172,14 @@ async fn exposes_datafusion_service(node: &ClusterNode) -> bool { fn insert_datafusion_worker( node: &ClusterNode, max_message_size: ByteSize, -) -> Change { +) -> Change { let grpc_addr = node.grpc_advertise_addr; Change::Insert( grpc_addr, - create_search_client_from_grpc_addr(grpc_addr, max_message_size), + SearcherNode { + node_id: node.node_id.clone(), + client: create_search_client_from_grpc_addr(grpc_addr, max_message_size), + }, ) } diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index 1bd34f3f4ac..151bf61a1f6 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -116,8 +116,8 @@ use quickwit_proto::metastore::{ use quickwit_proto::search::ReportSplitsRequest; use quickwit_proto::types::NodeId; use quickwit_search::{ - SearchJobPlacer, SearchService, SearchServiceClient, SearcherContext, SearcherPool, - create_search_client_from_channel, start_searcher_service, + SearchJobPlacer, SearchService, SearchServiceClient, SearcherContext, SearcherNode, + SearcherPool, create_search_client_from_channel, start_searcher_service, }; use quickwit_storage::{SearchSplitCache, StorageResolver}; pub use quickwit_telemetry_exporters::{EnvFilterReloadFn, do_nothing_env_filter_reload_fn}; @@ -1349,20 +1349,23 @@ async fn setup_searcher( chitchat_id.node_id, ); let grpc_addr = node.grpc_advertise_addr; - - if node.is_self_node() { - let search_client = - SearchServiceClient::from_service(search_service_clone, grpc_addr); - Some(Change::Insert(grpc_addr, search_client)) + let client = if node.is_self_node() { + SearchServiceClient::from_service(search_service_clone, grpc_addr) } else { let timeout_channel = Timeout::new(node.channel(), request_timeout); - let search_client = create_search_client_from_channel( + create_search_client_from_channel( grpc_addr, timeout_channel, max_message_size, - ); - Some(Change::Insert(grpc_addr, search_client)) - } + ) + }; + Some(Change::Insert( + grpc_addr, + SearcherNode { + node_id: node.node_id.clone(), + client, + }, + )) } ClusterChange::Remove(node) if node.is_searcher() => { let chitchat_id = node.chitchat_id(); From e7659f87481fd87d996f07d6e3444c2e64968bed Mon Sep 17 00:00:00 2001 From: "cong.xie" Date: Wed, 19 Aug 2026 17:30:36 -0400 Subject: [PATCH 2/2] refactor(search): rename SocketAddrAndClient to NodeIdAndClient The helper no longer stores a socket address; rendezvous hashing uses node_id only. --- quickwit/quickwit-search/src/search_job_placer.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/quickwit/quickwit-search/src/search_job_placer.rs b/quickwit/quickwit-search/src/search_job_placer.rs index 0d31b30e469..91ce37d27b0 100644 --- a/quickwit/quickwit-search/src/search_job_placer.rs +++ b/quickwit/quickwit-search/src/search_job_placer.rs @@ -116,12 +116,12 @@ impl SearchJobPlacer { } } -struct SocketAddrAndClient { +struct NodeIdAndClient { node_id: NodeId, client: SearchServiceClient, } -impl Hash for SocketAddrAndClient { +impl Hash for NodeIdAndClient { fn hash(&self, hasher: &mut H) { self.node_id.hash(hasher); } @@ -134,11 +134,11 @@ impl SearchJobPlacer { &self, affinity_key: &[u8], ) -> impl Iterator { - let mut nodes: Vec = self + let mut nodes: Vec = self .searcher_pool .pairs() .into_iter() - .map(|(_socket_addr, searcher_node)| SocketAddrAndClient { + .map(|(_socket_addr, searcher_node)| NodeIdAndClient { node_id: searcher_node.node_id, client: searcher_node.client, }) @@ -146,7 +146,7 @@ impl SearchJobPlacer { sort_by_rendez_vous_hash(&mut nodes[..], affinity_key); nodes .into_iter() - .map(|socket_addr_and_client| socket_addr_and_client.client) + .map(|node_id_and_client| node_id_and_client.client) } /// Returns searcher node IDs ordered by decreasing affinity with `affinity_key`.