Simplify representation of TimestampRange & make RepeatablePersistence::load_documents pre-filter the range (#40853)

Makes TimestampRange a simple inclusive range internally and make operations on it infallible.
Then add `TimestampRange::intersect` so that we can get rid of post-filtering on `RepeatablePersistence::load_documents` & similar methods,
by intersecting the provided TimestampRange with the upper bound.

GitOrigin-RevId: a3259d4556f23602e981a8f7cc38a29969e409bd

Author: goffrie

crates/common/src/persistence.rs

@@ -272,75 +272,91 @@ pub trait Persistence: Sync + Send + 'static {
 
 #[derive(Debug, Clone, Copy)]
 pub struct TimestampRange {
-    start_bound: Bound<Timestamp>,
-    end_bound: Bound<Timestamp>,
+    start_inclusive: Timestamp,
+    end_inclusive: Timestamp,
 }
 
 impl TimestampRange {
-    pub fn new<T: RangeBounds<Timestamp>>(range: T) -> anyhow::Result<Self> {
-        // Bounds check.
-        Self::min_inclusive(&range.start_bound().cloned())?;
-        Self::max_exclusive(&range.end_bound().cloned())?;
-        Ok(Self {
-            start_bound: range.start_bound().cloned(),
-            end_bound: range.end_bound().cloned(),
-        })
-    }
-
-    pub fn snapshot(ts: Timestamp) -> Self {
+    #[inline]
+    pub fn new<T: RangeBounds<Timestamp>>(range: T) -> Self {
+        let start_inclusive = match range.start_bound() {
+            Bound::Included(t) => *t,
+            Bound::Excluded(t) => {
+                if let Some(succ) = t.succ_opt() {
+                    succ
+                } else {
+                    return Self::empty();
+                }
+            },
+            Bound::Unbounded => Timestamp::MIN,
+        };
+        let end_inclusive = match range.end_bound() {
+            Bound::Included(t) => *t,
+            Bound::Excluded(t) => {
+                if let Some(pred) = t.pred_opt() {
+                    pred
+                } else {
+                    return Self::empty();
+                }
+            },
+            Bound::Unbounded => Timestamp::MAX,
+        };
         Self {
-            start_bound: Bound::Unbounded,
-            end_bound: Bound::Included(ts),
+            start_inclusive,
+            end_inclusive,
         }
     }
 
-    pub fn all() -> Self {
+    #[inline]
+    pub fn empty() -> Self {
         Self {
-            start_bound: Bound::Unbounded,
-            end_bound: Bound::Unbounded,
+            start_inclusive: Timestamp::MAX,
+            end_inclusive: Timestamp::MIN,
         }
     }
 
-    pub fn at(ts: Timestamp) -> Self {
-        Self {
-            start_bound: Bound::Included(ts),
-            end_bound: Bound::Included(ts),
-        }
+    #[inline]
+    pub fn snapshot(ts: Timestamp) -> Self {
+        Self::new(..=ts)
     }
 
-    pub fn greater_than(t: Timestamp) -> Self {
-        Self {
-            start_bound: Bound::Excluded(t),
-            end_bound: Bound::Unbounded,
-        }
+    #[inline]
+    pub fn all() -> Self {
+        Self::new(..)
     }
 
-    fn min_inclusive(start_bound: &Bound<Timestamp>) -> anyhow::Result<Timestamp> {
-        Ok(match start_bound {
-            Bound::Included(t) => *t,
-            Bound::Excluded(t) => t.succ()?,
-            Bound::Unbounded => Timestamp::MIN,
-        })
+    #[inline]
+    pub fn at(ts: Timestamp) -> Self {
+        Self::new(ts..=ts)
     }
 
-    pub fn min_timestamp_inclusive(&self) -> Timestamp {
-        Self::min_inclusive(&self.start_bound).unwrap()
+    #[inline]
+    pub fn greater_than(t: Timestamp) -> Self {
+        Self::new((Bound::Excluded(t), Bound::Unbounded))
     }
 
-    fn max_exclusive(end_bound: &Bound<Timestamp>) -> anyhow::Result<Timestamp> {
-        Ok(match end_bound {
-            Bound::Included(t) => t.succ()?,
-            Bound::Excluded(t) => *t,
-            Bound::Unbounded => Timestamp::MAX,
-        })
+    #[inline]
+    pub fn min_timestamp_inclusive(&self) -> Timestamp {
+        self.start_inclusive
     }
 
+    #[inline]
     pub fn max_timestamp_exclusive(&self) -> Timestamp {
-        Self::max_exclusive(&self.end_bound).unwrap()
+        // assumes that Timestamp::MAX never actually exists
+        self.end_inclusive.succ_opt().unwrap_or(Timestamp::MAX)
     }
 
+    #[inline]
     pub fn contains(&self, ts: Timestamp) -> bool {
-        self.min_timestamp_inclusive() <= ts && ts < self.max_timestamp_exclusive()
+        self.start_inclusive <= ts && ts <= self.end_inclusive
+    }
+
+    #[inline]
+    pub fn intersect(&self, other: Self) -> Self {
+        Self {
+            start_inclusive: self.start_inclusive.max(other.start_inclusive),
+            end_inclusive: self.end_inclusive.min(other.end_inclusive),
+        }
     }
 }
 
@@ -582,13 +598,12 @@ impl RepeatablePersistence {
     /// Same as [`Persistence::load_documents`] but only including documents in
     /// the snapshot range.
     pub fn load_documents(&self, range: TimestampRange, order: Order) -> DocumentStream<'_> {
-        let stream = self.reader.load_documents(
-            range,
+        self.reader.load_documents(
+            range.intersect(TimestampRange::snapshot(*self.upper_bound)),
             order,
             *DEFAULT_DOCUMENTS_PAGE_SIZE,
             self.retention_validator.clone(),
-        );
-        Box::pin(stream.try_filter(|entry| future::ready(entry.ts <= *self.upper_bound)))
+        )
     }
 
     /// Same as [`Persistence::load_documents_from_table`] but only including
@@ -599,14 +614,13 @@ impl RepeatablePersistence {
         range: TimestampRange,
         order: Order,
     ) -> DocumentStream<'_> {
-        let stream = self.reader.load_documents_from_table(
+        self.reader.load_documents_from_table(
             tablet_id,
-            range,
+            range.intersect(TimestampRange::snapshot(*self.upper_bound)),
             order,
             *DEFAULT_DOCUMENTS_PAGE_SIZE,
             self.retention_validator.clone(),
-        );
-        Box::pin(stream.try_filter(|entry| future::ready(entry.ts <= *self.upper_bound)))
+        )
     }
 
     /// Same as `load_documents` but doesn't use the `RetentionValidator` from
@@ -618,13 +632,12 @@ impl RepeatablePersistence {
         order: Order,
         retention_validator: Arc<dyn RetentionValidator>,
     ) -> DocumentStream<'_> {
-        let stream = self.reader.load_documents(
-            range,
+        self.reader.load_documents(
+            range.intersect(TimestampRange::snapshot(*self.upper_bound)),
             order,
             *DEFAULT_DOCUMENTS_PAGE_SIZE,
             retention_validator,
-        );
-        Box::pin(stream.try_filter(|entry| future::ready(entry.ts <= *self.upper_bound)))
+        )
     }
 
     pub async fn previous_revisions(

crates/common/src/testing/persistence_test_suite.rs

@@ -346,7 +346,7 @@ pub async fn write_and_load_from_table<P: Persistence>(p: Arc<P>) -> anyhow::Res
     test_load_documents_from_table(
         &p,
         doc1.id().tablet_id,
-        TimestampRange::new(Timestamp::must(1)..)?,
+        TimestampRange::new(Timestamp::must(1)..),
         Order::Asc,
         vec![DocumentLogEntry {
             ts: Timestamp::must(1),
@@ -360,7 +360,7 @@ pub async fn write_and_load_from_table<P: Persistence>(p: Arc<P>) -> anyhow::Res
     test_load_documents_from_table(
         &p,
         doc2.id().tablet_id,
-        TimestampRange::new(Timestamp::must(1)..)?,
+        TimestampRange::new(Timestamp::must(1)..),
         Order::Asc,
         vec![DocumentLogEntry {
             ts: Timestamp::must(1),
@@ -374,7 +374,7 @@ pub async fn write_and_load_from_table<P: Persistence>(p: Arc<P>) -> anyhow::Res
     test_load_documents_from_table(
         &p,
         doc1.id().tablet_id,
-        TimestampRange::new(..Timestamp::must(1))?,
+        TimestampRange::new(..Timestamp::must(1)),
         Order::Asc,
         vec![DocumentLogEntry {
             ts: Timestamp::must(0),
@@ -388,7 +388,7 @@ pub async fn write_and_load_from_table<P: Persistence>(p: Arc<P>) -> anyhow::Res
     test_load_documents_from_table(
         &p,
         doc2.id().tablet_id,
-        TimestampRange::new(..Timestamp::must(1))?,
+        TimestampRange::new(..Timestamp::must(1)),
         Order::Asc,
         vec![DocumentLogEntry {
             ts: Timestamp::must(0),
@@ -457,7 +457,7 @@ pub async fn write_and_load<P: Persistence>(p: Arc<P>) -> anyhow::Result<()> {
     test_load_documents(
         &p,
         &id_generator,
-        TimestampRange::new(Timestamp::must(1)..)?,
+        TimestampRange::new(Timestamp::must(1)..),
         Order::Asc,
         vec![DocumentLogEntry {
             ts: Timestamp::must(1),
@@ -485,7 +485,7 @@ pub async fn write_and_load<P: Persistence>(p: Arc<P>) -> anyhow::Result<()> {
     test_load_documents(
         &p,
         &id_generator,
-        TimestampRange::new(..Timestamp::must(2))?,
+        TimestampRange::new(..Timestamp::must(2)),
         Order::Desc,
         vec![
             DocumentLogEntry {

crates/convex/sync_types/src/timestamp.rs

@@ -7,7 +7,10 @@ use std::{
     },
 };
 
-use anyhow::Context;
+use anyhow::{
+    anyhow,
+    Context,
+};
 use derive_more::FromStr;
 use serde::Serialize;
 use serde_json::json;
@@ -25,18 +28,32 @@ impl Timestamp {
     pub const MAX: Self = Self(i64::MAX as u64);
     pub const MIN: Self = Self(0);
 
-    pub fn succ(&self) -> anyhow::Result<Self> {
+    #[inline]
+    pub fn succ_opt(&self) -> Option<Self> {
         if *self >= Self::MAX {
-            anyhow::bail!("timestamp {self} already at max");
+            None
+        } else {
+            Some(Self(self.0 + 1))
         }
-        Ok(Self(self.0 + 1))
     }
 
-    pub fn pred(&self) -> anyhow::Result<Self> {
+    pub fn succ(&self) -> anyhow::Result<Self> {
+        self.succ_opt()
+            .ok_or_else(|| anyhow!("timestamp {self} already at max"))
+    }
+
+    #[inline]
+    pub fn pred_opt(&self) -> Option<Self> {
         if *self <= Self::MIN {
-            anyhow::bail!("timestamp {self} already at min");
+            None
+        } else {
+            Some(Self(self.0 - 1))
         }
-        Ok(Self(self.0 - 1))
+    }
+
+    pub fn pred(&self) -> anyhow::Result<Self> {
+        self.pred_opt()
+            .ok_or_else(|| anyhow!("timestamp {self} already at min"))
     }
 
     pub fn add(&self, duration: Duration) -> anyhow::Result<Self> {

crates/database/src/committer.rs

@@ -421,7 +421,7 @@ impl<RT: Runtime> Committer<RT> {
             !vector_index_manager.is_bootstrapping(),
             "Trying to update vector index while it's still bootstrapping"
         );
-        let range = TimestampRange::new((Bound::Excluded(bootstrap_ts), Bound::Unbounded))?;
+        let range = TimestampRange::new((Bound::Excluded(bootstrap_ts), Bound::Unbounded));
 
         let revision_stream =
             stream_revision_pairs_for_indexes(tables_with_indexes, &persistence, range);

crates/database/src/database.rs

@@ -1821,7 +1821,7 @@ impl<RT: Runtime> Database<RT> {
             self.retention_validator(),
         );
         let range = match cursor {
-            Some(ts) => TimestampRange::new((Bound::Excluded(ts), Bound::Unbounded))?,
+            Some(ts) => TimestampRange::new((Bound::Excluded(ts), Bound::Unbounded)),
             None => TimestampRange::all(),
         };
         let mut document_stream = repeatable_persistence.load_documents(range, Order::Asc);

crates/database/src/database_index_workers/index_writer.rs

@@ -405,7 +405,7 @@ impl<RT: Runtime> IndexWriter<RT> {
         let producer = async {
             let revision_stream = self.stream_revision_pairs(
                 &repeatable_persistence,
-                TimestampRange::new(start_ts..=*end_ts)?,
+                TimestampRange::new(start_ts..=*end_ts),
                 Order::Asc,
                 index_selector,
             );
@@ -466,7 +466,7 @@ impl<RT: Runtime> IndexWriter<RT> {
         let producer = async {
             let revision_stream = self.stream_revision_pairs(
                 &repeatable_persistence,
-                TimestampRange::new(end_ts..*start_ts)?,
+                TimestampRange::new(end_ts..*start_ts),
                 Order::Desc,
                 index_selector,
             );

crates/database/src/retention.rs

@@ -575,7 +575,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         );
         let reader_ = &reader;
         let mut index_entry_chunks = reader
-            .load_documents(TimestampRange::new(*cursor..*min_snapshot_ts)?, Order::Asc)
+            .load_documents(TimestampRange::new(*cursor..*min_snapshot_ts), Order::Asc)
             .try_chunks2(*RETENTION_READ_CHUNK)
             .map(move |chunk| async move {
                 let chunk = chunk?;
@@ -815,7 +815,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         let reader_ = &reader;
         let mut document_chunks = reader
             .load_documents_with_retention_validator(
-                TimestampRange::new(*cursor..*min_document_snapshot_ts)?,
+                TimestampRange::new(*cursor..*min_document_snapshot_ts),
                 Order::Asc,
                 Arc::new(NoopRetentionValidator),
             )
@@ -1503,7 +1503,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         let reader = persistence.reader();
         let mut document_stream = reader.load_documents_from_table(
             index_table_id,
-            TimestampRange::new(**cursor..*latest_ts)?,
+            TimestampRange::new(**cursor..*latest_ts),
             Order::Asc,
             *DEFAULT_DOCUMENTS_PAGE_SIZE,
             retention_validator,

crates/database/src/search_index_bootstrap.rs

@@ -256,7 +256,7 @@ impl IndexesToBootstrap {
         let range = TimestampRange::new((
             Bound::Excluded(self.oldest_index_ts),
             Bound::Included(*upper_bound),
-        ))?;
+        ));
         let tables_with_indexes = self.tables_with_indexes();
         let revision_stream =
             stream_revision_pairs_for_indexes(&tables_with_indexes, persistence, range);

crates/database/src/search_index_workers/search_flusher.rs

@@ -604,7 +604,7 @@ impl<RT: Runtime, T: SearchIndex + 'static> SearchFlusher<RT, T> {
                         TimestampRange::new((
                             Bound::Excluded(*last_ts),
                             Bound::Included(*snapshot_ts),
-                        ))?,
+                        )),
                         T::partial_document_order(),
                         &row_rate_limiter,
                     ),

crates/database/src/search_index_workers/writer.rs

@@ -635,7 +635,7 @@ impl<RT: Runtime, T: SearchIndex> Inner<RT, T> {
             T::download_previous_segments(storage.clone(), segments_to_update).await?;
         let documents = database.load_documents_in_table(
             *index_name.table(),
-            TimestampRange::new((Bound::Excluded(start_ts), Bound::Included(*current_ts)))?,
+            TimestampRange::new((Bound::Excluded(start_ts), Bound::Included(*current_ts))),
             Order::Asc,
             &row_rate_limiter,
         );