lock tables before capturing lsn, quick and dirty don't judge

Author: martykulma

src/sql-server-util/src/cdc.rs

@@ -123,6 +123,8 @@ impl<'a> CdcStream<'a> {
         ),
         SqlServerError,
     > {
+        static SAVEPOINT_NAME: &str = "_mz_snap_";
+
         // Determine what table we need to snapshot.
         let instances = self
             .capture_instances
@@ -147,10 +149,26 @@ impl<'a> CdcStream<'a> {
         let mut txn = self.client.transaction().await?;
         // Get the current LSN of the database. This operation assigns a savepoint to our
         // current transaction.
-        let lsn = crate::inspect::lsn_at_savepoint(&mut txn, "_mz_snapshot").await?;
 
+        // create a savepoint that we will lock the tables under and collect an LSN
+        // this allows us to rollback the savepoint while maintaining the outer transaction
+        txn.create_savepoint(SAVEPOINT_NAME).await?;
+        // lock all the tables we are planning to snapshot so that we can ensure that
+        // writes that might be in progress are properly ordered before or after this snapshot
+        // in addition to the LSN being properly ordered.
+        // TODO (maz): we should considering a timeout here because we may lock some tables,
+        // and the next table may be locked for some extended period, resulting in a traffic
+        // jam.
+        for (_capture_instance, schema, table) in &tables {
+            tracing::trace!(%schema, %table, "locking table");
+            crate::inspect::lock_table(&mut txn, &*schema, &*table).await?;
+        }
+
+        let lsn = crate::inspect::get_lsn(&mut txn).await?;
         tracing::info!(?tables, ?lsn, "starting snapshot");
 
+        txn.rollback_savepoint(SAVEPOINT_NAME).await?;
+
         // Get the size of each table we're about to snapshot.
         //
         // TODO(sql_server3): To expose a more "generic" interface it would be nice to
@@ -564,11 +582,11 @@ impl TryFrom<Numeric> for Lsn {
 
         let vlf_id = u32::try_from(decimal_lsn / 10_i128.pow(15))
             .map_err(|e| format!("Failed to decode vlf_id for lsn {decimal_lsn}: {e:?}"))?;
-        decimal_lsn -= vlf_id as i128 * 10_i128.pow(15);
+        decimal_lsn -= i128::try_from(vlf_id).unwrap() * 10_i128.pow(15);
 
         let block_id = u32::try_from(decimal_lsn / 10_i128.pow(5))
             .map_err(|e| format!("Failed to decode block_id for lsn {decimal_lsn}: {e:?}"))?;
-        decimal_lsn -= block_id as i128 * 10_i128.pow(5);
+        decimal_lsn -= i128::try_from(block_id).unwrap() * 10_i128.pow(5);
 
         let record_id = u16::try_from(decimal_lsn)
             .map_err(|e| format!("Failed to decode record_id for lsn {decimal_lsn}: {e:?}"))?;

src/sql-server-util/src/inspect.rs

@@ -68,26 +68,65 @@ pub async fn increment_lsn(client: &mut Client, lsn: Lsn) -> Result<Lsn, SqlServ
 /// - only contains letters, digits, and underscores
 /// - no resserved words
 /// - 32 char max
-pub async fn lsn_at_savepoint(
+pub async fn create_savepoint(
     txn: &mut Transaction<'_>,
     savepoint_name: &str,
-) -> Result<Lsn, SqlServerError> {
-    let current_lsn_query: &str = "
+) -> Result<(), SqlServerError> {
+    // TODO (maz): make sure savepoint name is safe
+    let _result = txn
+        .client
+        .simple_query(format!("SAVE TRANSACTION [{savepoint_name}]"))
+        .await?;
+    Ok(())
+}
+
+pub async fn get_lsn(txn: &mut Transaction<'_>) -> Result<Lsn, SqlServerError> {
+    static CURRENT_LSN_QUERY: &str = "
 SELECT dt.database_transaction_begin_lsn
 FROM sys.dm_tran_database_transactions AS dt
 JOIN sys.dm_tran_session_transactions AS st
   ON dt.transaction_id = st.transaction_id
 WHERE st.session_id = @@SPID
 ";
-    // TODO (maz): make sure savepoint name is safe
-
-    txn.client
-        .simple_query(format!("SAVE TRANSACTION [{savepoint_name}]"))
-        .await?;
-    let result = txn.client.simple_query(current_lsn_query).await?;
+    let result = txn.client.simple_query(CURRENT_LSN_QUERY).await?;
     parse_numeric_lsn(&result)
 }
 
+pub async fn lock_table(
+    txn: &mut Transaction<'_>,
+    schema: &str,
+    table: &str,
+) -> Result<(), SqlServerError> {
+    // This query probably seems odd, but there is no LOCK command in MS SQL. Locks are specified
+    // in SELECT using the WITH keyword.  This query does not need to return any rows to lock the table,
+    // hence the 1=0, which is something short that always evaluates to false in this universe;
+    let query = format!("SELECT * FROM {schema}.{table} WITH (TABLOCKX) WHERE 1=0;");
+    let _result = txn.client.query(query, &[]).await?;
+    Ok(())
+}
+
+/// Parse an [`Lsn`] in Decimal(25,0) format of the provided [`tiberius::Row`].
+///
+/// Returns an error if the provided slice doesn't have exactly one row.
+fn parse_numeric_lsn(row: &[tiberius::Row]) -> Result<Lsn, SqlServerError> {
+    match row {
+        [r] => {
+            let numeric_lsn = r
+                .try_get::<Numeric, _>(0)?
+                .ok_or_else(|| SqlServerError::NullLsn)?;
+            let lsn = Lsn::try_from(numeric_lsn).map_err(|msg| SqlServerError::InvalidData {
+                column_name: "lsn".to_string(),
+                error: msg,
+            })?;
+            Ok(lsn)
+        }
+        other => Err(SqlServerError::InvalidData {
+            column_name: "lsn".to_string(),
+            error: format!("expected 1 column, got {other:?}"),
+        }),
+    }
+}
+
 /// Parse an [`Lsn`] from the first column of the provided [`tiberius::Row`].
 ///
 /// Returns an error if the provided slice doesn't have exactly one row.
@@ -114,28 +153,6 @@ fn parse_lsn(result: &[tiberius::Row]) -> Result<Lsn, SqlServerError> {
     }
 }
 
-/// Parse an [`Lsn`] in Decimal(25,0) format of the provided [`tiberius::Row`].
-///
-/// Returns an error if the provided slice doesn't have exactly one row.
-fn parse_numeric_lsn(row: &[tiberius::Row]) -> Result<Lsn, SqlServerError> {
-    match row {
-        [r] => {
-            let numeric_lsn = r
-                .try_get::<Numeric, _>(0)?
-                .ok_or_else(|| SqlServerError::NullLsn)?;
-            let lsn = Lsn::try_from(numeric_lsn).map_err(|msg| SqlServerError::InvalidData {
-                column_name: "lsn".to_string(),
-                error: msg,
-            })?;
-            Ok(lsn)
-        }
-        other => Err(SqlServerError::InvalidData {
-            column_name: "lsn".to_string(),
-            error: format!("expected 1 column, got {other:?}"),
-        }),
-    }
-}
-
 /// Queries the specified capture instance and returns all changes from
 /// `[start_lsn, end_lsn)`, ordered by `start_lsn` in an ascending fashion.
 ///

src/sql-server-util/src/lib.rs

@@ -336,6 +336,7 @@ pub type RowStream<'a> =
 pub struct Transaction<'a> {
     client: &'a mut Client,
     closed: bool,
+    nested_xact_names: Vec<String>,
 }
 
 impl<'a> Transaction<'a> {
@@ -352,10 +353,37 @@ impl<'a> Transaction<'a> {
             Ok(Transaction {
                 client,
                 closed: false,
+                nested_xact_names: Default::default(),
             })
         }
     }
 
+    /// Creates a savepoint with a transaction that can be committed or rolled back
+    /// without affecting the out transaction.
+    pub async fn create_savepoint(&mut self, savepoint_name: &str) -> Result<(), SqlServerError> {
+        let stmt = format!("SAVE TRANSACTION [{savepoint_name}]");
+        let _result = self.client.simple_query(stmt).await?;
+        self.nested_xact_names.push(savepoint_name.to_string());
+        Ok(())
+    }
+
+    pub async fn rollback_savepoint(&mut self, savepoint_name: &str) -> Result<(), SqlServerError> {
+        let last_xact_name = self.nested_xact_names.pop();
+        if last_xact_name
+            .as_ref()
+            .is_none_or(|last_xact_name| *last_xact_name != savepoint_name)
+        {
+            panic!(
+                "Attempt to rollback savepoint {savepoint_name} doesn't match last savepoint {:?}",
+                last_xact_name
+            );
+        }
+        let stmt = format!("ROLLBACK TRANSACTION [{savepoint_name}]");
+        let _result = self.client.simple_query(stmt).await?;
+        self.nested_xact_names.push(savepoint_name.to_string());
+        Ok(())
+    }
+
     /// See [`Client::execute`].
     pub async fn execute<'q>(
         &mut self,