Thread prev_ts through DocumentUpdate (#33295)

This adds prev_ts to the old_document of `tx.apply_validated_write`, requiring the caller to thread the ts from whatever document they queried. This is then threaded all the way through DocumentUpdate and into the persistence layer.

In particular, the caller passes a `WriteTimestamp`, which may be `Pending`; `Writes::update` is responsible for collapsing the prev_ts chain in that case.

Some users of DocumentUpdate - primarily around index updates - do not care about the prev_ts. I have split those into a separate type, `IndexDocumentUpdate`.

There is a push safety consideration because we serialize the write set and send it around. As such I have marked the prev_ts as optional almost everywhere to account for the case where it gets lost during an RPC. This can get cleaned up after a push cycle.

GitOrigin-RevId: f6cca7c59dce42d159780486d30142ebdae7d3c3

Author: goffrie

crates/common/src/document.rs

@@ -28,6 +28,7 @@ use packed_value::{
 };
 use pb::common::{
     DocumentUpdate as DocumentUpdateProto,
+    DocumentUpdateWithPrevTs as DocumentUpdateWithPrevTsProto,
     ResolvedDocument as ResolvedDocumentProto,
 };
 #[cfg(any(test, feature = "testing"))]
@@ -588,18 +589,91 @@ impl ResolvedDocument {
 
 #[derive(Clone, Debug, Eq, PartialEq)]
 #[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
-pub struct DocumentUpdate {
+pub struct DocumentUpdateWithPrevTs {
     pub id: ResolvedDocumentId,
-    pub old_document: Option<ResolvedDocument>,
+    /// The old document and its timestamp in the document log.
+    /// The timestamp will become the update's `prev_ts`.
+    // TODO: make the timestamp non-optional after everything has pushed
+    pub old_document: Option<(ResolvedDocument, Option<Timestamp>)>,
     pub new_document: Option<ResolvedDocument>,
 }
 
-impl HeapSize for DocumentUpdate {
+impl DocumentUpdateWithPrevTs {
+    /// Checks if two DocumentUpdates are almost equal, ignoring the case where
+    /// one has a missing old timestamp
+    // TODO: remove this once the old timestamp is non-optional
+    pub fn eq_ignoring_none_old_ts(&self, other: &DocumentUpdateWithPrevTs) -> bool {
+        self.id == other.id
+            && self.old_document.as_ref().map(|(d, _)| d)
+                == other.old_document.as_ref().map(|(d, _)| d)
+            && self
+                .old_document
+                .as_ref()
+                .map(|(_, ts)| ts)
+                .zip(other.old_document.as_ref().map(|(_, ts)| ts))
+                .is_none_or(|(a, b)| a == b)
+            && self.new_document == other.new_document
+    }
+}
+
+impl HeapSize for DocumentUpdateWithPrevTs {
     fn heap_size(&self) -> usize {
         self.old_document.heap_size() + self.new_document.heap_size()
     }
 }
 
+impl TryFrom<DocumentUpdateWithPrevTs> for DocumentUpdateWithPrevTsProto {
+    type Error = anyhow::Error;
+
+    fn try_from(
+        DocumentUpdateWithPrevTs {
+            id,
+            old_document,
+            new_document,
+        }: DocumentUpdateWithPrevTs,
+    ) -> anyhow::Result<Self> {
+        let (old_document, old_ts) = old_document.unzip();
+        Ok(Self {
+            id: Some(id.into()),
+            old_document: old_document.map(|d| d.try_into()).transpose()?,
+            old_ts: old_ts.flatten().map(|ts| ts.into()),
+            new_document: new_document.map(|d| d.try_into()).transpose()?,
+        })
+    }
+}
+
+impl TryFrom<DocumentUpdateWithPrevTsProto> for DocumentUpdateWithPrevTs {
+    type Error = anyhow::Error;
+
+    fn try_from(
+        DocumentUpdateWithPrevTsProto {
+            id,
+            old_document,
+            old_ts,
+            new_document,
+        }: DocumentUpdateWithPrevTsProto,
+    ) -> anyhow::Result<Self> {
+        let id = id
+            .context("Document updates missing document id")?
+            .try_into()?;
+        Ok(Self {
+            id,
+            old_document: old_document
+                .map(|d| anyhow::Ok((d.try_into()?, old_ts.map(Timestamp::try_from).transpose()?)))
+                .transpose()?,
+            new_document: new_document.map(|d| d.try_into()).transpose()?,
+        })
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct DocumentUpdate {
+    pub id: ResolvedDocumentId,
+    pub old_document: Option<ResolvedDocument>,
+    pub new_document: Option<ResolvedDocument>,
+}
+
 impl TryFrom<DocumentUpdate> for DocumentUpdateProto {
     type Error = anyhow::Error;
 
@@ -639,6 +713,51 @@ impl TryFrom<DocumentUpdateProto> for DocumentUpdate {
     }
 }
 
+impl From<DocumentUpdateWithPrevTs> for DocumentUpdate {
+    fn from(update: DocumentUpdateWithPrevTs) -> Self {
+        Self {
+            id: update.id,
+            old_document: update.old_document.map(|(d, _)| d),
+            new_document: update.new_document,
+        }
+    }
+}
+
+/// Either a [`DocumentUpdate`] or a [`DocumentUpdateWithPrevTs`]
+pub trait DocumentUpdateRef {
+    fn id(&self) -> ResolvedDocumentId;
+    fn old_document(&self) -> Option<&ResolvedDocument>;
+    fn new_document(&self) -> Option<&ResolvedDocument>;
+}
+
+impl DocumentUpdateRef for DocumentUpdateWithPrevTs {
+    fn id(&self) -> ResolvedDocumentId {
+        self.id
+    }
+
+    fn old_document(&self) -> Option<&ResolvedDocument> {
+        self.old_document.as_ref().map(|(d, _)| d)
+    }
+
+    fn new_document(&self) -> Option<&ResolvedDocument> {
+        self.new_document.as_ref()
+    }
+}
+
+impl DocumentUpdateRef for DocumentUpdate {
+    fn id(&self) -> ResolvedDocumentId {
+        self.id
+    }
+
+    fn old_document(&self) -> Option<&ResolvedDocument> {
+        self.old_document.as_ref()
+    }
+
+    fn new_document(&self) -> Option<&ResolvedDocument> {
+        self.new_document.as_ref()
+    }
+}
+
 impl DeveloperDocument {
     pub fn new(
         id: DeveloperDocumentId,
@@ -899,13 +1018,15 @@ mod tests {
 
     use super::{
         DocumentUpdateProto,
+        DocumentUpdateWithPrevTsProto,
         ResolvedDocumentProto,
     };
     use crate::{
         assert_obj,
         document::{
             CreationTime,
             DocumentUpdate,
+            DocumentUpdateWithPrevTs,
             ResolvedDocument,
         },
         paths::FieldPath,
@@ -951,7 +1072,13 @@ mod tests {
 
 
         #[test]
-        fn test_document_update_proto_roundtrips(left in any::<DocumentUpdate>()) {
+        fn test_document_update_proto_roundtrips(left in any::<DocumentUpdateWithPrevTs>()) {
+            assert_roundtrips::<DocumentUpdateWithPrevTs, DocumentUpdateWithPrevTsProto>(left);
+        }
+
+
+        #[test]
+        fn test_index_document_update_proto_roundtrips(left in any::<DocumentUpdate>()) {
             assert_roundtrips::<DocumentUpdate, DocumentUpdateProto>(left);
         }
     }

crates/database/src/bootstrap_model/import_facing.rs

@@ -187,7 +187,7 @@ impl<'a, RT: Runtime> ImportFacingModel<'a, RT> {
             .table_mapping()
             .tablet_namespace(table_id.tablet_id)?;
 
-        let existing_doc = self.tx.get(id).await?;
+        let existing_doc = self.tx.get_with_ts(id).await?;
 
         let creation_time_field = FieldName::from(CREATION_TIME_FIELD.clone());
         let creation_time = if let Some(ConvexValue::Float64(f)) = value.get(&creation_time_field) {
@@ -236,7 +236,7 @@ impl<'a, RT: Runtime> ImportFacingModel<'a, RT> {
         }
 
         let id = ResolvedDocumentId::new(table_id.tablet_id, developer_id);
-        let existing_doc = self.tx.get(id).await?;
+        let existing_doc = self.tx.get_with_ts(id).await?;
 
         self.tx.apply_validated_write(id, existing_doc, None)?;
 

crates/database/src/committer.rs

@@ -20,7 +20,7 @@ use common::{
         ComponentPath,
     },
     document::{
-        DocumentUpdate,
+        DocumentUpdateWithPrevTs,
         ParsedDocument,
         ResolvedDocument,
     },
@@ -657,7 +657,7 @@ impl<RT: Runtime> Committer<RT> {
             commit_ts,
             ordered_updates
                 .into_iter()
-                .map(|(id, update)| (id, PackedDocumentUpdate::pack(update)))
+                .map(|(id, update)| (id, PackedDocumentUpdate::pack(update.into())))
                 .collect(),
             write_source,
         );
@@ -673,7 +673,7 @@ impl<RT: Runtime> Committer<RT> {
     fn compute_writes(
         &self,
         commit_ts: Timestamp,
-        ordered_updates: &Vec<(ResolvedDocumentId, DocumentUpdate)>,
+        ordered_updates: &Vec<(ResolvedDocumentId, DocumentUpdateWithPrevTs)>,
     ) -> anyhow::Result<(
         Vec<ValidatedDocumentWrite>,
         BTreeSet<(Timestamp, DatabaseIndexUpdate)>,
@@ -696,6 +696,10 @@ impl<RT: Runtime> Committer<RT> {
                     document: document_update.new_document.clone(),
                 },
                 doc_in_vector_index,
+                prev_ts: document_update
+                    .old_document
+                    .as_ref()
+                    .and_then(|&(_, ts)| ts),
             });
         }
         let index_writes = index_writes
@@ -736,14 +740,11 @@ impl<RT: Runtime> Committer<RT> {
         let timer = metrics::commit_persistence_write_timer();
         let document_writes = document_writes
             .into_iter()
-            .map(|write| {
-                DocumentLogEntry {
-                    ts: write.commit_ts,
-                    id: write.id,
-                    value: write.write.document,
-                    // TODO: fill in prev_ts
-                    prev_ts: None,
-                }
+            .map(|write| DocumentLogEntry {
+                ts: write.commit_ts,
+                id: write.id,
+                value: write.write.document,
+                prev_ts: write.prev_ts,
             })
             .collect();
         persistence
@@ -999,6 +1000,7 @@ struct ValidatedDocumentWrite {
     id: InternalDocumentId,
     write: DocumentWrite,
     doc_in_vector_index: DocInVectorIndex,
+    prev_ts: Option<Timestamp>,
 }
 
 #[derive(Clone)]

crates/database/src/snapshot_manager.rs

@@ -12,7 +12,7 @@ use common::{
         ComponentPath,
     },
     document::{
-        DocumentUpdate,
+        DocumentUpdateRef,
         ResolvedDocument,
     },
     knobs::MAX_TRANSACTION_WINDOW,
@@ -211,13 +211,13 @@ pub struct Snapshot {
 impl Snapshot {
     pub(crate) fn update(
         &mut self,
-        document_update: &DocumentUpdate,
+        document_update: &impl DocumentUpdateRef,
         commit_ts: Timestamp,
     ) -> anyhow::Result<(Vec<DatabaseIndexUpdate>, DocInVectorIndex)> {
         block_in_place(|| {
-            let removal = document_update.old_document.as_ref();
-            let insertion = document_update.new_document.as_ref();
-            let document_id = document_update.id;
+            let removal = document_update.old_document();
+            let insertion = document_update.new_document();
+            let document_id = document_update.id();
             let table_update = self
                 .table_registry
                 .update(