perf(log-ingestor): Reduce compression-submission lock contention with READ COMMITTED and smaller update batches (fixes y-scope#2358). (y-scope#2369)

Author: LinZhihao-723

components/log-ingestor/src/ingestion_job_manager/clp_ingestion.rs

@@ -22,7 +22,7 @@ use clp_rust_utils::{
 };
 use const_format::formatcp;
 use non_empty_string::NonEmptyString;
-use sqlx::MySqlPool;
+use sqlx::{Connection, MySqlPool};
 use strum_macros::{AsRefStr, Display, EnumIter, EnumString};
 use tokio::sync::mpsc;
 
@@ -1071,8 +1071,9 @@ impl ClpCompressionState {
     /// * [`anyhow::Error`] if one or more object metadata rows fail to be updated in the DB.
     /// * Forwards [`clp_rust_utils::serde::BrotliMsgpack::serialize`]'s return values on failure.
     /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure.
-    /// * Forwards [`sqlx::Pool::begin`]'s return values on failure.
+    /// * Forwards [`sqlx::Connection::begin`]'s return values on failure.
     /// * Forwards [`sqlx::Transaction::commit`]'s return values on failure.
+    /// * Forwards [`run_read_committed_tx`]'s return values on failure.
     ///
     /// # Panics
     ///
@@ -1093,56 +1094,61 @@ impl ClpCompressionState {
             }
         };
 
-        let mut tx = self.db_pool.begin().await?;
+        // Run the submission transaction at `READ COMMITTED` to avoid the gap/next-key locks that
+        // serialize concurrent submissions under the default `REPEATABLE READ`.
+        run_read_committed_tx(self.db_pool.clone(), async |conn| {
+            let mut tx = Connection::begin(conn).await?;
 
-        // Submit compression job
-        let result = sqlx::query(COMPRESSION_JOB_SUBMISSION_QUERY)
-            .bind(clp_rust_utils::serde::BrotliMsgpack::serialize(&io_config)?)
-            .execute(&mut *tx)
-            .await?;
-        let compression_job_id =
-            CompressionJobId::try_from(result.last_insert_id()).map_err(|_| {
-                anyhow::anyhow!("The retrieved ID overflows: {}", result.last_insert_id())
-            })?;
-
-        // Update compression job ID for ingested objects.
-        // NOTE: We batch the update to avoid hitting the maximum placeholder limit of MySQL. The
-        // batch size is chosen to be 10000, which is conservative enough to avoid hitting the limit
-        // while also minimizing the number of batches for typical use cases. If the number of
-        // placeholders per update changes, we may need to adjust the batch size accordingly.
-        for chunk in object_metadata_ids.chunks(10000) {
-            let mut query_builder = sqlx::QueryBuilder::<sqlx::MySql>::new(formatcp!(
-                r"UPDATE `{table}` ",
-                table = INGESTED_S3_OBJECT_METADATA_TABLE_NAME,
-            ));
-            query_builder
-                .push("SET `compression_job_id` = ")
-                .push_bind(compression_job_id);
-            query_builder
-                .push(", `status` = ")
-                .push_bind(IngestedS3ObjectMetadataStatus::Submitted);
-            query_builder.push(" WHERE `id` IN (");
-            let mut separated_ids = query_builder.separated(", ");
-            for id in chunk {
-                separated_ids.push_bind(id);
-            }
-            query_builder.push(")");
-            query_builder
-                .push(" AND `status` = ")
-                .push_bind(IngestedS3ObjectMetadataStatus::Buffered);
-
-            let result = query_builder.build().execute(&mut *tx).await?;
-            if result.rows_affected()
-                != u64::try_from(chunk.len()).expect("size conversion should always succeed")
-            {
-                return Err(anyhow::anyhow!(
-                    "Failed to update compression job ID for some objects."
+            // Submit compression job
+            let result = sqlx::query(COMPRESSION_JOB_SUBMISSION_QUERY)
+                .bind(clp_rust_utils::serde::BrotliMsgpack::serialize(&io_config)?)
+                .execute(&mut *tx)
+                .await?;
+            let compression_job_id =
+                CompressionJobId::try_from(result.last_insert_id()).map_err(|_| {
+                    anyhow::anyhow!("The retrieved ID overflows: {}", result.last_insert_id())
+                })?;
+
+            // Update compression job ID for ingested objects.
+            // NOTE: We batch the update to avoid hitting the maximum placeholder limit of MySQL.
+            // The batch size of 1000 is conservative enough to avoid the limit while keeping each
+            // UPDATE's lock footprint small under concurrency. If the number of placeholders per
+            // update changes, we may need to adjust the batch size accordingly.
+            for chunk in object_metadata_ids.chunks(1000) {
+                let mut query_builder = sqlx::QueryBuilder::<sqlx::MySql>::new(formatcp!(
+                    r"UPDATE `{table}` ",
+                    table = INGESTED_S3_OBJECT_METADATA_TABLE_NAME,
                 ));
+                query_builder
+                    .push("SET `compression_job_id` = ")
+                    .push_bind(compression_job_id);
+                query_builder
+                    .push(", `status` = ")
+                    .push_bind(IngestedS3ObjectMetadataStatus::Submitted);
+                query_builder.push(" WHERE `id` IN (");
+                let mut separated_ids = query_builder.separated(", ");
+                for id in chunk {
+                    separated_ids.push_bind(id);
+                }
+                query_builder.push(")");
+                query_builder
+                    .push(" AND `status` = ")
+                    .push_bind(IngestedS3ObjectMetadataStatus::Buffered);
+
+                let result = query_builder.build().execute(&mut *tx).await?;
+                if result.rows_affected()
+                    != u64::try_from(chunk.len()).expect("size conversion should always succeed")
+                {
+                    return Err(anyhow::anyhow!(
+                        "Failed to update compression job ID for some objects."
+                    ));
+                }
             }
-        }
 
-        tx.commit().await?;
-        Ok(compression_job_id)
+            tx.commit().await?;
+            Ok(compression_job_id)
+        })
+        .await
     }
 
     /// Waits for the compression job to finish and updates the status of submitted object metadata.
@@ -1509,6 +1515,53 @@ async fn update_job_status(
     Ok(())
 }
 
+/// Runs `tx` on a freshly acquired pooled connection whose next transaction uses the `READ
+/// COMMITTED` isolation level.
+///
+/// A `SET TRANSACTION ISOLATION LEVEL READ COMMITTED` statement is issued on the connection before
+/// `tx` runs. Because it omits `SESSION`/`GLOBAL`, it applies only to the next transaction started
+/// on that connection. `tx` is expected to begin exactly one transaction to consume the setting,
+/// and to commit or roll it back itself.
+///
+/// If `tx` returns an error, the connection is detached from the pool and closed rather than being
+/// released back into it. This ensures a failed attempt can never hand a later, unrelated borrower
+/// a connection still carrying the pending isolation change (or a half-open transaction).
+///
+/// # Type Parameters
+///
+/// * `ReturnType` - The return type of `tx`.
+/// * `TransactionType` - The type of `tx`, which is an async function that takes a mutable
+///   reference to the connection.
+///
+/// # Returns
+///
+/// The value returned by `tx` on success.
+///
+/// # Errors
+///
+/// Returns an error if:
+///
+/// * Forwards [`sqlx::Pool::acquire`]'s return values on failure.
+/// * Forwards [`sqlx::query::Query::execute`]'s return values on failure.
+/// * Forwards `tx`'s return values on failure.
+async fn run_read_committed_tx<ReturnType, TransactionType>(
+    pool: MySqlPool,
+    tx: TransactionType,
+) -> anyhow::Result<ReturnType>
+where
+    for<'connection_lifetime> TransactionType:
+        AsyncFnOnce(&'connection_lifetime mut sqlx::MySqlConnection) -> anyhow::Result<ReturnType>,
+{
+    const SET_READ_COMMITTED: &str = "SET TRANSACTION ISOLATION LEVEL READ COMMITTED";
+    let mut conn = pool.acquire().await?;
+    sqlx::query(SET_READ_COMMITTED).execute(&mut *conn).await?;
+    let result = tx(&mut *conn).await;
+    if result.is_err() {
+        let _ = conn.detach().close().await;
+    }
+    result
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;