diff --git a/nidx/src/searcher/grpc.rs b/nidx/src/searcher/grpc.rs index 8e65ff4fc4..b8419c3de8 100644 --- a/nidx/src/searcher/grpc.rs +++ b/nidx/src/searcher/grpc.rs @@ -154,31 +154,33 @@ 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(); +/// 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(); - 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 - self.distributed_query(request, shards).await? + $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. // We won't hop to any other node - shard_search::search(shards, Arc::clone(&self.index_cache), request) + <$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> { + shards_request!(self, SearchOp, request) } async fn suggest(&self, request: Request) -> Result> { @@ -271,7 +273,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 +315,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 +358,36 @@ 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), - ) - }; + let merged = Op::merge(&request, responses); Ok(merged) } - 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(request: &Self::Request, partitions: Vec) -> 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,16 @@ impl SearchServer { }; Ok(result) } + + 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) + } + } } struct QueryPartitioner { @@ -462,3 +503,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) +}