Add stream splits gRPC.

quickwit-oss · Nov 9, 2023 · eb3af54 · eb3af54
1 parent 83c2cf5
commit eb3af54
Show file tree

Hide file tree

Showing 31 changed files with 481 additions and 461 deletions.
diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock
diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml
@@ -111,6 +111,7 @@ openssl = { version = "0.10.55", default-features = false }
 openssl-probe = "0.1.5"
 opentelemetry = { version = "0.19", features = ["rt-tokio"] }
 opentelemetry-otlp = "0.12.0"
+ouroboros = "0.18.0"
 pin-project = "1.1.0"
 pnet = { version = "0.33.0", features = ["std"] }
 postcard = { version = "1.0.4", features = ["use-std"], default-features = false}

diff --git a/quickwit/quickwit-cli/tests/cli.rs b/quickwit/quickwit-cli/tests/cli.rs
@@ -42,8 +42,7 @@ use quickwit_common::rand::append_random_suffix;
 use quickwit_common::uri::Uri;
 use quickwit_config::{SourceInputFormat, CLI_INGEST_SOURCE_ID};
 use quickwit_metastore::{
-    ListSplitsRequestExt, ListSplitsResponseExt, MetastoreResolver, MetastoreServiceExt,
-    SplitState, StageSplitsRequestExt,
+    ListSplitsRequestExt, MetastoreResolver, MetastoreServiceExt, SplitState, StageSplitsRequestExt,
 };
 use quickwit_proto::metastore::{
     DeleteSplitsRequest, EntityKind, IndexMetadataRequest, ListSplitsRequest,
@@ -254,8 +253,6 @@ async fn test_ingest_docs_cli() {
         .await
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
 
     assert_eq!(splits.len(), 1);
@@ -668,8 +665,6 @@ async fn test_garbage_collect_cli_no_grace() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     assert_eq!(splits.len(), 1);
 
@@ -718,8 +713,6 @@ async fn test_garbage_collect_cli_no_grace() {
             .list_splits(ListSplitsRequest::try_from_index_uid(index_uid).unwrap())
             .await
             .unwrap()
-            .deserialize_splits()
-            .unwrap()
             .len(),
         0
     );
@@ -779,8 +772,6 @@ async fn test_garbage_collect_index_cli() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     assert_eq!(splits.len(), 1);
 
@@ -798,8 +789,6 @@ async fn test_garbage_collect_index_cli() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     assert_eq!(splits.len(), 1);
 
@@ -839,8 +828,6 @@ async fn test_garbage_collect_index_cli() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     assert_eq!(splits[0].split_state, SplitState::Staged);
 
@@ -854,8 +841,6 @@ async fn test_garbage_collect_index_cli() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     assert_eq!(splits.len(), 1);
     assert_eq!(splits[0].split_state, SplitState::Staged);
@@ -872,8 +857,6 @@ async fn test_garbage_collect_index_cli() {
     let splits = metastore
         .list_splits(ListSplitsRequest::try_from_index_uid(index_uid).unwrap())
         .await
-        .unwrap()
-        .deserialize_splits()
         .unwrap();
     // Splits should be deleted from both metastore and file system.
     assert_eq!(splits.len(), 0);

diff --git a/quickwit/quickwit-common/src/stream_utils.rs b/quickwit/quickwit-common/src/stream_utils.rs
@@ -33,6 +33,30 @@ pub struct ServiceStream<T> {
     inner: BoxStream<T>,
 }
 
+impl<T> ServiceStream<T>
+where T: Send + 'static
+{
+    pub fn new(inner: BoxStream<T>) -> Self {
+        Self { inner }
+    }
+
+    pub fn empty() -> Self {
+        Self {
+            inner: Box::pin(stream::empty()),
+        }
+    }
+}
+
+impl<T> From<Vec<T>> for ServiceStream<T>
+where T: Send + 'static
+{
+    fn from(values: Vec<T>) -> Self {
+        Self {
+            inner: Box::pin(stream::iter(values)),
+        }
+    }
+}
+
 impl<T> fmt::Debug for ServiceStream<T>
 where T: 'static
 {

diff --git a/quickwit/quickwit-index-management/src/garbage_collection.rs b/quickwit/quickwit-index-management/src/garbage_collection.rs
@@ -25,7 +25,7 @@ use std::time::Duration;
 use futures::Future;
 use quickwit_common::{PrettySample, Progress};
 use quickwit_metastore::{
-    ListSplitsQuery, ListSplitsRequestExt, ListSplitsResponseExt, SplitInfo, SplitMetadata,
+    ListSplitsQuery, ListSplitsRequestExt, MetastoreServiceExt, SplitInfo, SplitMetadata,
     SplitState,
 };
 use quickwit_proto::metastore::{
@@ -106,7 +106,9 @@ pub async fn run_garbage_collect(
         metastore.list_splits(list_deletable_staged_request),
     )
     .await?
-    .deserialize_splits_metadata()?;
+    .into_iter()
+    .map(|split| split.split_metadata)
+    .collect();
 
     if dry_run {
         let marked_for_deletion_query = ListSplitsQuery::for_index(index_uid.clone())
@@ -118,7 +120,9 @@ pub async fn run_garbage_collect(
             metastore.list_splits(marked_for_deletion_request),
         )
         .await?
-        .deserialize_splits_metadata()?;
+        .into_iter()
+        .map(|split| split.split_metadata)
+        .collect();
         splits_marked_for_deletion.extend(deletable_staged_splits);
 
         let candidate_entries: Vec<SplitInfo> = splits_marked_for_deletion
@@ -192,9 +196,7 @@ async fn delete_splits_marked_for_deletion(
             }
         };
         let list_splits_result =
-            protect_future(progress_opt, metastore.list_splits(list_splits_request))
-                .await
-                .and_then(|list_splits_response| list_splits_response.deserialize_splits());
+            protect_future(progress_opt, metastore.list_splits(list_splits_request)).await;
 
         let splits_to_delete: Vec<SplitMetadata> = match list_splits_result {
             Ok(splits) => splits
@@ -345,14 +347,13 @@ mod tests {
     use std::time::Duration;
 
     use itertools::Itertools;
+    use quickwit_common::ServiceStream;
     use quickwit_config::IndexConfig;
     use quickwit_metastore::{
         metastore_for_test, CreateIndexRequestExt, ListSplitsQuery, SplitMetadata, SplitState,
         StageSplitsRequestExt,
     };
-    use quickwit_proto::metastore::{
-        CreateIndexRequest, EntityKind, ListSplitsResponse, StageSplitsRequest,
-    };
+    use quickwit_proto::metastore::{CreateIndexRequest, EntityKind, StageSplitsRequest};
     use quickwit_proto::types::IndexUid;
     use quickwit_storage::{
         storage_for_test, BulkDeleteError, DeleteFailure, MockStorage, PutPayload,
@@ -395,8 +396,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             1
         );
@@ -422,8 +421,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             1
         );
@@ -449,8 +446,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             1
         );
@@ -496,8 +491,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             1
         );
@@ -523,8 +516,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             1
         );
@@ -549,8 +540,6 @@ mod tests {
                 .list_splits(list_splits_request)
                 .await
                 .unwrap()
-                .deserialize_splits()
-                .unwrap()
                 .len(),
             0
         );
@@ -562,9 +551,9 @@ mod tests {
         let storage = storage_for_test();
         let mut metastore = MetastoreServiceClient::mock();
         metastore
-            .expect_list_splits()
+            .expect_stream_splits()
             .times(2)
-            .returning(|_| Ok(ListSplitsResponse::empty()));
+            .returning(|_| Ok(ServiceStream::empty()));
         run_garbage_collect(
             IndexUid::new_with_random_ulid("index-test-gc-deletes"),
             storage.clone(),
@@ -620,8 +609,6 @@ mod tests {
         let splits = metastore
             .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
             .await
-            .unwrap()
-            .deserialize_splits()
             .unwrap();
         assert_eq!(splits.len(), 1);
 
@@ -646,8 +633,6 @@ mod tests {
             .list_splits(ListSplitsRequest::try_from_index_uid(index_uid).unwrap())
             .await
             .unwrap()
-            .deserialize_splits()
-            .unwrap()
             .is_empty());
     }
 
@@ -738,8 +723,6 @@ mod tests {
         let splits = metastore
             .list_splits(ListSplitsRequest::try_from_index_uid(index_uid).unwrap())
             .await
-            .unwrap()
-            .deserialize_splits()
             .unwrap();
         assert_eq!(splits.len(), 1);
         assert_eq!(splits[0].split_id(), split_id_1);

diff --git a/quickwit/quickwit-index-management/src/index.rs b/quickwit/quickwit-index-management/src/index.rs
@@ -25,7 +25,7 @@ use quickwit_config::{validate_identifier, IndexConfig, SourceConfig};
 use quickwit_indexing::check_source_connectivity;
 use quickwit_metastore::{
     AddSourceRequestExt, CreateIndexRequestExt, IndexMetadata, IndexMetadataResponseExt,
-    ListSplitsQuery, ListSplitsRequestExt, ListSplitsResponseExt, SplitInfo, SplitMetadata,
+    ListSplitsQuery, ListSplitsRequestExt, MetastoreServiceExt, SplitInfo, SplitMetadata,
     SplitState,
 };
 use quickwit_proto::metastore::{
@@ -174,25 +174,26 @@ impl IndexService {
 
         if dry_run {
             let list_splits_request = ListSplitsRequest::try_from_index_uid(index_uid)?;
-            let splits_to_delete = self
+            let splits_to_delete: Vec<SplitInfo> = self
                 .metastore
                 .list_splits(list_splits_request)
                 .await?
-                .deserialize_splits()?
                 .into_iter()
                 .map(|split| split.split_metadata.as_split_info())
-                .collect::<Vec<_>>();
+                .collect();
             return Ok(splits_to_delete);
         }
         // Schedule staged and published splits for deletion.
         let query = ListSplitsQuery::for_index(index_uid.clone())
             .with_split_states([SplitState::Staged, SplitState::Published]);
         let list_splits_request = ListSplitsRequest::try_from_list_splits_query(query)?;
-        let split_ids = self
+        let split_ids: Vec<SplitId> = self
             .metastore
             .list_splits(list_splits_request)
             .await?
-            .deserialize_split_ids()?;
+            .into_iter()
+            .map(|split| split.split_metadata.split_id)
+            .collect();
         let mark_splits_for_deletion_request =
             MarkSplitsForDeletionRequest::new(index_uid.clone(), split_ids);
         self.metastore
@@ -203,11 +204,13 @@ impl IndexService {
         let query = ListSplitsQuery::for_index(index_uid.clone())
             .with_split_state(SplitState::MarkedForDeletion);
         let list_splits_request = ListSplitsRequest::try_from_list_splits_query(query)?;
-        let splits_to_delete = self
+        let splits_to_delete: Vec<SplitMetadata> = self
             .metastore
             .list_splits(list_splits_request)
             .await?
-            .deserialize_splits_metadata()?;
+            .into_iter()
+            .map(|split| split.split_metadata)
+            .collect();
 
         let deleted_splits = delete_splits_from_storage_and_metastore(
             index_uid.clone(),
@@ -291,7 +294,9 @@ impl IndexService {
             .metastore
             .list_splits(list_splits_request)
             .await?
-            .deserialize_splits_metadata()?;
+            .into_iter()
+            .map(|split| split.split_metadata)
+            .collect();
         let split_ids: Vec<SplitId> = splits_metadata
             .iter()
             .map(|split| split.split_id.to_string())
@@ -496,8 +501,6 @@ mod tests {
         let splits = metastore
             .list_splits(ListSplitsRequest::try_from_index_uid(index_uid.clone()).unwrap())
             .await
-            .unwrap()
-            .deserialize_splits()
             .unwrap();
         assert_eq!(splits.len(), 1);
 

diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs
@@ -589,13 +589,13 @@ mod tests {
     use std::sync::Arc;
 
     use quickwit_actors::{Command, Universe};
+    use quickwit_common::ServiceStream;
     use quickwit_config::{IndexingSettings, SourceInputFormat, SourceParams, VoidSourceParams};
     use quickwit_doc_mapper::{default_doc_mapper_for_test, DefaultDocMapper};
     use quickwit_metastore::checkpoint::IndexCheckpointDelta;
-    use quickwit_metastore::{IndexMetadata, ListSplitsResponseExt, PublishSplitsRequestExt};
+    use quickwit_metastore::{IndexMetadata, PublishSplitsRequestExt};
     use quickwit_proto::metastore::{
-        EmptyResponse, IndexMetadataResponse, LastDeleteOpstampResponse, ListSplitsResponse,
-        MetastoreError,
+        EmptyResponse, IndexMetadataResponse, LastDeleteOpstampResponse, MetastoreError,
     };
     use quickwit_proto::types::IndexUid;
     use quickwit_storage::RamStorage;
@@ -827,8 +827,8 @@ mod tests {
                 Ok(IndexMetadataResponse::try_from_index_metadata(index_metadata).unwrap())
             });
         mock_metastore
-            .expect_list_splits()
-            .returning(|_| Ok(ListSplitsResponse::empty()));
+            .expect_stream_splits()
+            .returning(|_| Ok(ServiceStream::empty()));
         let universe = Universe::with_accelerated_time();
         let node_id = "test-node";
         let doc_mapper = Arc::new(default_doc_mapper_for_test());