add sequence_watermark(name) function and track sequence allocations in mvcc mode

Author: PThorpe92

core/connection.rs

@@ -3097,6 +3097,10 @@ impl Connection {
                 continue;
             }
             let seq = self.read_sequence_descriptor_via_sql(&backing_table_name, &seq_name)?;
+            let watermark = self.read_sequence_watermark_via_sql(&backing_table_name, &seq)?;
+            if let Some(mv_store) = self.db.get_mv_store().as_ref() {
+                mv_store.set_sequence_watermark(&normalized, watermark);
+            }
             self.with_database_schema_mut(MAIN_DB_ID, |schema| {
                 schema.sequences.insert(normalized.clone(), Arc::new(seq));
             });
@@ -3162,6 +3166,53 @@ impl Connection {
         })
     }
 
+    /// Read the current exclusive sequence watermark from the backing table.
+    ///
+    /// The returned value is the upper bound used by `sequence_watermark()`:
+    /// callers can consume ascending sequence rows with values lower than it.
+    fn read_sequence_watermark_via_sql(
+        self: &Arc<Connection>,
+        backing_table_name: &str,
+        seq: &crate::schema::Sequence,
+    ) -> Result<i64> {
+        let escaped = backing_table_name.replace('"', "\"\"");
+        let direction = if seq.increment_by >= 0 { "DESC" } else { "ASC" };
+        let sql = format!(
+            "SELECT value, is_called FROM \"{escaped}\" ORDER BY value {direction} LIMIT 1"
+        );
+        let mut stmt = self.prepare_internal(sql).map_err(|err| {
+            LimboError::Corrupt(format!(
+                "internal sequence backing table \"{backing_table_name}\" for sequence \
+                 \"{}\": cannot prepare watermark SELECT: {err}",
+                seq.name
+            ))
+        })?;
+        let mut row: Option<(i64, bool)> = None;
+        stmt.run_with_row_callback(|r| {
+            let value = r.get::<i64>(0)?;
+            let is_called = r.get::<i64>(1)? != 0;
+            row = Some((value, is_called));
+            Ok(())
+        })
+        .map_err(|err| {
+            LimboError::Corrupt(format!(
+                "internal sequence backing table \"{backing_table_name}\" for sequence \
+                 \"{}\": watermark row read failed: {err}",
+                seq.name
+            ))
+        })?;
+        let (value, is_called) = row.ok_or_else(|| {
+            LimboError::Corrupt(format!(
+                "internal sequence backing table \"{backing_table_name}\" for sequence \
+                 \"{}\" is empty; cannot derive sequence watermark",
+                seq.name
+            ))
+        })?;
+        Ok(crate::mvcc::database::first_unsafe_sequence_watermark(
+            seq, value, is_called,
+        ))
+    }
+
     /// Sync AUTOINCREMENT backing-table watermarks from `sqlite_sequence`.
     ///
     /// Covers the WAL→MVCC mode-switch compatibility path: a WAL-mode
@@ -3258,6 +3309,13 @@ impl Connection {
             );
             let mut insert_stmt = self.prepare_internal(insert_sql)?;
             insert_stmt.run_with_row_callback(|_| Ok(()))?;
+            if let Some(mv_store) = self.db.get_mv_store().as_ref() {
+                let first_unsafe = watermark.checked_add(1).unwrap_or(watermark);
+                mv_store.set_sequence_watermark(
+                    &autoincrement_sequence_name(&table_name),
+                    first_unsafe,
+                );
+            }
         }
 
         Ok(())

core/function.rs

@@ -784,6 +784,7 @@ pub enum ScalarFunc {
     StatGet,
     ConnTxnId,
     IsAutocommit,
+    SequenceWatermark,
     // Test type functions (for custom type system testing)
     TestUintEncode,
     TestUintDecode,
@@ -910,6 +911,7 @@ impl Deterministic for ScalarFunc {
             ScalarFunc::StatGet => false,  // internal ANALYZE function
             ScalarFunc::ConnTxnId => false, // depends on connection state
             ScalarFunc::IsAutocommit => false, // depends on connection state
+            ScalarFunc::SequenceWatermark => false, // depends on active MVCC transactions
             ScalarFunc::TestUintEncode
             | ScalarFunc::TestUintDecode
             | ScalarFunc::TestUintAdd
@@ -1049,6 +1051,7 @@ impl Display for ScalarFunc {
             Self::StatGet => "stat_get",
             Self::ConnTxnId => "conn_txn_id",
             Self::IsAutocommit => "is_autocommit",
+            Self::SequenceWatermark => "sequence_watermark",
             Self::TestUintEncode => "test_uint_encode",
             Self::TestUintDecode => "test_uint_decode",
             Self::TestUintAdd => "test_uint_add",
@@ -1152,7 +1155,8 @@ impl ScalarFunc {
             | Self::Upper
             | Self::ZeroBlob
             | Self::Likely
-            | Self::Unlikely => &[1],
+            | Self::Unlikely
+            | Self::SequenceWatermark => &[1],
             // 2-arg
             Self::Glob
             | Self::Instr
@@ -1711,6 +1715,7 @@ impl Func {
             "bin_record_json_object" => Ok(Some(Self::Scalar(ScalarFunc::BinRecordJsonObject))),
             "conn_txn_id" => Ok(Some(Self::Scalar(ScalarFunc::ConnTxnId))),
             "is_autocommit" => Ok(Some(Self::Scalar(ScalarFunc::IsAutocommit))),
+            "sequence_watermark" => Ok(Some(Self::Scalar(ScalarFunc::SequenceWatermark))),
             "acos" => Ok(Some(Self::Math(MathFunc::Acos))),
             "acosh" => Ok(Some(Self::Math(MathFunc::Acosh))),
             "asin" => Ok(Some(Self::Math(MathFunc::Asin))),

core/mvcc/database/mod.rs

@@ -3,7 +3,7 @@ use crate::mvcc::cursor::{static_iterator_hack, MvccIterator};
 #[cfg(any(test, injected_yields))]
 use crate::mvcc::yield_hooks::{ProvidesYieldContext, YieldContext, YieldPointMarker};
 use crate::mvcc::yield_points::{inject_transition_failure, inject_transition_yield};
-use crate::schema::{Schema, Table};
+use crate::schema::{Schema, Sequence, Table};
 use crate::state_machine::StateMachine;
 use crate::state_machine::StateTransition;
 use crate::state_machine::TransitionResult;
@@ -44,7 +44,7 @@ use crossbeam_skiplist::map::Entry;
 use crossbeam_skiplist::SkipMap;
 use rustc_hash::FxHashMap as HashMap;
 use rustc_hash::FxHashSet as HashSet;
-use std::collections::BTreeSet;
+use std::collections::{BTreeSet, HashMap as StdHashMap};
 use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::ops::Bound;
@@ -80,6 +80,20 @@ pub mod tests;
 /// Sentinel value for `MvStore::exclusive_tx` indicating no exclusive transaction is active.
 const NO_EXCLUSIVE_TX: u64 = 0;
 
+/// Convert a sequence backing-table row into the exclusive upper bound used by
+/// sync scans. This is intentionally tailored to ascending non-CYCLE sequences,
+/// which is the shape used by AUTOINCREMENT CDC ids.
+pub(crate) fn first_unsafe_sequence_watermark(seq: &Sequence, value: i64, is_called: bool) -> i64 {
+    if !is_called {
+        return value;
+    }
+    if seq.increment_by > 0 && !seq.cycle {
+        value.checked_add(seq.increment_by).unwrap_or(value)
+    } else {
+        value
+    }
+}
+
 #[cfg(not(any(test, injected_yields)))]
 struct YieldContext;
 
@@ -3437,6 +3451,15 @@ pub struct MvStore<Clock: LogicalClock> {
     /// deadlock.
     last_global_header_ts: AtomicU64,
     table_id_to_last_rowid: RwLock<HashMap<MVTableId, Arc<RowidAllocator>>>,
+    /// Per-sequence first value not guaranteed safe to read past based only on
+    /// durable/current sequence state. Active allocations can lower this.
+    sequence_watermarks: Mutex<HashMap<String, i64>>,
+    /// Per-sequence minimum allocated value for each active transaction.
+    ///
+    /// This is in-memory and therefore only correct while all MVCC writers for a
+    /// database live in one process. Multi-process MVCC will need a shared
+    /// coordination mechanism before sync can rely on this watermark.
+    sequence_allocations: Mutex<HashMap<String, StdHashMap<TxID, i64>>>,
 }
 
 impl<Clock: LogicalClock> MvStore<Clock> {
@@ -3512,6 +3535,8 @@ impl<Clock: LogicalClock> MvStore<Clock> {
             last_committed_tx_ts: AtomicU64::new(0),
             last_global_header_ts: AtomicU64::new(0),
             table_id_to_last_rowid: RwLock::new(HashMap::default()),
+            sequence_watermarks: Mutex::new(HashMap::default()),
+            sequence_allocations: Mutex::new(HashMap::default()),
         }
     }
 
@@ -4744,6 +4769,7 @@ impl<Clock: LogicalClock> MvStore<Clock> {
     }
 
     pub fn remove_tx(&self, tx_id: TxID) {
+        self.remove_sequence_allocations(tx_id);
         if let Some(entry) = self.txs.get(&tx_id) {
             let tx = entry.value();
             if let TransactionState::Committed(commit_ts) = tx.state.load() {
@@ -4768,6 +4794,87 @@ impl<Clock: LogicalClock> MvStore<Clock> {
         self.blocking_checkpoint_lock.unlock();
     }
 
+    pub fn register_sequence_allocation(
+        &self,
+        tx_id: TxID,
+        sequence_name: &str,
+        sequence_value: i64,
+    ) -> Result<()> {
+        let Some(tx) = self.txs.get(&tx_id) else {
+            return Err(LimboError::NoSuchTransactionID(tx_id.to_string()));
+        };
+        turso_assert!(
+            matches!(
+                tx.value().state.load(),
+                TransactionState::Active | TransactionState::Preparing(_)
+            ),
+            "sequence allocation must be registered while the transaction is active or preparing"
+        );
+
+        let sequence_name = crate::util::normalize_ident(sequence_name);
+        let mut allocations = self.sequence_allocations.lock();
+        let tx_allocations = allocations.entry(sequence_name).or_default();
+        tx_allocations
+            .entry(tx_id)
+            .and_modify(|value| *value = (*value).min(sequence_value))
+            .or_insert(sequence_value);
+        Ok(())
+    }
+
+    pub fn set_sequence_watermark(&self, sequence_name: &str, watermark: i64) {
+        let sequence_name = crate::util::normalize_ident(sequence_name);
+        self.sequence_watermarks
+            .lock()
+            .insert(sequence_name, watermark);
+    }
+
+    /// Returns the first sequence value that is not safe for cursor scans to pass.
+    ///
+    /// Readers can safely consume rows with sequence values less than this
+    /// watermark. The value is the minimum of the current sequence boundary and
+    /// any lower value already allocated by an active transaction.
+    pub fn sequence_watermark(&self, sequence_name: &str) -> Option<i64> {
+        let sequence_name = crate::util::normalize_ident(sequence_name);
+        let mut allocations = self.sequence_allocations.lock();
+        let mut remove_allocations = false;
+        let active_watermark = {
+            allocations
+                .get_mut(&sequence_name)
+                .and_then(|tx_allocations| {
+                    tx_allocations.retain(|tx_id, _| {
+                        self.txs.get(tx_id).is_some_and(|tx| {
+                            matches!(
+                                tx.value().state.load(),
+                                TransactionState::Active | TransactionState::Preparing(_)
+                            )
+                        })
+                    });
+                    let watermark = tx_allocations.values().copied().min();
+                    if tx_allocations.is_empty() {
+                        remove_allocations = true;
+                    }
+                    watermark
+                })
+        };
+        if remove_allocations {
+            allocations.remove(&sequence_name);
+        }
+        let current_watermark = self.sequence_watermarks.lock().get(&sequence_name).copied();
+        match (current_watermark, active_watermark) {
+            (Some(current), Some(active)) => Some(current.min(active)),
+            (Some(current), None) => Some(current),
+            (None, active) => active,
+        }
+    }
+
+    fn remove_sequence_allocations(&self, tx_id: TxID) {
+        let mut allocations = self.sequence_allocations.lock();
+        allocations.retain(|_, tx_allocations| {
+            tx_allocations.remove(&tx_id);
+            !tx_allocations.is_empty()
+        });
+    }
+
     /// Atomically retire a committed tx: clear the connection's mv_tx_id cache
     /// for `db_id`, then remove the tx from `txs`. Pairs the two mutations so
     /// no concurrent observer (or in-flight statement) can see the divergent

core/mvcc/database/tests.rs

@@ -4070,6 +4070,100 @@ pub(crate) fn commit_tx_no_conn(
     Ok(())
 }
 
+#[test]
+fn test_sequence_watermark_tracks_lowest_active_allocation() {
+    let db = MvccTestDbNoConn::new_with_random_db();
+    let conn = db.connect();
+    let mv_store = db.get_mvcc_store();
+    let pager = conn.pager.load().clone();
+
+    mv_store.set_sequence_watermark("turso_cdc_pk_autoincrement", 13);
+    let tx1 = mv_store.begin_tx(pager.clone()).unwrap();
+    let tx2 = mv_store.begin_tx(pager.clone()).unwrap();
+    mv_store
+        .register_sequence_allocation(tx1, "turso_cdc_pk_autoincrement", 10)
+        .unwrap();
+    mv_store
+        .register_sequence_allocation(tx2, "turso_cdc_pk_autoincrement", 12)
+        .unwrap();
+    mv_store
+        .register_sequence_allocation(tx1, "turso_cdc_pk_autoincrement", 11)
+        .unwrap();
+
+    assert_eq!(
+        mv_store.sequence_watermark("turso_cdc_pk_autoincrement"),
+        Some(10)
+    );
+
+    commit_tx(mv_store.clone(), &conn, tx1).unwrap();
+    assert_eq!(
+        mv_store.sequence_watermark("turso_cdc_pk_autoincrement"),
+        Some(12)
+    );
+
+    mv_store.rollback_tx(tx2, pager, conn.as_ref(), crate::MAIN_DB_ID);
+    assert_eq!(
+        mv_store.sequence_watermark("turso_cdc_pk_autoincrement"),
+        Some(13)
+    );
+}
+
+#[test]
+fn test_sequence_watermark_function_returns_current_watermark_without_active_allocations() {
+    let db = MvccTestDbNoConn::new_with_random_db();
+    let conn = db.connect();
+    let mv_store = db.get_mvcc_store();
+    let pager = conn.pager.load().clone();
+
+    conn.execute("CREATE SEQUENCE s").unwrap();
+    mv_store.set_sequence_watermark("s", 42);
+    let rows = get_rows(&conn, "SELECT sequence_watermark('s')");
+
+    assert_eq!(rows.len(), 1);
+    assert_eq!(rows[0][0].as_int().unwrap(), 42);
+
+    let tx_id = mv_store.begin_tx(pager.clone()).unwrap();
+    mv_store
+        .register_sequence_allocation(tx_id, "s", 10)
+        .unwrap();
+    let rows = get_rows(&conn, "SELECT sequence_watermark('s')");
+
+    assert_eq!(rows[0][0].as_int().unwrap(), 10);
+
+    mv_store.rollback_tx(tx_id, pager, conn.as_ref(), crate::MAIN_DB_ID);
+    let rows = get_rows(&conn, "SELECT sequence_watermark('s')");
+
+    assert_eq!(rows[0][0].as_int().unwrap(), 42);
+}
+
+#[test]
+fn test_sequence_watermark_tracks_nextval_allocations() {
+    let db = MvccTestDbNoConn::new_with_random_db();
+    let setup = db.connect();
+    setup.execute("CREATE SEQUENCE s START WITH 1").unwrap();
+
+    let rows = get_rows(&setup, "SELECT sequence_watermark('s')");
+    assert!(matches!(rows[0][0], Value::Null));
+
+    let rows = get_rows(&setup, "SELECT nextval('s')");
+    assert_eq!(rows[0][0].as_int().unwrap(), 1);
+    let rows = get_rows(&setup, "SELECT sequence_watermark('s')");
+    assert_eq!(rows[0][0].as_int().unwrap(), 2);
+
+    let writer = db.connect();
+    writer.execute("BEGIN CONCURRENT").unwrap();
+    let rows = get_rows(&writer, "SELECT nextval('s')");
+    assert_eq!(rows[0][0].as_int().unwrap(), 2);
+
+    let observer = db.connect();
+    let rows = get_rows(&observer, "SELECT sequence_watermark('s')");
+    assert_eq!(rows[0][0].as_int().unwrap(), 2);
+
+    writer.execute("COMMIT").unwrap();
+    let rows = get_rows(&observer, "SELECT sequence_watermark('s')");
+    assert_eq!(rows[0][0].as_int().unwrap(), 3);
+}
+
 /// What this test checks: Cursor traversal and seek operations honor MVCC visibility and key ordering under updates/deletes.
 /// Why this matters: Read-path correctness is critical: wrong cursor semantics directly surface as wrong query answers.
 #[test]

core/translate/expr/translator.rs

@@ -1235,7 +1235,8 @@ pub fn translate_expr(
                         | ScalarFunc::RandomBlob
                         | ScalarFunc::Sign
                         | ScalarFunc::Soundex
-                        | ScalarFunc::ZeroBlob => {
+                        | ScalarFunc::ZeroBlob
+                        | ScalarFunc::SequenceWatermark => {
                             let args = expect_arguments_exact!(args, 1, srf);
                             let start_reg = program.alloc_register();
                             translate_expr(