Avoid huge search RPC#41
Draft
rdettai-sk wants to merge 2 commits intosekoiafrom
Draft
Conversation
Implements StreamFetchDocs RPC that streams document contents in batches to avoid hitting gRPC message size limits. The batching size is controlled by SCROLL_BATCH_LEN (1000 documents per batch). Changes: - Add StreamFetchDocs RPC to search.proto - Implement fetch_docs_stream() function in fetch_docs.rs - Add stream_fetch_docs() to SearchService trait - Implement stream_fetch_docs() in SearchServiceImpl - Add StreamFetchDocs handler in gRPC adapter
40bc5d9 to
9a154a6
Compare
There was a problem hiding this comment.
Pull request overview
This PR introduces mechanisms to avoid oversized gRPC search responses by adding a server-streaming fetch-docs RPC, and clarifies the HTTP/gRPC separation by removing unused “root”/scroll/search-plan endpoints from the gRPC SearchService. It also adds cluster maintenance-mode APIs (control plane RPCs, REST endpoints, CLI/rest-client support) backed by a metastore key-value store.
Changes:
- Add
StreamFetchDocsgRPC endpoint and corresponding search-service streaming implementation. - Add control-plane “maintenance mode” feature with metastore KV persistence + REST/CLI/rest-client surfaces.
- Remove gRPC
RootSearch,RootListTerms,Scroll, andSearchPlanfrom the SearchService API.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| quickwit/quickwit-serve/src/search_api/grpc_adapter.rs | Implements new stream_fetch_docs gRPC adapter and removes obsolete gRPC endpoints. |
| quickwit/quickwit-serve/src/rest.rs | Wires cluster handler with a control-plane client. |
| quickwit/quickwit-serve/src/lib.rs | Minor module formatting change. |
| quickwit/quickwit-serve/src/cluster_api/rest_handler.rs | Adds REST maintenance-mode endpoints under /cluster/maintenance. |
| quickwit/quickwit-search/src/service.rs | Extends SearchService with stream_fetch_docs and implements it. |
| quickwit/quickwit-search/src/lib.rs | Exposes fetch_docs_stream internally alongside fetch_docs. |
| quickwit/quickwit-search/src/fetch_docs.rs | Adds fetch_docs_stream to stream FetchDocsResponse in batches. |
| quickwit/quickwit-search/src/client.rs | Removes gRPC root-search client method. |
| quickwit/quickwit-rest-client/src/rest_client.rs | Adds a maintenance REST client (status/enable/disable). |
| quickwit/quickwit-proto/src/control_plane/mod.rs | Adds RpcName implementations for maintenance RPCs. |
| quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs | Regenerates search gRPC stubs with StreamFetchDocs and removed endpoints. |
| quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs | Adds KV RPC/messages and tower plumbing in generated metastore API. |
| quickwit/quickwit-proto/src/codegen/quickwit/quickwit.control_plane.rs | Adds maintenance RPC/messages and tower plumbing in generated control-plane API. |
| quickwit/quickwit-proto/protos/quickwit/search.proto | Updates SearchService proto: remove root/scroll/search-plan; add StreamFetchDocs. |
| quickwit/quickwit-proto/protos/quickwit/metastore.proto | Adds metastore KV RPCs/messages. |
| quickwit/quickwit-proto/protos/quickwit/control_plane.proto | Adds maintenance-mode RPCs/messages. |
| quickwit/quickwit-metastore/src/tests/mod.rs | Registers new KV API test suite module. |
| quickwit/quickwit-metastore/src/tests/kv.rs | Adds metastore KV API tests (set/get/delete, edge cases). |
| quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs | Implements metastore KV operations for Postgres backend. |
| quickwit/quickwit-metastore/src/metastore/file_backed/state.rs | Extends file-backed metastore state with an in-memory KV store. |
| quickwit/quickwit-metastore/src/metastore/file_backed/mod.rs | Implements KV operations for file-backed metastore. |
| quickwit/quickwit-metastore/src/metastore/file_backed/manifest.rs | Persists KV store in the file-backed manifest format. |
| quickwit/quickwit-metastore/src/metastore/control_plane_metastore.rs | Delegates KV operations through the control-plane metastore wrapper. |
| quickwit/quickwit-control-plane/src/tests.rs | Adjusts control-plane test scaffolding to mock get_kv. |
| quickwit/quickwit-control-plane/src/metrics.rs | Adds maintenance_mode gauge metric. |
| quickwit/quickwit-control-plane/src/maintenance.rs | New maintenance-mode state + metastore KV persistence + tests. |
| quickwit/quickwit-control-plane/src/lib.rs | Exposes the new maintenance module. |
| quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs | Skips shard rebalance during maintenance mode. |
| quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs | Adds frozen-plan loading and gates plan rebuild/control in maintenance mode. |
| quickwit/quickwit-control-plane/src/indexing_plan.rs | Adds Deserialize to PhysicalIndexingPlan for persistence. |
| quickwit/quickwit-control-plane/src/control_plane.rs | Adds maintenance-mode RPC handlers and integrates persistence + gating logic. |
| quickwit/quickwit-control-plane/Cargo.toml | Adds dependencies for maintenance persistence (base64/prost/time). |
| quickwit/quickwit-cli/src/maintenance.rs | Adds CLI commands: `maintenance enable |
| quickwit/quickwit-cli/src/lib.rs | Exposes CLI maintenance module. |
| quickwit/quickwit-cli/src/cli.rs | Registers maintenance as a top-level CLI subcommand. |
| quickwit/Cargo.lock | Updates lockfile for new dependencies. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+166
to
+180
| futures::stream::once(async move { | ||
| let global_doc_addrs: Vec<GlobalDocAddress> = partial_hits | ||
| .iter() | ||
| .map(GlobalDocAddress::from_partial_hit) | ||
| .collect(); | ||
|
|
||
| let mut global_doc_addr_to_doc_json = fetch_docs_to_map( | ||
| searcher_context, | ||
| global_doc_addrs, | ||
| index_storage, | ||
| &splits, | ||
| doc_mapper, | ||
| snippet_request_opt.as_ref(), | ||
| ) | ||
| .await?; |
Comment on lines
+74
to
+78
| // Convert SearchError to tonic::Status | ||
| let grpc_stream = stream | ||
| .map_err(|err| tonic::Status::internal(format!("Stream fetch docs error: {}", err))); | ||
|
|
||
| Ok(tonic::Response::new(Box::pin(grpc_stream))) |
|
|
||
| let storage = match storage_resolver.resolve(&index_uri).await { | ||
| Ok(s) => s, | ||
| Err(e) => return Err(SearchError::Internal(e.to_string())), |
Comment on lines
+152
to
+157
| /// Creates a stream of `FetchDocsResponse` that yields documents in batches | ||
| /// to avoid hitting gRPC message size limits. | ||
| /// | ||
| /// This is the streaming version of `fetch_docs` that processes partial hits | ||
| /// and yields them in configurable batch sizes (default: SCROLL_BATCH_LEN). | ||
| pub fn fetch_docs_stream( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR contains 2 commits:
Describe how you tested this PR.