quickwit-oss · fmassot · Aug 25, 2021 · fulmicoton · Aug 25, 2021
diff --git a/quickwit-search/src/client.rs b/quickwit-search/src/client.rs
@@ -19,6 +19,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
+use futures::StreamExt;
 use http::Uri;
 use quickwit_proto::LeafSearchStreamResult;
 use std::fmt;
@@ -126,21 +127,20 @@ impl SearchServiceClient {
     pub async fn leaf_search_stream(
         &mut self,
         request: quickwit_proto::LeafSearchStreamRequest,
-    ) -> crate::Result<UnboundedReceiverStream<Result<LeafSearchStreamResult, tonic::Status>>> {
+    ) -> crate::Result<UnboundedReceiverStream<crate::Result<LeafSearchStreamResult>>> {
+        // It could be possible to return a stream... but the stream should be Sync + Send and tonic returns
+        // a stream which is only Send.
+        let (result_sender, result_receiver) = tokio::sync::mpsc::unbounded_channel();
         match &mut self.client_impl {
             SearchServiceClientImpl::Grpc(grpc_client) => {
                 let mut grpc_client_clone = grpc_client.clone();
-                let (result_sender, result_receiver) = tokio::sync::mpsc::unbounded_channel();
+                let tonic_request = Request::new(request);
+                let mut results_stream = grpc_client_clone
+                    .leaf_search_stream(tonic_request)
+                    .await
+                    .map_err(|tonic_error| parse_grpc_error(&tonic_error))?
+                    .into_inner();
                 tokio::spawn(async move {
-                    let tonic_request = Request::new(request);
-                    let mut results_stream = grpc_client_clone
-                        .leaf_search_stream(tonic_request)
-                        .await
-                        .map_err(|tonic_error| parse_grpc_error(&tonic_error))?
-                        .into_inner();
-
-                    // TODO: returning stream instead of a channel may be better.
-                    // But this seems to be difficult. Try it at your own expense.
                     while let Some(result) = results_stream
                         .message()
                         .await
@@ -156,11 +156,25 @@ impl SearchServiceClient {
                     }
                     Result::<_, SearchError>::Ok(())
                 });
-
-                Ok(UnboundedReceiverStream::new(result_receiver))
             }
-            SearchServiceClientImpl::Local(service) => service.leaf_search_stream(request).await,
+            SearchServiceClientImpl::Local(service) => {
+                let mut results_stream = service.leaf_search_stream(request).await?;
+                tokio::spawn(async move {
+                    while let Some(result) = results_stream.next().await {
+                        // We want to stop doing unnecessary work on the leaves as soon as
+                        // there is an issue sending the result.
+                        // Terminating the task will drop the `result_stream` consequently
+                        // canceling the gRPC request.
+                        result_sender.send(result).map_err(|_| {
+                            SearchError::InternalError("Could not send leaf result".into())
+                        })?;
+                    }
+                    Result::<_, SearchError>::Ok(())
+                });
+            }
         }
+
+        Ok(UnboundedReceiverStream::new(result_receiver))
     }
 
     /// Perform fetch docs.

diff --git a/quickwit-search/src/error.rs b/quickwit-search/src/error.rs
@@ -29,7 +29,7 @@ use quickwit_storage::StorageResolverError;
 
 /// Possible SearchError
 #[allow(missing_docs)]
-#[derive(Error, Debug, Serialize, Deserialize)]
+#[derive(Error, Debug, Serialize, Deserialize, Clone)]
 pub enum SearchError {
     #[error("Index `{index_id}` does not exist.")]
     IndexDoesNotExist { index_id: String },

diff --git a/quickwit-search/src/search_stream/leaf.rs b/quickwit-search/src/search_stream/leaf.rs
@@ -22,6 +22,7 @@ use super::FastFieldCollectorBuilder;
 use crate::leaf::open_index;
 use crate::leaf::warmup;
 use crate::SearchError;
+use futures::{FutureExt, StreamExt};
 use quickwit_index_config::IndexConfig;
 use quickwit_proto::LeafSearchStreamResult;
 use quickwit_proto::OutputFormat;
@@ -31,7 +32,8 @@ use quickwit_storage::Storage;
 use std::sync::Arc;
 use tantivy::schema::Type;
 use tantivy::ReloadPolicy;
-use tokio_stream::wrappers::UnboundedReceiverStream;
+
+const CONCURRENT_SPLIT_SEARCH_STREAM: usize = 2;
 
 /// `leaf` step of search stream.
 // Note: we return a stream of a result with a tonic::Status error
@@ -42,38 +44,29 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 // signature defined by proto generated code.
 pub async fn leaf_search_stream(
     index_config: Arc<dyn IndexConfig>,
-    request: &SearchStreamRequest,
+    request: SearchStreamRequest,
     split_ids: Vec<String>,
     storage: Arc<dyn Storage>,
-) -> UnboundedReceiverStream<Result<LeafSearchStreamResult, tonic::Status>> {
-    let (result_sender, result_receiver) = tokio::sync::mpsc::unbounded_channel();
-    for split_id in split_ids {
-        let split_storage: Arc<dyn Storage> =
-            quickwit_storage::add_prefix_to_storage(storage.clone(), split_id.clone());
-        let index_config_clone = index_config.clone();
-        let result_sender_clone = result_sender.clone();
-        let request_clone = request.clone();
-        tokio::spawn(async move {
-            let leaf_split_result =
-                leaf_search_stream_single_split(index_config_clone, &request_clone, split_storage)
-                    .await
-                    .map_err(SearchError::convert_to_tonic_status);
-            result_sender_clone.send(leaf_split_result).map_err(|_| {
-                SearchError::InternalError(format!(
-                    "Unable to send leaf export result for split `{}`",
-                    split_id
-                ))
-            })?;
-            Result::<(), SearchError>::Ok(())
-        });
-    }
-    UnboundedReceiverStream::new(result_receiver)
+) -> impl futures::Stream<Item = crate::Result<LeafSearchStreamResult>> + Sync + Send + 'static {
+    futures::stream::iter(split_ids)
+        .map(move |split_id| {
+            let index_config_clone = index_config.clone();
+            let request_clone = request.clone();
+            let split_storage: Arc<dyn Storage> =
+                quickwit_storage::add_prefix_to_storage(storage.clone(), split_id);
+            (index_config_clone, request_clone, split_storage)
+        })
+        .map(|(index_config_clone, request_clone, split_storage)| {
+            leaf_search_stream_single_split(index_config_clone, request_clone, split_storage)
+                .shared()
+        })
+        .buffered(CONCURRENT_SPLIT_SEARCH_STREAM)
 }
 
 /// Apply a leaf search on a single split.
 async fn leaf_search_stream_single_split(
     index_config: Arc<dyn IndexConfig>,
-    stream_request: &SearchStreamRequest,
+    stream_request: SearchStreamRequest,
     storage: Arc<dyn Storage>,
 ) -> crate::Result<LeafSearchStreamResult> {
     let index = open_index(storage).await?;

diff --git a/quickwit-search/src/search_stream/root.rs b/quickwit-search/src/search_stream/root.rs
@@ -35,7 +35,6 @@ use quickwit_metastore::Metastore;
 use quickwit_proto::SearchRequest;
 
 use crate::client_pool::Job;
-use crate::error::parse_grpc_error;
 use crate::list_relevant_splits;
 use crate::root::job_for_splits;
 use crate::root::NodeSearchError;
@@ -88,8 +87,8 @@ pub async fn root_search_stream(
 
             let mut leaf_bytes: Vec<Bytes> = Vec::new();
             while let Some(leaf_result) = receiver.next().await {
-                let leaf_data = leaf_result.map_err(|status| NodeSearchError {
-                    search_error: parse_grpc_error(&status),
+                let leaf_data = leaf_result.map_err(|search_error| NodeSearchError {
+                    search_error,
                     split_ids: split_ids.clone(),
                 })?;
                 leaf_bytes.push(Bytes::from(leaf_data.data));
@@ -188,7 +187,7 @@ mod tests {
         }))?;
         mock_search_service.expect_leaf_search_stream().return_once(
             |_leaf_search_req: quickwit_proto::LeafSearchStreamRequest| {
-                Ok(UnboundedReceiverStream::new(result_receiver))
+                Ok(Box::pin(UnboundedReceiverStream::new(result_receiver)))
             },
         );
         // The test will hang on indefinitely if we don't drop the receiver.

diff --git a/quickwit-search/src/service.rs b/quickwit-search/src/service.rs
@@ -28,8 +28,8 @@ use quickwit_proto::{
 };
 use quickwit_proto::{LeafSearchStreamRequest, LeafSearchStreamResult, SearchStreamRequest};
 use quickwit_storage::StorageUriResolver;
+use std::pin::Pin;
 use std::sync::Arc;
-use tokio_stream::wrappers::UnboundedReceiverStream;
 use tracing::info;
 
 use crate::fetch_docs;
@@ -89,7 +89,16 @@ pub trait SearchService: 'static + Send + Sync {
     async fn leaf_search_stream(
         &self,
         _request: LeafSearchStreamRequest,
-    ) -> crate::Result<UnboundedReceiverStream<Result<LeafSearchStreamResult, tonic::Status>>>;
+    ) -> crate::Result<
+        Pin<
+            Box<
+                dyn futures::Stream<Item = crate::Result<LeafSearchStreamResult>>
+                    + Sync
+                    + Send
+                    + 'static,
+            >,
+        >,
+    >;
 }
 
 impl SearchServiceImpl {
@@ -172,7 +181,16 @@ impl SearchService for SearchServiceImpl {
     async fn leaf_search_stream(
         &self,
         leaf_stream_request: LeafSearchStreamRequest,
-    ) -> crate::Result<UnboundedReceiverStream<Result<LeafSearchStreamResult, tonic::Status>>> {
+    ) -> crate::Result<
+        Pin<
+            Box<
+                dyn futures::Stream<Item = crate::Result<LeafSearchStreamResult>>
+                    + Sync
+                    + Send
+                    + 'static,
+            >,
+        >,
+    > {
         let stream_request = leaf_stream_request
             .request
             .ok_or_else(|| SearchError::InternalError("No search request.".to_string()))?;
@@ -184,8 +202,9 @@ impl SearchService for SearchServiceImpl {
         let storage = self.storage_resolver.resolve(&index_metadata.index_uri)?;
         let split_ids = leaf_stream_request.split_ids;
         let index_config = index_metadata.index_config;
-        let leaf_receiver =
-            leaf_search_stream(index_config, &stream_request, split_ids, storage.clone()).await;
+        let leaf_receiver = Box::pin(
+            leaf_search_stream(index_config, stream_request, split_ids, storage.clone()).await,
+        );
         Ok(leaf_receiver)
     }
 }
diff --git a/quickwit-serve/src/grpc_adapter/search_adapter.rs b/quickwit-serve/src/grpc_adapter/search_adapter.rs
@@ -19,10 +19,10 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-use std::sync::Arc;
+use std::{pin::Pin, sync::Arc};
 
 use async_trait::async_trait;
-use tokio_stream::wrappers::UnboundedReceiverStream;
+use futures::TryStreamExt;
 
 use quickwit_proto::{
     search_service_server as grpc, LeafSearchStreamRequest, LeafSearchStreamResult,
@@ -79,8 +79,14 @@ impl grpc::SearchService for GrpcSearchAdapter {
         Ok(tonic::Response::new(fetch_docs_result))
     }
 
-    type LeafSearchStreamStream =
-        UnboundedReceiverStream<Result<LeafSearchStreamResult, tonic::Status>>;
+    type LeafSearchStreamStream = Pin<
+        Box<
+            dyn futures::Stream<Item = Result<LeafSearchStreamResult, tonic::Status>>
+                + Sync
+                + Send
+                + 'static,
+        >,
+    >;
     async fn leaf_search_stream(
         &self,
         request: tonic::Request<LeafSearchStreamRequest>,
@@ -90,7 +96,8 @@ impl grpc::SearchService for GrpcSearchAdapter {
             .0
             .leaf_search_stream(leaf_search_request)
             .await
-            .map_err(SearchError::convert_to_tonic_status)?;
-        Ok(tonic::Response::new(leaf_search_result))
+            .map_err(SearchError::convert_to_tonic_status)?
+            .map_err(SearchError::convert_to_tonic_status);
+        Ok(tonic::Response::new(Box::pin(leaf_search_result)))
     }
 }
diff --git a/quickwit-storage/src/error.rs b/quickwit-storage/src/error.rs
@@ -42,7 +42,7 @@ pub enum StorageErrorKind {
 
 /// Generic Storage Resolver Error.
 #[allow(missing_docs)]
-#[derive(Error, Debug, Serialize, Deserialize)]
+#[derive(Error, Debug, Serialize, Deserialize, Clone)]
 pub enum StorageResolverError {
     /// The input is not a valid URI.
     /// A protocol is required for the URI.