split QueryNode::next so the caller handles IndexRangeRequests (#23784)

allow the caller of QueryNode::next to execute index ranges in batches by returning an IndexRangeRequest that should be executed. The response is passed back in to QueryNode::feed. Then the caller calls QueryNode::next again and loops.

considered alternatives:
1. #23743 inverts control completely, pulling the QueryNode logic into a batched function.
2. execute `next()` in parallel, passing in a channel instead of `&mut Transaction`. unfortunately `next()` can call six different methods on Transaction which would all have to be handled by the channel.

GitOrigin-RevId: 7479a27cb2ad6f54c931d0f9f2c029823e3ec3e6

Author: ldanilek

crates/database/src/bootstrap_model/user_facing.rs

@@ -9,7 +9,6 @@ use common::{
         DeveloperDocument,
         ResolvedDocument,
     },
-    index::IndexKeyBytes,
     query::CursorPosition,
     runtime::Runtime,
     types::{
@@ -30,13 +29,15 @@ use value::{
     DeveloperDocumentId,
     Size,
     TableName,
+    TableNumber,
 };
 
 use crate::{
     metrics::{
         log_virtual_table_get,
         log_virtual_table_query,
     },
+    query::IndexRangeResponse,
     transaction::{
         IndexRangeRequest,
         MAX_PAGE_SIZE,
@@ -322,17 +323,12 @@ impl<'a, RT: Runtime> UserFacingModel<'a, RT> {
     async fn start_index_range(
         &mut self,
         request: IndexRangeRequest,
-    ) -> anyhow::Result<
-        Result<
-            (
-                Vec<(IndexKeyBytes, DeveloperDocument, WriteTimestamp)>,
-                CursorPosition,
-            ),
-            RangeRequest,
-        >,
-    > {
+    ) -> anyhow::Result<Result<IndexRangeResponse<TableNumber>, RangeRequest>> {
         if request.interval.is_empty() {
-            return Ok(Ok((vec![], CursorPosition::End)));
+            return Ok(Ok(IndexRangeResponse {
+                page: vec![],
+                cursor: CursorPosition::End,
+            }));
         }
 
         let max_rows = cmp::min(request.max_rows, MAX_PAGE_SIZE);
@@ -357,7 +353,10 @@ impl<'a, RT: Runtime> UserFacingModel<'a, RT> {
                 return Ok(Ok(virtual_result));
             },
             StableIndexName::Missing => {
-                return Ok(Ok((vec![], CursorPosition::End)));
+                return Ok(Ok(IndexRangeResponse {
+                    page: vec![],
+                    cursor: CursorPosition::End,
+                }));
             },
         };
         let index_name = tablet_index_name
@@ -378,13 +377,7 @@ impl<'a, RT: Runtime> UserFacingModel<'a, RT> {
     pub async fn index_range_batch(
         &mut self,
         requests: BTreeMap<BatchKey, IndexRangeRequest>,
-    ) -> BTreeMap<
-        BatchKey,
-        anyhow::Result<(
-            Vec<(IndexKeyBytes, DeveloperDocument, WriteTimestamp)>,
-            CursorPosition,
-        )>,
-    > {
+    ) -> BTreeMap<BatchKey, anyhow::Result<IndexRangeResponse<TableNumber>>> {
         let batch_size = requests.len();
         let mut results = BTreeMap::new();
         let mut fetch_requests = BTreeMap::new();
@@ -409,12 +402,15 @@ impl<'a, RT: Runtime> UserFacingModel<'a, RT> {
             .await;
 
         for (batch_key, fetch_result) in fetch_results {
-            let result = fetch_result.map(|(resolved_results, cursor)| {
-                let developer_results = resolved_results
+            let result = fetch_result.map(|IndexRangeResponse { page, cursor }| {
+                let developer_results = page
                     .into_iter()
                     .map(|(key, doc, ts)| (key, doc.to_developer(), ts))
                     .collect();
-                (developer_results, cursor)
+                IndexRangeResponse {
+                    page: developer_results,
+                    cursor,
+                }
             });
             results.insert(batch_key, result);
         }

crates/database/src/query/filter.rs

@@ -1,17 +1,17 @@
 use async_trait::async_trait;
 use common::{
-    document::GenericDocument,
     query::{
         CursorPosition,
         Expression,
     },
     runtime::Runtime,
-    types::WriteTimestamp,
 };
 
 use super::{
+    IndexRangeResponse,
     QueryNode,
     QueryStream,
+    QueryStreamNext,
     QueryType,
 };
 use crate::Transaction;
@@ -51,17 +51,24 @@ impl<T: QueryType> QueryStream<T> for Filter<T> {
         &mut self,
         tx: &mut Transaction<RT>,
         _prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Option<(GenericDocument<T::T>, WriteTimestamp)>> {
+    ) -> anyhow::Result<QueryStreamNext<T>> {
         loop {
             let (document, write_timestamp) =
                 match self.inner.next(tx, Some(FILTER_QUERY_PREFETCH)).await? {
-                    Some(v) => v,
-                    None => return Ok(None),
+                    QueryStreamNext::Ready(Some(v)) => v,
+                    QueryStreamNext::Ready(None) => return Ok(QueryStreamNext::Ready(None)),
+                    QueryStreamNext::WaitingOn(request) => {
+                        return Ok(QueryStreamNext::WaitingOn(request))
+                    },
                 };
             let value = document.value().0.clone();
             if self.expr.eval(&value)?.into_boolean()? {
-                return Ok(Some((document, write_timestamp)));
+                return Ok(QueryStreamNext::Ready(Some((document, write_timestamp))));
             }
         }
     }
+
+    fn feed(&mut self, index_range_response: IndexRangeResponse<T::T>) -> anyhow::Result<()> {
+        self.inner.feed(index_range_response)
+    }
 }

crates/database/src/query/index_range.rs

@@ -3,7 +3,6 @@ use std::{
     collections::VecDeque,
 };
 
-use anyhow::Context;
 use async_trait::async_trait;
 use common::{
     document::GenericDocument,
@@ -25,12 +24,13 @@ use common::{
     },
     version::Version,
 };
-use maplit::btreemap;
 
 use super::{
     query_scanned_too_many_documents_error,
     query_scanned_too_much_data,
+    IndexRangeResponse,
     QueryStream,
+    QueryStreamNext,
     QueryType,
     DEFAULT_QUERY_PREFETCH,
     MAX_QUERY_FETCH,
@@ -145,8 +145,7 @@ impl<T: QueryType> IndexRange<T> {
         &mut self,
         tx: &mut Transaction<RT>,
         prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Result<Option<(GenericDocument<T::T>, WriteTimestamp)>, IndexRangeRequest>>
-    {
+    ) -> anyhow::Result<QueryStreamNext<T>> {
         // If we have an end cursor, for correctness we need to process
         // the entire interval, so ignore `maximum_rows_read` and `maximum_bytes_read`.
         let enforce_limits = self.cursor_interval.end_inclusive.is_none();
@@ -177,10 +176,10 @@ impl<T: QueryType> IndexRange<T> {
                 self.printable_index_name.is_system_owned(),
             );
             self.returned_bytes += v.size();
-            return Ok(Ok(Some((v, timestamp))));
+            return Ok(QueryStreamNext::Ready(Some((v, timestamp))));
         }
         if let Some(CursorPosition::End) = self.cursor_interval.curr_exclusive {
-            return Ok(Ok(None));
+            return Ok(QueryStreamNext::Ready(None));
         }
         if self.unfetched_interval.is_empty() {
             // We're out of results. If we have an end cursor then we must
@@ -192,7 +191,7 @@ impl<T: QueryType> IndexRange<T> {
                     .clone()
                     .unwrap_or(CursorPosition::End),
             );
-            return Ok(Ok(None));
+            return Ok(QueryStreamNext::Ready(None));
         }
 
         let mut max_rows = prefetch_hint
@@ -205,7 +204,7 @@ impl<T: QueryType> IndexRange<T> {
             }
             max_rows = cmp::min(max_rows, maximum_rows_read - self.rows_read);
         }
-        Ok(Err(IndexRangeRequest {
+        Ok(QueryStreamNext::WaitingOn(IndexRangeRequest {
             stable_index_name: self.stable_index_name.clone(),
             interval: self.unfetched_interval.clone(),
             order: self.order,
@@ -227,25 +226,6 @@ impl<T: QueryType> IndexRange<T> {
         self.page.extend(page);
         Ok(())
     }
-
-    #[convex_macro::instrument_future]
-    async fn _next<RT: Runtime>(
-        &mut self,
-        tx: &mut Transaction<RT>,
-        prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Option<(GenericDocument<T::T>, WriteTimestamp)>> {
-        loop {
-            let request = match self.start_next(tx, prefetch_hint)? {
-                Ok(result) => return Ok(result),
-                Err(request) => request,
-            };
-            let (page, fetch_cursor) = T::index_range_batch(tx, btreemap! {0 => request})
-                .await
-                .remove(&0)
-                .context("batch_key missing")??;
-            self.process_fetch(page, fetch_cursor)?;
-        }
-    }
 }
 
 pub const fn soft_data_limit(hard_limit: usize) -> usize {
@@ -276,8 +256,12 @@ impl<T: QueryType> QueryStream<T> for IndexRange<T> {
         &mut self,
         tx: &mut Transaction<RT>,
         prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Option<(GenericDocument<T::T>, WriteTimestamp)>> {
-        self._next(tx, prefetch_hint).await
+    ) -> anyhow::Result<QueryStreamNext<T>> {
+        self.start_next(tx, prefetch_hint)
+    }
+
+    fn feed(&mut self, index_range_response: IndexRangeResponse<T::T>) -> anyhow::Result<()> {
+        self.process_fetch(index_range_response.page, index_range_response.cursor)
     }
 }
 

crates/database/src/query/limit.rs

@@ -3,15 +3,15 @@ use std::cmp;
 use async_trait::async_trait;
 use common::{
     self,
-    document::GenericDocument,
     query::CursorPosition,
     runtime::Runtime,
-    types::WriteTimestamp,
 };
 
 use super::{
+    IndexRangeResponse,
     QueryNode,
     QueryStream,
+    QueryStreamNext,
     QueryType,
     DEFAULT_QUERY_PREFETCH,
 };
@@ -52,18 +52,22 @@ impl<T: QueryType> QueryStream<T> for Limit<T> {
         &mut self,
         tx: &mut Transaction<RT>,
         prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Option<(GenericDocument<T::T>, WriteTimestamp)>> {
+    ) -> anyhow::Result<QueryStreamNext<T>> {
         if self.rows_emitted >= self.limit {
-            return Ok(None);
+            return Ok(QueryStreamNext::Ready(None));
         }
         let inner_hint = cmp::min(
             prefetch_hint.unwrap_or(DEFAULT_QUERY_PREFETCH),
             self.limit - self.rows_emitted,
         );
-        let next_value = self.inner.next(tx, Some(inner_hint)).await?;
-        if next_value.is_some() {
+        let result = self.inner.next(tx, Some(inner_hint)).await?;
+        if let QueryStreamNext::Ready(Some(_)) = result {
             self.rows_emitted += 1;
         }
-        Ok(next_value)
+        Ok(result)
+    }
+
+    fn feed(&mut self, index_range_response: IndexRangeResponse<T::T>) -> anyhow::Result<()> {
+        self.inner.feed(index_range_response)
     }
 }