retention look up by exact prev_ts, checking ts against retention (#34449)

Redo of #34317 . Instead of querying in the documents log for exactly `(id, prev_ts)`, we change the interface to take in an `id, ts, prev_ts`. Then we can compare the `ts` against retention and use `prev_ts` for the direct equality lookup. And we can return a map that uses the entire `id, ts, prev_ts` as a key, to avoid ambiguity of which timestamp is used in the key.

GitOrigin-RevId: 1e8ac19ad499ad4751820b92cc2580b9bc433c78

Author: ldanilek

crates/common/src/persistence.rs

@@ -312,6 +312,13 @@ pub trait RetentionValidator: Sync + Send {
     fn fail_if_falling_behind(&self) -> anyhow::Result<()>;
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct DocumentPrevTsQuery {
+    pub id: InternalDocumentId,
+    pub ts: Timestamp,
+    pub prev_ts: Timestamp,
+}
+
 #[async_trait]
 pub trait PersistenceReader: Send + Sync + 'static {
     /// The persistence is required to load documents within the given timestamp
@@ -359,14 +366,14 @@ pub trait PersistenceReader: Send + Sync + 'static {
         retention_validator: Arc<dyn RetentionValidator>,
     ) -> anyhow::Result<BTreeMap<(InternalDocumentId, Timestamp), DocumentLogEntry>>;
 
-    /// Look up documents at exactly the specified timestamps, returning a map
-    /// where for each `(id, ts)` we have an entry only if a document exists
-    /// at exactly that timestamp.
-    async fn documents_multiget(
+    /// Look up documents at exactly the specified prev_ts timestamps, returning
+    /// a map where for each `DocumentPrevTsQuery` we have an entry only if
+    /// a document exists at `(id, prev_ts)`.
+    async fn previous_revisions_of_documents(
         &self,
-        ids: BTreeSet<(InternalDocumentId, Timestamp)>,
+        ids: BTreeSet<DocumentPrevTsQuery>,
         retention_validator: Arc<dyn RetentionValidator>,
-    ) -> anyhow::Result<BTreeMap<(InternalDocumentId, Timestamp), DocumentLogEntry>>;
+    ) -> anyhow::Result<BTreeMap<DocumentPrevTsQuery, DocumentLogEntry>>;
 
     /// Loads documentIds with respective timestamps that match the
     /// index query criteria.
@@ -586,16 +593,16 @@ impl RepeatablePersistence {
             .await
     }
 
-    pub async fn documents_multiget(
+    pub async fn previous_revisions_of_documents(
         &self,
-        ids: BTreeSet<(InternalDocumentId, Timestamp)>,
-    ) -> anyhow::Result<BTreeMap<(InternalDocumentId, Timestamp), DocumentLogEntry>> {
-        for (_, ts) in &ids {
-            // Reading documents ==ts, so ts needs to be repeatable.
+        ids: BTreeSet<DocumentPrevTsQuery>,
+    ) -> anyhow::Result<BTreeMap<DocumentPrevTsQuery, DocumentLogEntry>> {
+        for DocumentPrevTsQuery { ts, .. } in &ids {
+            // Reading previous revisions of documents at ts, so ts needs to be repeatable.
             anyhow::ensure!(*ts <= self.upper_bound);
         }
         self.reader
-            .documents_multiget(ids, self.retention_validator.clone())
+            .previous_revisions_of_documents(ids, self.retention_validator.clone())
             .await
     }
 

crates/common/src/testing/persistence_test_suite.rs

@@ -53,6 +53,7 @@ use crate::{
         fake_retention_validator::FakeRetentionValidator,
         ConflictStrategy,
         DocumentLogEntry,
+        DocumentPrevTsQuery,
         LatestDocument,
         NoopRetentionValidator,
         Persistence,
@@ -221,10 +222,13 @@ macro_rules! run_persistence_test_suite {
         }
 
         #[tokio::test]
-        async fn test_persistence_documents_multiget() -> anyhow::Result<()> {
+        async fn test_persistence_previous_revisions_of_documents() -> anyhow::Result<()> {
             let $db = $create_db;
             let p = $create_persistence;
-            persistence_test_suite::persistence_documents_multiget(::std::sync::Arc::new(p)).await
+            persistence_test_suite::persistence_previous_revisions_of_documents(
+                ::std::sync::Arc::new(p),
+            )
+            .await
         }
 
         #[tokio::test]
@@ -1756,7 +1760,9 @@ pub async fn persistence_delete_documents<P: Persistence>(p: Arc<P>) -> anyhow::
     Ok(())
 }
 
-pub async fn persistence_documents_multiget<P: Persistence>(p: Arc<P>) -> anyhow::Result<()> {
+pub async fn persistence_previous_revisions_of_documents<P: Persistence>(
+    p: Arc<P>,
+) -> anyhow::Result<()> {
     let mut id_generator = TestIdGenerator::new();
     let table: TableName = str::parse("table")?;
     let id1 = id_generator.user_generate(&table);
@@ -1806,38 +1812,87 @@ pub async fn persistence_documents_multiget<P: Persistence>(p: Arc<P>) -> anyhow
         id_generator.generate_internal(),
     );
 
+    // For the purposes of testing, set `ts` to be anything, because only `prev_ts`
+    // is used.
     let queries = btreeset![
         // Latest revision
-        (id1.into(), Timestamp::must(3)),
-        // Non-latest revision
-        (id1.into(), Timestamp::must(1)),
-        // Tombstone
-        (id2.into(), Timestamp::must(2)),
-        // Nonexistent revision
-        (id2.into(), Timestamp::must(3)),
+        DocumentPrevTsQuery {
+            id: id1.into(),
+            ts: Timestamp::must(4),
+            prev_ts: Timestamp::must(3),
+        },
+        // Previous revision of latest revision
+        DocumentPrevTsQuery {
+            id: id1.into(),
+            ts: Timestamp::must(3),
+            prev_ts: Timestamp::must(1)
+        },
+        // Tombstone (in this case ts doesn't actually exist but it's fine)
+        DocumentPrevTsQuery {
+            id: id2.into(),
+            ts: Timestamp::must(3),
+            prev_ts: Timestamp::must(2)
+        },
+        // Nonexistent revision at both ts and prev_ts
+        DocumentPrevTsQuery {
+            id: id2.into(),
+            ts: Timestamp::must(4),
+            prev_ts: Timestamp::must(3)
+        },
         // Unchanged document
-        (id3.into(), Timestamp::must(1)),
+        DocumentPrevTsQuery {
+            id: id3.into(),
+            ts: Timestamp::must(2),
+            prev_ts: Timestamp::must(1),
+        },
         // Nonexistent document
-        (nonexistent_id, Timestamp::must(1))
+        DocumentPrevTsQuery {
+            id: nonexistent_id,
+            ts: Timestamp::must(2),
+            prev_ts: Timestamp::must(1),
+        },
     ];
 
     // Test with NoopRetentionValidator
     // Note: Proper retention validation testing will be added in a separate PR
     let results = p
         .reader()
-        .documents_multiget(queries.clone(), Arc::new(NoopRetentionValidator))
+        .previous_revisions_of_documents(queries.clone(), Arc::new(NoopRetentionValidator))
         .await?;
 
     // Should get exact matches only
     assert_eq!(results.len(), 4); // id1@3, id1@1, id2@2, id3@1
-    assert!(results.contains_key(&(id1.into(), Timestamp::must(3))));
-    assert!(results.contains_key(&(id1.into(), Timestamp::must(1))));
-    assert!(results.contains_key(&(id2.into(), Timestamp::must(2))));
-    assert!(results.contains_key(&(id3.into(), Timestamp::must(1))));
+    assert!(results.contains_key(&DocumentPrevTsQuery {
+        id: id1.into(),
+        ts: Timestamp::must(3),
+        prev_ts: Timestamp::must(1)
+    }));
+    assert!(results.contains_key(&DocumentPrevTsQuery {
+        id: id2.into(),
+        ts: Timestamp::must(3),
+        prev_ts: Timestamp::must(2)
+    }));
+    assert!(results.contains_key(&DocumentPrevTsQuery {
+        id: id3.into(),
+        ts: Timestamp::must(2),
+        prev_ts: Timestamp::must(1)
+    }));
 
     // Verify document contents
-    let id1_at_3 = results.get(&(id1.into(), Timestamp::must(3))).unwrap();
-    let id1_at_1 = results.get(&(id1.into(), Timestamp::must(1))).unwrap();
+    let id1_at_3 = results
+        .get(&DocumentPrevTsQuery {
+            id: id1.into(),
+            ts: Timestamp::must(4),
+            prev_ts: Timestamp::must(3),
+        })
+        .unwrap();
+    let id1_at_1 = results
+        .get(&DocumentPrevTsQuery {
+            id: id1.into(),
+            ts: Timestamp::must(3),
+            prev_ts: Timestamp::must(1),
+        })
+        .unwrap();
 
     // Verify id1@3 has the correct document and prev_ts pointing to id1@1
     assert_eq!(id1_at_3.value, Some(doc(id1)));
@@ -1853,7 +1908,11 @@ pub async fn persistence_documents_multiget<P: Persistence>(p: Arc<P>) -> anyhow
     // Verify tombstone
     assert_eq!(
         results
-            .get(&(id2.into(), Timestamp::must(2)))
+            .get(&DocumentPrevTsQuery {
+                id: id2.into(),
+                ts: Timestamp::must(3),
+                prev_ts: Timestamp::must(2)
+            })
             .unwrap()
             .value,
         None
@@ -1862,21 +1921,25 @@ pub async fn persistence_documents_multiget<P: Persistence>(p: Arc<P>) -> anyhow
     let retention_validator = FakeRetentionValidator::new(Timestamp::must(4), Timestamp::must(0));
     // Min ts queried is 1, and min_document_ts is 0, so it's a valid query.
     p.reader()
-        .documents_multiget(queries.clone(), Arc::new(retention_validator))
+        .previous_revisions_of_documents(queries.clone(), Arc::new(retention_validator))
         .await?;
 
     let retention_validator = FakeRetentionValidator::new(Timestamp::must(4), Timestamp::must(4));
     // Min ts queried is 1, and min_document_ts is 4, so it's an invalid query.
     assert!(p
         .reader()
-        .documents_multiget(queries, Arc::new(retention_validator))
+        .previous_revisions_of_documents(queries, Arc::new(retention_validator))
         .await
         .is_err());
     // Errors even if there is no document at the timestamp.
     assert!(p
         .reader()
-        .documents_multiget(
-            btreeset![(nonexistent_id, Timestamp::must(1))],
+        .previous_revisions_of_documents(
+            btreeset![DocumentPrevTsQuery {
+                id: nonexistent_id,
+                ts: Timestamp::must(1),
+                prev_ts: Timestamp::must(1)
+            }],
             Arc::new(retention_validator)
         )
         .await

crates/common/src/testing/test_persistence.rs

@@ -39,6 +39,7 @@ use crate::{
     persistence::{
         ConflictStrategy,
         DocumentLogEntry,
+        DocumentPrevTsQuery,
         DocumentStream,
         IndexStream,
         LatestDocument,
@@ -282,25 +283,25 @@ impl PersistenceReader for TestPersistence {
         Ok(result)
     }
 
-    async fn documents_multiget(
+    async fn previous_revisions_of_documents(
         &self,
-        ids: BTreeSet<(InternalDocumentId, Timestamp)>,
+        ids: BTreeSet<DocumentPrevTsQuery>,
         retention_validator: Arc<dyn RetentionValidator>,
-    ) -> anyhow::Result<BTreeMap<(InternalDocumentId, Timestamp), DocumentLogEntry>> {
-        let min_ts = ids.iter().map(|(_, ts)| *ts).min();
+    ) -> anyhow::Result<BTreeMap<DocumentPrevTsQuery, DocumentLogEntry>> {
+        let min_ts = ids.iter().map(|DocumentPrevTsQuery { ts, .. }| *ts).min();
         let result = {
             let inner = self.inner.lock();
             let result = ids
                 .into_iter()
-                .filter_map(|(id, ts)| {
-                    inner.log.get(&(ts, id)).map(|(doc, prev_ts)| {
+                .filter_map(|DocumentPrevTsQuery { id, ts, prev_ts }| {
+                    inner.log.get(&(prev_ts, id)).map(|(doc, prev_prev_ts)| {
                         (
-                            (id, ts),
+                            DocumentPrevTsQuery { id, ts, prev_ts },
                             DocumentLogEntry {
                                 id,
-                                ts,
+                                ts: prev_ts,
                                 value: doc.clone(),
-                                prev_ts: *prev_ts,
+                                prev_ts: *prev_prev_ts,
                             },
                         )
                     })

crates/database/src/retention.rs

@@ -63,6 +63,7 @@ use common::{
     persistence::{
         new_static_repeatable_recent,
         DocumentLogEntry,
+        DocumentPrevTsQuery,
         NoopRetentionValidator,
         Persistence,
         PersistenceGlobalKey,
@@ -561,26 +562,47 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
                 // Each prev rev has 1 or 2 index entries to delete per index -- one entry at
                 // the prev rev's ts, and a tombstone at the current rev's ts if
                 // the document was deleted or its index key changed.
-                // TODO: use prev_ts when available
                 let prev_revs = reader_
-                    .previous_revisions(chunk.iter().map(|entry| (entry.id, entry.ts)).collect())
+                    .previous_revisions_of_documents(
+                        chunk
+                            .iter()
+                            .filter_map(|entry| {
+                                entry.prev_ts.map(|prev_ts| DocumentPrevTsQuery {
+                                    id: entry.id,
+                                    ts: entry.ts,
+                                    prev_ts,
+                                })
+                            })
+                            .collect(),
+                    )
                     .await?;
                 for DocumentLogEntry {
                     ts,
                     id,
                     value: maybe_doc,
+                    prev_ts,
                     ..
                 } in chunk
                 {
                     // If there is no prev rev, there's nothing to delete.
                     // If this happens for a tombstone, it means the document was created and
                     // deleted in the same transaction, with no index rows.
+                    let Some(prev_ts) = prev_ts else {
+                        log_retention_scanned_document(maybe_doc.is_none(), false);
+                        continue;
+                    };
                     let Some(DocumentLogEntry {
                         ts: prev_rev_ts,
                         value: maybe_prev_rev,
                         ..
-                    }) = prev_revs.get(&(id, ts))
+                    }) = prev_revs.get(&DocumentPrevTsQuery { id, ts, prev_ts })
                     else {
+                        // This is unexpected: if there is a prev_ts, there should be a prev_rev.
+                        report_error(&mut anyhow::anyhow!(
+                            "Skipping deleting indexes for {id}@{ts}. It has a prev_ts of \
+                             {prev_ts} but no previous revision."
+                        ))
+                        .await;
                         log_retention_scanned_document(maybe_doc.is_none(), false);
                         continue;
                     };

crates/mysql/src/chunks.rs

@@ -8,7 +8,10 @@ use common::{
         MYSQL_MAX_CHUNK_BYTES,
         MYSQL_MAX_DYNAMIC_SMART_CHUNK_SIZE,
     },
-    persistence::DocumentLogEntry,
+    persistence::{
+        DocumentLogEntry,
+        DocumentPrevTsQuery,
+    },
     types::{
         DatabaseIndexUpdate,
         Timestamp,
@@ -35,6 +38,12 @@ impl ApproxSize for Timestamp {
     }
 }
 
+impl ApproxSize for DocumentPrevTsQuery {
+    fn approx_size(&self) -> usize {
+        self.id.approx_size() + self.ts.approx_size() + self.prev_ts.approx_size()
+    }
+}
+
 impl ApproxSize for InternalDocumentId {
     fn approx_size(&self) -> usize {
         self.size()