From d2c75911b2bafc0979f43756edb41d9716f967dd Mon Sep 17 00:00:00 2001 From: Joan Antoni RE Date: Thu, 9 Jul 2026 12:57:49 +0200 Subject: [PATCH 1/3] Move towards a generic distributed query algorithm --- nidx/src/searcher/grpc.rs | 99 +++++++++++++++++++++++++++------------ 1 file changed, 69 insertions(+), 30 deletions(-) diff --git a/nidx/src/searcher/grpc.rs b/nidx/src/searcher/grpc.rs index 8e65ff4fc4..104a5fa3cd 100644 --- a/nidx/src/searcher/grpc.rs +++ b/nidx/src/searcher/grpc.rs @@ -168,12 +168,15 @@ impl NidxSearcher for SearchServer { let response = if coordinator { // We are the requested searcher to perform a distributed query - self.distributed_query(request, shards).await? + let order_by = OrderBy::from(&request); + let limit = Limit(request.result_per_page as usize); + let partitions = self.distributed_query::(request, shards).await?; + SearchOp::merge(partitions, order_by, limit) } else { // This is a hopped node, i.e., another searcher has delegated this part of the query to // us. We must query our shards and return either a full or partial response or an error. // We won't hop to any other node - shard_search::search(shards, Arc::clone(&self.index_cache), request) + SearchOp::local_query(Arc::clone(&self.index_cache), request, shards) .await? .data }; @@ -271,7 +274,10 @@ impl SearchServer { /// requested to query several partitions and one or more can fail, it will then return a /// partial result and is our responsibility to retry on other nodes. /// - async fn distributed_query(&self, request: SearchRequest, shards: Vec) -> NidxResult { + async fn distributed_query(&self, request: Op::Request, shards: Vec) -> NidxResult> + where + Op: SearcherOp, + { let mut partitioner = QueryPartitioner::new(&self.shard_selector, &shards)?; // Use the query partitioner to group requests by the most preferred @@ -310,7 +316,7 @@ impl SearchServer { let mut tasks = JoinSet::new(); for (node, shard_ids) in groups { tasks.spawn( - self.node_query(node, shard_ids, request.clone()) + self.node_query::(node, shard_ids, request.clone()) .await? .instrument(Span::current()), ); @@ -353,45 +359,35 @@ impl SearchServer { } } - let merged = if responses.len() == 1 { - responses.pop().unwrap() - } else { - shard_merge::merge( - responses, - OrderBy::from(&request), - Limit(request.result_per_page as usize), - ) - }; - Ok(merged) + Ok(responses) } - async fn node_query( + async fn node_query( &self, node: SearcherNode, - shard_ids: Vec, - request: SearchRequest, - ) -> NidxResult>> + Send + 'static>>> { + shards: Vec, + request: Op::Request, + ) -> NidxResult>> + Send + 'static>>> + where + Op: SearcherOp, + { match node { SearcherNode::This => { let index_cache = Arc::clone(&self.index_cache); - Ok(Box::pin(async move { - shard_search::search(shard_ids, index_cache, request).await - })) + Ok(Box::pin( + async move { Op::local_query(index_cache, request, shards).await }, + )) } SearcherNode::Remote(ref hostname) => { match self.get_client(hostname).await { - Ok(client) => Ok(Box::pin(async move { - Self::remote_query(client, request, shard_ids) - .await - .map_err(NidxError::GrpcError) - })), + Ok(client) => Ok(Box::pin(async move { Op::remote_query(client, request, shards).await })), Err(e) => { // A client for this node is not available, although we've seen it // in the k8s cluster. We skip it and will retry those shards on // another node to complete the request warn!( ?node, - ?shard_ids, + ?shards, "{}", format!("Error getting a client for node '{hostname}': {e}") ); @@ -401,12 +397,47 @@ impl SearchServer { } } } +} + +trait SearcherOp: Clone + Send { + type Request: Clone + Send + 'static; + type Response: Send + 'static; + + fn local_query( + index_cache: Arc, + request: Self::Request, + shards: Vec, + ) -> impl std::future::Future>> + Send; + + fn remote_query( + client: SearcherClient, + request: Self::Request, + shards: Vec, + ) -> impl std::future::Future>> + Send; + + fn merge(partitions: Vec, order_by: OrderBy, limit: Limit) -> Self::Response; +} + +#[derive(Clone)] +struct SearchOp; + +impl SearcherOp for SearchOp { + type Request = SearchRequest; + type Response = SearchResponse; + + async fn local_query( + index_cache: Arc, + request: Self::Request, + shards: Vec, + ) -> NidxResult> { + shard_search::search(shards, index_cache, request).await + } async fn remote_query( mut client: SearcherClient, - request: SearchRequest, + request: Self::Request, shards: Vec, - ) -> tonic::Result> { + ) -> NidxResult> { // Send the query to a different node specifying only a subset of shards let mut search_request = request.clone(); search_request.shard_ids = shards.iter().map(|x| x.to_string()).collect(); @@ -416,7 +447,7 @@ impl SearchServer { let mut request = Request::new(search_request); request.metadata_mut().insert(HEADER_LOCAL_ONLY, "1".parse().unwrap()); - let response = client.search(request).await?.into_inner(); + let response = client.search(request).await.map_err(NidxError::GrpcError)?.into_inner(); let result = if response.shard_ids.len() == shards.len() { PartialResponse { data: response, shards } @@ -433,6 +464,14 @@ impl SearchServer { }; Ok(result) } + + fn merge(mut partitions: Vec, order_by: OrderBy, limit: Limit) -> Self::Response { + if partitions.len() == 1 { + partitions.pop().unwrap() + } else { + shard_merge::merge(partitions, order_by, limit) + } + } } struct QueryPartitioner { From ff2b3eb459381ff252330421bbcdf705cbcdd424 Mon Sep 17 00:00:00 2001 From: Joan Antoni RE Date: Thu, 9 Jul 2026 13:25:53 +0200 Subject: [PATCH 2/3] Add macro to simplify gRPC function code --- nidx/src/searcher/grpc.rs | 38 ++++++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/nidx/src/searcher/grpc.rs b/nidx/src/searcher/grpc.rs index 104a5fa3cd..e09e65505e 100644 --- a/nidx/src/searcher/grpc.rs +++ b/nidx/src/searcher/grpc.rs @@ -154,34 +154,35 @@ macro_rules! shard_request { }}; } -#[tonic::async_trait] -impl NidxSearcher for SearchServer { - async fn search(&self, request: Request) -> Result> { - let coordinator = !request.metadata().contains_key(HEADER_LOCAL_ONLY); - let request = request.into_inner(); +macro_rules! distributed_query { + ($self:ident, $op:ty, $request:ident) => {{ + let coordinator = !$request.metadata().contains_key(HEADER_LOCAL_ONLY); + let request = $request.into_inner(); - let mut shards = vec![]; - for shard_id in &request.shard_ids { - let shard_id = Uuid::parse_str(shard_id).map_err(NidxError::from)?; - shards.push(shard_id); - } + let shards = validate_shards(&request.shard_ids)?; let response = if coordinator { // We are the requested searcher to perform a distributed query let order_by = OrderBy::from(&request); let limit = Limit(request.result_per_page as usize); - let partitions = self.distributed_query::(request, shards).await?; - SearchOp::merge(partitions, order_by, limit) + let partitions = $self.distributed_query::<$op>(request, shards).await?; + <$op>::merge(partitions, order_by, limit) } else { // This is a hopped node, i.e., another searcher has delegated this part of the query to // us. We must query our shards and return either a full or partial response or an error. // We won't hop to any other node - SearchOp::local_query(Arc::clone(&self.index_cache), request, shards) + <$op>::local_query(Arc::clone(&$self.index_cache), request, shards) .await? .data }; - Ok(Response::new(response)) + }}; +} + +#[tonic::async_trait] +impl NidxSearcher for SearchServer { + async fn search(&self, request: Request) -> Result> { + distributed_query!(self, SearchOp, request) } async fn suggest(&self, request: Request) -> Result> { @@ -501,3 +502,12 @@ impl QueryPartitioner { if !nodes.is_empty() { Some(nodes.remove(0)) } else { None } } } + +fn validate_shards(shards: &[String]) -> NidxResult> { + let mut valid = Vec::with_capacity(shards.len()); + for shard_id in shards { + let shard_id = Uuid::parse_str(shard_id).map_err(NidxError::from)?; + valid.push(shard_id); + } + Ok(valid) +} From 3445466fe5a945c7a3e788799fd24b25da3f5f7e Mon Sep 17 00:00:00 2001 From: Joan Antoni RE Date: Thu, 9 Jul 2026 15:38:13 +0200 Subject: [PATCH 3/3] Rename macro. Move merge back inside distributed_query --- nidx/src/searcher/grpc.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/nidx/src/searcher/grpc.rs b/nidx/src/searcher/grpc.rs index e09e65505e..b8419c3de8 100644 --- a/nidx/src/searcher/grpc.rs +++ b/nidx/src/searcher/grpc.rs @@ -154,7 +154,8 @@ macro_rules! shard_request { }}; } -macro_rules! distributed_query { +/// Operation to multiple shards distributed across the searcher cluster +macro_rules! shards_request { ($self:ident, $op:ty, $request:ident) => {{ let coordinator = !$request.metadata().contains_key(HEADER_LOCAL_ONLY); let request = $request.into_inner(); @@ -163,10 +164,7 @@ macro_rules! distributed_query { let response = if coordinator { // We are the requested searcher to perform a distributed query - let order_by = OrderBy::from(&request); - let limit = Limit(request.result_per_page as usize); - let partitions = $self.distributed_query::<$op>(request, shards).await?; - <$op>::merge(partitions, order_by, limit) + $self.distributed_query::<$op>(request, shards).await? } else { // This is a hopped node, i.e., another searcher has delegated this part of the query to // us. We must query our shards and return either a full or partial response or an error. @@ -182,7 +180,7 @@ macro_rules! distributed_query { #[tonic::async_trait] impl NidxSearcher for SearchServer { async fn search(&self, request: Request) -> Result> { - distributed_query!(self, SearchOp, request) + shards_request!(self, SearchOp, request) } async fn suggest(&self, request: Request) -> Result> { @@ -275,7 +273,7 @@ impl SearchServer { /// requested to query several partitions and one or more can fail, it will then return a /// partial result and is our responsibility to retry on other nodes. /// - async fn distributed_query(&self, request: Op::Request, shards: Vec) -> NidxResult> + async fn distributed_query(&self, request: Op::Request, shards: Vec) -> NidxResult where Op: SearcherOp, { @@ -360,7 +358,8 @@ impl SearchServer { } } - Ok(responses) + let merged = Op::merge(&request, responses); + Ok(merged) } async fn node_query( @@ -416,7 +415,7 @@ trait SearcherOp: Clone + Send { shards: Vec, ) -> impl std::future::Future>> + Send; - fn merge(partitions: Vec, order_by: OrderBy, limit: Limit) -> Self::Response; + fn merge(request: &Self::Request, partitions: Vec) -> Self::Response; } #[derive(Clone)] @@ -466,10 +465,12 @@ impl SearcherOp for SearchOp { Ok(result) } - fn merge(mut partitions: Vec, order_by: OrderBy, limit: Limit) -> Self::Response { + fn merge(request: &Self::Request, mut partitions: Vec) -> Self::Response { if partitions.len() == 1 { partitions.pop().unwrap() } else { + let order_by = OrderBy::from(request); + let limit = Limit(request.result_per_page as usize); shard_merge::merge(partitions, order_by, limit) } }