[Transaction] batch index_range (#23706)

refactor, not changing behavior yet.
we thread the batching up, from TransactionIndex::index_range up through UserFacingModel::index_range.

the Query layer is still unbatched, but this refactor unblocks that.

GitOrigin-RevId: 9704db01e647a72f519b28a6e92de2da7ed33523

Author: ldanilek

crates/common/src/types/index.rs

@@ -112,7 +112,7 @@ pub type TabletIndexName = GenericIndexName<TableId>;
 
 /// Like TabletIndexName in that it refers to a stable underlying index,
 /// but it works for virtual tables too.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum StableIndexName {
     Physical(TabletIndexName),
     Virtual(IndexName, TabletIndexName),

crates/database/src/bootstrap_model/user_facing.rs

@@ -10,11 +10,7 @@ use common::{
         ResolvedDocument,
     },
     index::IndexKeyBytes,
-    interval::Interval,
-    query::{
-        CursorPosition,
-        Order,
-    },
+    query::CursorPosition,
     runtime::Runtime,
     types::{
         StableIndexName,
@@ -41,7 +37,10 @@ use crate::{
         log_virtual_table_get,
         log_virtual_table_query,
     },
-    transaction::MAX_PAGE_SIZE,
+    transaction::{
+        IndexRangeRequest,
+        MAX_PAGE_SIZE,
+    },
     unauthorized_error,
     virtual_tables::VirtualTable,
     PatchValue,
@@ -320,72 +319,106 @@ impl<'a, RT: Runtime> UserFacingModel<'a, RT> {
         )
     }
 
-    /// NOTE: returns a page of results. Callers must call record_read_document
-    /// for all documents returned from the index stream.
-    #[convex_macro::instrument_future]
-    pub async fn index_range(
+    async fn start_index_range(
         &mut self,
-        stable_index_name: &StableIndexName,
-        interval: &Interval,
-        order: Order,
-        mut max_rows: usize,
-        version: Option<Version>,
-    ) -> anyhow::Result<(
-        Vec<(IndexKeyBytes, DeveloperDocument, WriteTimestamp)>,
-        CursorPosition,
-    )> {
-        if interval.is_empty() {
-            return Ok((vec![], CursorPosition::End));
+        request: IndexRangeRequest,
+    ) -> anyhow::Result<
+        Result<
+            (
+                Vec<(IndexKeyBytes, DeveloperDocument, WriteTimestamp)>,
+                CursorPosition,
+            ),
+            RangeRequest,
+        >,
+    > {
+        if request.interval.is_empty() {
+            return Ok(Ok((vec![], CursorPosition::End)));
         }
 
-        max_rows = cmp::min(max_rows, MAX_PAGE_SIZE);
+        let max_rows = cmp::min(request.max_rows, MAX_PAGE_SIZE);
 
-        let tablet_index_name = match stable_index_name {
+        let tablet_index_name = match request.stable_index_name {
             StableIndexName::Physical(tablet_index_name) => tablet_index_name,
             StableIndexName::Virtual(index_name, tablet_index_name) => {
                 log_virtual_table_query();
-                return VirtualTable::new(self.tx)
+                // TODO(lee) batch virtual table queryStreamNext
+                let virtual_result = VirtualTable::new(self.tx)
                     .index_range(
                         RangeRequest {
                             index_name: tablet_index_name.clone(),
                             printable_index_name: index_name.clone(),
-                            interval: interval.clone(),
-                            order,
+                            interval: request.interval.clone(),
+                            order: request.order,
                             max_size: max_rows,
                         },
-                        version,
+                        request.version,
                     )
-                    .await;
+                    .await?;
+                return Ok(Ok(virtual_result));
             },
             StableIndexName::Missing => {
-                return Ok((vec![], CursorPosition::End));
+                return Ok(Ok((vec![], CursorPosition::End)));
             },
         };
         let index_name = tablet_index_name
             .clone()
             .map_table(&self.tx.table_mapping().tablet_to_name())?;
+        Ok(Err(RangeRequest {
+            index_name: tablet_index_name.clone(),
+            printable_index_name: index_name,
+            interval: request.interval.clone(),
+            order: request.order,
+            max_size: max_rows,
+        }))
+    }
+
+    /// NOTE: returns a page of results. Callers must call record_read_document
+    /// for all documents returned from the index stream.
+    #[convex_macro::instrument_future]
+    pub async fn index_range_batch(
+        &mut self,
+        requests: BTreeMap<BatchKey, IndexRangeRequest>,
+    ) -> BTreeMap<
+        BatchKey,
+        anyhow::Result<(
+            Vec<(IndexKeyBytes, DeveloperDocument, WriteTimestamp)>,
+            CursorPosition,
+        )>,
+    > {
+        let batch_size = requests.len();
+        let mut results = BTreeMap::new();
+        let mut fetch_requests = BTreeMap::new();
+        for (batch_key, request) in requests {
+            match self.start_index_range(request).await {
+                Err(e) => {
+                    results.insert(batch_key, Err(e));
+                },
+                Ok(Ok(result)) => {
+                    results.insert(batch_key, Ok(result));
+                },
+                Ok(Err(request)) => {
+                    fetch_requests.insert(batch_key, request);
+                },
+            }
+        }
 
-        let (results, cursor) = self
+        let fetch_results = self
             .tx
             .index
-            .range(
-                &mut self.tx.reads,
-                RangeRequest {
-                    index_name: tablet_index_name.clone(),
-                    printable_index_name: index_name,
-                    interval: interval.clone(),
-                    order,
-                    max_size: max_rows,
-                },
-            )
-            .await?;
-        let developer_results = results
-            .into_iter()
-            .map(|(key, doc, ts)| {
-                let doc = doc.to_developer();
-                anyhow::Ok((key, doc, ts))
-            })
-            .try_collect()?;
-        Ok((developer_results, cursor))
+            .range_batch(&mut self.tx.reads, fetch_requests)
+            .await;
+
+        for (batch_key, fetch_result) in fetch_results {
+            let result = fetch_result.map(|(resolved_results, cursor)| {
+                let developer_results = resolved_results
+                    .into_iter()
+                    .map(|(key, doc, ts)| (key, doc.to_developer(), ts))
+                    .collect();
+                (developer_results, cursor)
+            });
+            results.insert(batch_key, result);
+        }
+        assert_eq!(results.len(), batch_size);
+        results
     }
 }

crates/database/src/query/index_range.rs

@@ -3,6 +3,7 @@ use std::{
     collections::VecDeque,
 };
 
+use anyhow::Context;
 use async_trait::async_trait;
 use common::{
     document::GenericDocument,
@@ -24,6 +25,7 @@ use common::{
     },
     version::Version,
 };
+use maplit::btreemap;
 
 use super::{
     query_scanned_too_many_documents_error,
@@ -35,6 +37,7 @@ use super::{
 };
 use crate::{
     metrics,
+    transaction::IndexRangeRequest,
     Transaction,
 };
 
@@ -138,80 +141,94 @@ impl<T: QueryType> IndexRange<T> {
         }
     }
 
-    #[convex_macro::instrument_future]
-    async fn _next<RT: Runtime>(
+    fn start_next<RT: Runtime>(
         &mut self,
         tx: &mut Transaction<RT>,
         prefetch_hint: Option<usize>,
-    ) -> anyhow::Result<Option<(GenericDocument<T::T>, WriteTimestamp)>> {
-        loop {
-            // 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();
+    ) -> anyhow::Result<Result<Option<(GenericDocument<T::T>, WriteTimestamp)>, IndexRangeRequest>>
+    {
+        // 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();
 
-            if enforce_limits
-                && let Some(maximum_bytes_read) = self.maximum_bytes_read
-                && self.returned_bytes >= maximum_bytes_read
-            {
-                // If we're over our data budget, throw an error.
-                // We do this after we've already exceeded the limit to ensure that
-                // paginated queries always scan at least one item so they can
-                // make progress.
-                return Err(query_scanned_too_much_data(self.returned_bytes).into());
-            }
+        if enforce_limits
+            && let Some(maximum_bytes_read) = self.maximum_bytes_read
+            && self.returned_bytes >= maximum_bytes_read
+        {
+            // If we're over our data budget, throw an error.
+            // We do this after we've already exceeded the limit to ensure that
+            // paginated queries always scan at least one item so they can
+            // make progress.
+            return Err(query_scanned_too_much_data(self.returned_bytes).into());
+        }
 
-            if let Some((index_position, v, timestamp)) = self.page.pop_front() {
-                let index_bytes = index_position.len();
-                if let Some(intermediate_cursors) = &mut self.intermediate_cursors {
-                    intermediate_cursors.push(CursorPosition::After(index_position.clone()));
-                }
-                self.cursor_interval.curr_exclusive = Some(CursorPosition::After(index_position));
-                self.returned_results += 1;
-                T::record_read_document(tx, &v, self.printable_index_name.table())?;
-                // Database bandwidth for index reads
-                tx.usage_tracker.track_database_egress_size(
-                    self.printable_index_name.table().to_string(),
-                    index_bytes as u64,
-                    self.printable_index_name.is_system_owned(),
-                );
-                self.returned_bytes += v.size();
-                return Ok(Some((v, timestamp)));
-            }
-            if let Some(CursorPosition::End) = self.cursor_interval.curr_exclusive {
-                return Ok(None);
-            }
-            if self.unfetched_interval.is_empty() {
-                // We're out of results. If we have an end cursor then we must
-                // have reached it. Otherwise we're at the end of the entire
-                // query.
-                self.cursor_interval.curr_exclusive = Some(
-                    self.cursor_interval
-                        .end_inclusive
-                        .clone()
-                        .unwrap_or(CursorPosition::End),
-                );
-                return Ok(None);
+        if let Some((index_position, v, timestamp)) = self.page.pop_front() {
+            let index_bytes = index_position.len();
+            if let Some(intermediate_cursors) = &mut self.intermediate_cursors {
+                intermediate_cursors.push(CursorPosition::After(index_position.clone()));
             }
+            self.cursor_interval.curr_exclusive = Some(CursorPosition::After(index_position));
+            self.returned_results += 1;
+            T::record_read_document(tx, &v, self.printable_index_name.table())?;
+            // Database bandwidth for index reads
+            tx.usage_tracker.track_database_egress_size(
+                self.printable_index_name.table().to_string(),
+                index_bytes as u64,
+                self.printable_index_name.is_system_owned(),
+            );
+            self.returned_bytes += v.size();
+            return Ok(Ok(Some((v, timestamp))));
+        }
+        if let Some(CursorPosition::End) = self.cursor_interval.curr_exclusive {
+            return Ok(Ok(None));
+        }
+        if self.unfetched_interval.is_empty() {
+            // We're out of results. If we have an end cursor then we must
+            // have reached it. Otherwise we're at the end of the entire
+            // query.
+            self.cursor_interval.curr_exclusive = Some(
+                self.cursor_interval
+                    .end_inclusive
+                    .clone()
+                    .unwrap_or(CursorPosition::End),
+            );
+            return Ok(Ok(None));
+        }
 
-            let mut max_rows = prefetch_hint
-                .unwrap_or(DEFAULT_QUERY_PREFETCH)
-                .clamp(1, MAX_QUERY_FETCH);
+        let mut max_rows = prefetch_hint
+            .unwrap_or(DEFAULT_QUERY_PREFETCH)
+            .clamp(1, MAX_QUERY_FETCH);
 
-            if enforce_limits && let Some(maximum_rows_read) = self.maximum_rows_read {
-                if self.rows_read >= maximum_rows_read {
-                    return Err(query_scanned_too_many_documents_error(self.rows_read).into());
-                }
-                max_rows = cmp::min(max_rows, maximum_rows_read - self.rows_read);
+        if enforce_limits && let Some(maximum_rows_read) = self.maximum_rows_read {
+            if self.rows_read >= maximum_rows_read {
+                return Err(query_scanned_too_many_documents_error(self.rows_read).into());
             }
-            let (page, fetch_cursor) = T::index_range(
-                tx,
-                &self.stable_index_name,
-                &self.unfetched_interval,
-                self.order,
-                max_rows,
-                self.version.clone(),
-            )
-            .await?;
+            max_rows = cmp::min(max_rows, maximum_rows_read - self.rows_read);
+        }
+        Ok(Err(IndexRangeRequest {
+            stable_index_name: self.stable_index_name.clone(),
+            interval: self.unfetched_interval.clone(),
+            order: self.order,
+            max_rows,
+            version: self.version.clone(),
+        }))
+    }
+
+    #[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")??;
             let (_, new_unfetched_interval) =
                 self.unfetched_interval.split(fetch_cursor, self.order);
             anyhow::ensure!(self.unfetched_interval != new_unfetched_interval);