Make prev_ts non-optional in DocumentUpdateWithPrevTs (#34636)

GitOrigin-RevId: fd9ea18f01f57c65b1d8d0946a0b71dd2906e678

Author: goffrie

crates/common/src/document.rs

@@ -593,29 +593,10 @@ pub struct DocumentUpdateWithPrevTs {
     pub id: ResolvedDocumentId,
     /// 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 old_document: Option<(ResolvedDocument, Timestamp)>,
     pub new_document: Option<ResolvedDocument>,
 }
 
-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()
@@ -636,7 +617,7 @@ impl TryFrom<DocumentUpdateWithPrevTs> for DocumentUpdateWithPrevTsProto {
         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()),
+            old_ts: old_ts.map(|ts| ts.into()),
             new_document: new_document.map(|d| d.try_into()).transpose()?,
         })
     }
@@ -659,7 +640,12 @@ impl TryFrom<DocumentUpdateWithPrevTsProto> for DocumentUpdateWithPrevTs {
         Ok(Self {
             id,
             old_document: old_document
-                .map(|d| anyhow::Ok((d.try_into()?, old_ts.map(Timestamp::try_from).transpose()?)))
+                .map(|d| {
+                    anyhow::Ok((
+                        d.try_into()?,
+                        Timestamp::try_from(old_ts.context("old_ts missing")?)?,
+                    ))
+                })
                 .transpose()?,
             new_document: new_document.map(|d| d.try_into()).transpose()?,
         })

crates/database/src/committer.rs

@@ -696,10 +696,7 @@ 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),
+                prev_ts: document_update.old_document.as_ref().map(|&(_, ts)| ts),
             });
         }
         let index_writes = index_writes

crates/database/src/transaction.rs

@@ -474,7 +474,7 @@ impl<RT: Runtime> Transaction<RT> {
             // already written to in this transaction.
             if let Some(existing_update) = existing_updates.get(&id) {
                 anyhow::ensure!(
-                    existing_update.eq_ignoring_none_old_ts(&update),
+                    *existing_update == update,
                     "Conflicting updates for document {id}"
                 );
                 preserved_update_count += 1;
@@ -490,11 +490,11 @@ impl<RT: Runtime> Transaction<RT> {
                     self.next_creation_time.increment()?;
                 }
             }
-            self.apply_validated_write_maybe_old_ts(
+            self.apply_validated_write(
                 id,
                 update
                     .old_document
-                    .map(|(d, ts)| (d, ts.map(WriteTimestamp::Committed))),
+                    .map(|(d, ts)| (d, WriteTimestamp::Committed(ts))),
                 update.new_document,
             )?;
         }
@@ -935,21 +935,6 @@ impl<RT: Runtime> Transaction<RT> {
         id: ResolvedDocumentId,
         old_document_and_ts: Option<(ResolvedDocument, WriteTimestamp)>,
         new_document: Option<ResolvedDocument>,
-    ) -> anyhow::Result<()> {
-        self.apply_validated_write_maybe_old_ts(
-            id,
-            old_document_and_ts.map(|(d, ts)| (d, Some(ts))),
-            new_document,
-        )
-    }
-
-    // TODO: make WriteTimestamp non-optional and merge with
-    // `apply_validated_write`
-    pub(crate) fn apply_validated_write_maybe_old_ts(
-        &mut self,
-        id: ResolvedDocumentId,
-        old_document_and_ts: Option<(ResolvedDocument, Option<WriteTimestamp>)>,
-        new_document: Option<ResolvedDocument>,
     ) -> anyhow::Result<()> {
         // Implement something like two-phase commit between the index and the document
         // store. We first guarantee that the changes are valid for the index and

crates/database/src/writes.rs

@@ -208,7 +208,7 @@ impl Writes {
         is_system_document: bool,
         reads: &mut TransactionReadSet,
         document_id: ResolvedDocumentId,
-        old_document: Option<(ResolvedDocument, Option<WriteTimestamp>)>,
+        old_document: Option<(ResolvedDocument, WriteTimestamp)>,
         new_document: Option<ResolvedDocument>,
     ) -> anyhow::Result<()> {
         if old_document.is_none() {
@@ -281,7 +281,7 @@ impl Writes {
                  document's old update"
             );
             anyhow::ensure!(
-                [None, Some(WriteTimestamp::Pending)].contains(&old_document_ts.flatten()),
+                [None, Some(WriteTimestamp::Pending)].contains(&old_document_ts),
                 "Inconsistent update: The new document's old update timestamp should be Pending \
                  but is {:?}",
                 old_document_ts
@@ -295,14 +295,13 @@ impl Writes {
                     old_document: match old_document {
                         Some((d, ts)) => Some((
                             d,
-                            ts.map(|ts| match ts {
-                                WriteTimestamp::Committed(ts) => Ok(ts),
+                            match ts {
+                                WriteTimestamp::Committed(ts) => ts,
                                 WriteTimestamp::Pending => anyhow::bail!(
                                     "Old document timestamp is Pending, but there is no pending \
                                      write"
                                 ),
-                            })
-                            .transpose()?,
+                            },
                         )),
                         None => None,
                     },
@@ -635,7 +634,7 @@ mod tests {
             id,
             Some((
                 old_document.clone(),
-                Some(WriteTimestamp::Committed(Timestamp::must(123))),
+                WriteTimestamp::Committed(Timestamp::must(123)),
             )),
             Some(new_document.clone()),
         )?;
@@ -649,7 +648,7 @@ mod tests {
             false,
             &mut reads,
             id,
-            Some((new_document, Some(WriteTimestamp::Pending))),
+            Some((new_document, WriteTimestamp::Pending)),
             Some(newer_document.clone()),
         )?;
 
@@ -660,7 +659,7 @@ mod tests {
                 id,
                 DocumentUpdateWithPrevTs {
                     id,
-                    old_document: Some((old_document, Some(Timestamp::must(123)))),
+                    old_document: Some((old_document, Timestamp::must(123))),
                     new_document: Some(newer_document),
                 }
             )

crates/pb/protos/common.proto

@@ -48,6 +48,7 @@ message DocumentUpdate {
 message DocumentUpdateWithPrevTs {
     ResolvedDocumentId id = 3;
     ResolvedDocument old_document = 1;
+    // `old_ts` should be present iff `old_document` is present
     optional uint64 old_ts = 4;
     ResolvedDocument new_document = 2;
 }