Lift stream_revision_pairs into PersistenceReader (#40957)

GitOrigin-RevId: ea665cfe28ce651ca4d64a32a7d2945affc9f768

Author: goffrie

crates/common/src/persistence.rs

@@ -420,6 +420,36 @@ pub trait PersistenceReader: Send + Sync + 'static {
             .boxed()
     }
 
+    /// Loads revision pairs from the document log in the given timestamp range.
+    ///
+    /// If a tablet id is provided, the results are filtered to a single table.
+    fn load_revision_pairs(
+        &self,
+        tablet_id: Option<TabletId>,
+        range: TimestampRange,
+        order: Order,
+        page_size: u32,
+        retention_validator: Arc<dyn RetentionValidator>,
+    ) -> DocumentRevisionStream<'_> {
+        let stream = if let Some(tablet_id) = tablet_id {
+            self.load_documents_from_table(
+                tablet_id,
+                range,
+                order,
+                page_size,
+                retention_validator.clone(),
+            )
+        } else {
+            self.load_documents(range, order, page_size, retention_validator.clone())
+        };
+        crate::persistence_helpers::persistence_reader_stream_revision_pairs(
+            stream,
+            self,
+            retention_validator,
+        )
+        .boxed()
+    }
+
     /// Look up the previous revision of `(id, ts)`, returning a map where for
     /// each `(id, ts)` we have...
     ///
@@ -595,8 +625,8 @@ impl RepeatablePersistence {
         self.upper_bound
     }
 
-    /// Same as [`Persistence::load_documents`] but only including documents in
-    /// the snapshot range.
+    /// Same as [`PersistenceReader::load_documents`] but only including
+    /// documents in the snapshot range.
     pub fn load_documents(&self, range: TimestampRange, order: Order) -> DocumentStream<'_> {
         self.reader.load_documents(
             range.intersect(TimestampRange::snapshot(*self.upper_bound)),
@@ -606,8 +636,8 @@ impl RepeatablePersistence {
         )
     }
 
-    /// Same as [`Persistence::load_documents_from_table`] but only including
-    /// documents in the snapshot range.
+    /// Same as [`PersistenceReader::load_documents_from_table`] but only
+    /// including documents in the snapshot range.
     pub fn load_documents_from_table(
         &self,
         tablet_id: TabletId,
@@ -623,6 +653,23 @@ impl RepeatablePersistence {
         )
     }
 
+    /// Same as [`PersistenceReader::load_revision_pairs`] but only including
+    /// revisions in the snapshot range.
+    pub fn load_revision_pairs(
+        &self,
+        tablet_id: Option<TabletId>,
+        range: TimestampRange,
+        order: Order,
+    ) -> DocumentRevisionStream<'_> {
+        self.reader.load_revision_pairs(
+            tablet_id,
+            range.intersect(TimestampRange::snapshot(*self.upper_bound)),
+            order,
+            *DEFAULT_DOCUMENTS_PAGE_SIZE,
+            self.retention_validator.clone(),
+        )
+    }
+
     pub async fn previous_revisions(
         &self,
         ids: BTreeSet<(InternalDocumentId, Timestamp)>,

crates/common/src/persistence_helpers.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use anyhow::Context as _;
 use futures::{
     Stream,
@@ -12,19 +14,23 @@ use crate::{
     persistence::{
         DocumentLogEntry,
         DocumentPrevTsQuery,
+        PersistenceReader,
         RepeatablePersistence,
+        RetentionValidator,
     },
     try_chunks::TryChunksExt,
     types::Timestamp,
 };
 
 #[derive(Debug)]
+#[cfg_attr(any(test, feature = "testing"), derive(PartialEq))]
 pub struct DocumentRevision {
     pub ts: Timestamp,
     pub document: Option<ResolvedDocument>,
 }
 
 #[derive(Debug)]
+#[cfg_attr(any(test, feature = "testing"), derive(PartialEq))]
 pub struct RevisionPair {
     pub id: InternalDocumentId,
     pub rev: DocumentRevision,
@@ -47,6 +53,57 @@ impl RevisionPair {
 
 type RevisionStreamEntry = anyhow::Result<DocumentLogEntry>;
 
+/// Exposed as PersistenceReader::load_revision_pairs
+#[allow(clippy::needless_lifetimes)]
+#[try_stream(ok = RevisionPair, error = anyhow::Error)]
+pub(crate) async fn persistence_reader_stream_revision_pairs<'a, P: PersistenceReader + ?Sized>(
+    documents: impl Stream<Item = RevisionStreamEntry> + 'a,
+    reader: &'a P,
+    retention_validator: Arc<dyn RetentionValidator>,
+) {
+    let documents = documents.try_chunks2(*DOCUMENTS_IN_MEMORY);
+    futures::pin_mut!(documents);
+
+    while let Some(read_chunk) = documents.try_next().await? {
+        let queries = read_chunk
+            .iter()
+            .filter_map(|entry| {
+                entry.prev_ts.map(|prev_ts| DocumentPrevTsQuery {
+                    id: entry.id,
+                    ts: entry.ts,
+                    prev_ts,
+                })
+            })
+            .collect();
+        let mut prev_revs = reader
+            .previous_revisions_of_documents(queries, retention_validator.clone())
+            .await?;
+        for DocumentLogEntry {
+            ts,
+            prev_ts,
+            id,
+            value: document,
+            ..
+        } in read_chunk
+        {
+            let rev = DocumentRevision { ts, document };
+            let prev_rev = prev_ts
+                .map(|prev_ts| {
+                    let entry = prev_revs
+                        .remove(&DocumentPrevTsQuery { id, ts, prev_ts })
+                        .with_context(|| format!("prev_ts is missing for {id}@{ts}: {prev_ts}"))?;
+                    anyhow::Ok(DocumentRevision {
+                        ts: entry.ts,
+                        document: entry.value,
+                    })
+                })
+                .transpose()?;
+            yield RevisionPair { id, rev, prev_rev };
+        }
+    }
+}
+
+// TODO: remove this and make users go through PersistenceReader
 #[allow(clippy::needless_lifetimes)]
 #[try_stream(ok = RevisionPair, error = anyhow::Error)]
 pub async fn stream_revision_pairs<'a>(

crates/database/src/database_index_workers/index_writer.rs

@@ -31,10 +31,6 @@ use common::{
         RetentionValidator,
         TimestampRange,
     },
-    persistence_helpers::{
-        stream_revision_pairs,
-        RevisionPair,
-    },
     query::Order,
     runtime::{
         new_rate_limiter,
@@ -57,7 +53,6 @@ use futures::{
         self,
         FusedStream,
     },
-    Stream,
     StreamExt,
     TryStreamExt,
 };
@@ -403,11 +398,10 @@ impl<RT: Runtime> IndexWriter<RT> {
         );
         let (tx, rx) = mpsc::channel(32);
         let producer = async {
-            let revision_stream = self.stream_revision_pairs(
-                &repeatable_persistence,
+            let revision_stream = repeatable_persistence.load_revision_pairs(
+                index_selector.tablet_id(),
                 TimestampRange::new(start_ts..=*end_ts),
                 Order::Asc,
-                index_selector,
             );
             futures::pin_mut!(revision_stream);
             while let Some(revision_pair) = revision_stream.try_next().await? {
@@ -464,11 +458,10 @@ impl<RT: Runtime> IndexWriter<RT> {
             self.retention_validator.clone(),
         );
         let producer = async {
-            let revision_stream = self.stream_revision_pairs(
-                &repeatable_persistence,
+            let revision_stream = repeatable_persistence.load_revision_pairs(
+                index_selector.tablet_id(),
                 TimestampRange::new(end_ts..*start_ts),
                 Order::Desc,
-                index_selector,
             );
             futures::pin_mut!(revision_stream);
             while let Some(revision_pair) = revision_stream.try_next().await? {
@@ -528,21 +521,6 @@ impl<RT: Runtime> IndexWriter<RT> {
         start_ts.prior_ts(backfilled_ts)
     }
 
-    fn stream_revision_pairs<'a>(
-        &'a self,
-        reader: &'a RepeatablePersistence,
-        range: TimestampRange,
-        order: Order,
-        index_selector: &'a IndexSelector,
-    ) -> impl Stream<Item = anyhow::Result<RevisionPair>> + 'a {
-        let document_stream = if let Some(tablet_id) = index_selector.tablet_id() {
-            reader.load_documents_from_table(tablet_id, range, order)
-        } else {
-            reader.load_documents(range, order)
-        };
-        stream_revision_pairs(document_stream, reader)
-    }
-
     async fn write_index_entries(
         &self,
         updates: impl FusedStream<Item = (Timestamp, DatabaseIndexUpdate)>,

crates/database/src/persistence_helpers.rs

@@ -3,10 +3,7 @@ use common::{
         RepeatablePersistence,
         TimestampRange,
     },
-    persistence_helpers::{
-        stream_revision_pairs,
-        RevisionPair,
-    },
+    persistence_helpers::RevisionPair,
     query::Order,
     types::Timestamp,
 };
@@ -48,8 +45,7 @@ pub async fn stream_transactions<'a>(
     range: TimestampRange,
     order: Order,
 ) {
-    let document_stream = reader.load_documents(range, order);
-    let revision_stream = stream_revision_pairs(document_stream, reader);
+    let revision_stream = reader.load_revision_pairs(None /* tablet_id */, range, order);
     futures::pin_mut!(revision_stream);
 
     if let Some(first_pair) = revision_stream.try_next().await? {

crates/database/src/search_index_bootstrap.rs

@@ -28,10 +28,7 @@ use common::{
         RepeatablePersistence,
         TimestampRange,
     },
-    persistence_helpers::{
-        stream_revision_pairs,
-        RevisionPair,
-    },
+    persistence_helpers::RevisionPair,
     query::Order,
     runtime::{
         try_join_buffer_unordered,
@@ -441,16 +438,15 @@ pub fn stream_revision_pairs_for_indexes<'a>(
     persistence: &'a RepeatablePersistence,
     range: TimestampRange,
 ) -> impl Stream<Item = anyhow::Result<RevisionPair>> + 'a {
-    let document_stream = persistence
-        .load_documents(range, Order::Asc)
-        .try_filter(|entry| {
-            let is_in_indexed_table = tables_with_indexes.contains(&entry.id.table());
+    persistence
+        .load_revision_pairs(None /* tablet_id */, range, Order::Asc)
+        .try_filter(|revision| {
+            let is_in_indexed_table = tables_with_indexes.contains(&revision.id.table());
             if !is_in_indexed_table {
                 log_document_skipped();
             }
-            future::ready(tables_with_indexes.contains(&entry.id.table()))
-        });
-    stream_revision_pairs(document_stream, persistence)
+            future::ready(is_in_indexed_table)
+        })
 }
 
 impl<RT: Runtime> SearchIndexBootstrapWorker<RT> {