fix(tee): fix race condition in batch locking

After scaling zksync-tee-prover to two instances/replicas on Azure for
azure-mainnet2, azure-testnet2, and azure-mainnet2, we started
experiencing duplicated proving for some batches. This was due to a race
condition in batch locking. This PR fixes the issue by adding atomic
batch locking.

core/lib/dal/src/models/storage_tee_proof.rs

@@ -13,6 +13,11 @@ pub struct StorageTeeProof {
     pub attestation: Option<Vec<u8>>,
 }
 
+#[derive(Debug, Clone, sqlx::FromRow)]
+pub struct StorageBatch {
+    pub l1_batch_number: i64,
+}
+
 #[derive(Debug, Clone, sqlx::FromRow)]
 pub struct StorageLockedBatch {
     pub l1_batch_number: i64,

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -13,7 +13,7 @@ use zksync_db_connection::{
 use zksync_types::{tee_types::TeeType, L1BatchNumber};
 
 use crate::{
-    models::storage_tee_proof::{StorageLockedBatch, StorageTeeProof},
+    models::storage_tee_proof::{StorageBatch, StorageLockedBatch, StorageTeeProof},
     Core,
 };
 
@@ -64,72 +64,93 @@ impl TeeProofGenerationDal<'_, '_> {
     ) -> DalResult<Option<LockedBatch>> {
         let processing_timeout = pg_interval_from_duration(processing_timeout);
         let min_batch_number = i64::from(min_batch_number.0);
+        let mut transaction = self.storage.start_transaction().await?;
+        let batch_number = sqlx::query_as!(
+            StorageBatch,
+            r#"
+            SELECT
+                p.l1_batch_number
+            FROM
+                proof_generation_details p
+            LEFT JOIN
+                tee_proof_generation_details tee
+                ON
+                    p.l1_batch_number = tee.l1_batch_number
+                    AND tee.tee_type = $1
+            WHERE
+                (
+                    p.l1_batch_number >= $5
+                    AND p.vm_run_data_blob_url IS NOT NULL
+                    AND p.proof_gen_data_blob_url IS NOT NULL
+                )
+                AND (
+                    tee.l1_batch_number IS NULL
+                    OR (
+                        (tee.status = $2 OR tee.status = $3)
+                        AND tee.prover_taken_at < NOW() - $4::INTERVAL
+                    )
+                )
+            FETCH FIRST ROW ONLY
+            "#,
+            tee_type.to_string(),
+            TeeProofGenerationJobStatus::PickedByProver.to_string(),
+            TeeProofGenerationJobStatus::Failed.to_string(),
+            processing_timeout,
+            min_batch_number
+        )
+        .instrument("lock_batch_for_proving#get_batch_no")
+        .with_arg("tee_type", &tee_type)
+        .with_arg("processing_timeout", &processing_timeout)
+        .with_arg("min_batch_number", &min_batch_number)
+        .fetch_optional(&mut transaction)
+        .await?;
+
+        let batch_number = match batch_number {
+            Some(batch) => batch.l1_batch_number,
+            None => {
+                transaction.commit().await?;
+                return Ok(None);
+            }
+        };
+
         let locked_batch = sqlx::query_as!(
             StorageLockedBatch,
             r#"
-            WITH upsert AS (
-                SELECT
-                    p.l1_batch_number
-                FROM
-                    proof_generation_details p
-                LEFT JOIN
-                    tee_proof_generation_details tee
-                    ON
-                        p.l1_batch_number = tee.l1_batch_number
-                        AND tee.tee_type = $1
-                WHERE
-                    (
-                        p.l1_batch_number >= $5
-                        AND p.vm_run_data_blob_url IS NOT NULL
-                        AND p.proof_gen_data_blob_url IS NOT NULL
-                    )
-                    AND (
-                        tee.l1_batch_number IS NULL
-                        OR (
-                            (tee.status = $2 OR tee.status = $3)
-                            AND tee.prover_taken_at < NOW() - $4::INTERVAL
-                        )
-                    )
-                FETCH FIRST ROW ONLY
-            )
-            
             INSERT INTO
             tee_proof_generation_details (
                 l1_batch_number, tee_type, status, created_at, updated_at, prover_taken_at
             )
-            SELECT
-                l1_batch_number,
+            VALUES
+            (
                 $1,
                 $2,
+                $3,
                 NOW(),
                 NOW(),
                 NOW()
-            FROM
-                upsert
+            )
             ON CONFLICT (l1_batch_number, tee_type) DO
             UPDATE
             SET
-            status = $2,
+            status = $3,
             updated_at = NOW(),
             prover_taken_at = NOW()
             RETURNING
             l1_batch_number,
             created_at
             "#,
+            batch_number,
             tee_type.to_string(),
             TeeProofGenerationJobStatus::PickedByProver.to_string(),
-            TeeProofGenerationJobStatus::Failed.to_string(),
-            processing_timeout,
-            min_batch_number
         )
-        .instrument("lock_batch_for_proving")
+        .instrument("lock_batch_for_proving#insert")
+        .with_arg("batch_number", &batch_number)
         .with_arg("tee_type", &tee_type)
-        .with_arg("processing_timeout", &processing_timeout)
-        .with_arg("l1_batch_number", &min_batch_number)
-        .fetch_optional(self.storage)
+        .fetch_optional(&mut transaction)
         .await?
         .map(Into::into);
 
+        transaction.commit().await?;
         Ok(locked_batch)
     }