refactor: add IdempotentKVTxnSender (#18585)

* refactor: add IdempotentKVTxnSender

- make kv txn of `commit_table_meta` idempotent
- refactor `update_multi_table_meta` and `create_lock_revision`
  by using IdempotentKVTxnSender

* refine comments

* refine comments

Author: dantengsky

src/meta/api/src/schema_api.rs

@@ -105,6 +105,7 @@ use databend_common_proto_conv::FromToProto;
 use crate::errors::TableError;
 use crate::kv_app_error::KVAppError;
 use crate::meta_txn_error::MetaTxnError;
+use crate::util::IdempotentKVTxnSender;
 
 /// SchemaApi defines APIs that provides schema storage, such as database, table.
 #[async_trait::async_trait]
@@ -276,10 +277,10 @@ pub trait SchemaApi: Send + Sync {
         req: UpdateMultiTableMetaReq,
     ) -> Result<UpdateMultiTableMetaResult, KVAppError>;
 
-    async fn update_multi_table_meta_with_txn_id(
+    async fn update_multi_table_meta_with_sender(
         &self,
         req: UpdateMultiTableMetaReq,
-        txn_id: String,
+        kv_txn_sender: &IdempotentKVTxnSender,
     ) -> Result<UpdateMultiTableMetaResult, KVAppError>;
 
     async fn set_table_column_mask_policy(

src/meta/api/src/schema_api_impl.rs

@@ -193,7 +193,6 @@ use databend_common_meta_types::MatchSeqExt;
 use databend_common_meta_types::MetaError;
 use databend_common_meta_types::MetaId;
 use databend_common_meta_types::SeqV;
-use databend_common_meta_types::TxnCondition;
 use databend_common_meta_types::TxnGetRequest;
 use databend_common_meta_types::TxnGetResponse;
 use databend_common_meta_types::TxnOp;
@@ -240,6 +239,8 @@ use crate::util::txn_op_put_pb;
 use crate::util::txn_put_pb;
 use crate::util::txn_replace_exact;
 use crate::util::unknown_database_error;
+use crate::util::IdempotentKVTxnResponse;
+use crate::util::IdempotentKVTxnSender;
 use crate::SchemaApi;
 use crate::DEFAULT_MGET_SIZE;
 
@@ -1783,6 +1784,8 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
 
         let tenant_dbname_tbname = &req.name_ident;
 
+        let txn_sender = IdempotentKVTxnSender::new();
+
         let mut trials = txn_backoff(None, func_name!());
         loop {
             trials.next().unwrap()?.await;
@@ -1917,7 +1920,15 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
                     ],
                 );
 
-                let (succ, _responses) = send_txn(self, txn_req).await?;
+                let txn_response = txn_sender.send_txn(self, txn_req).await?;
+                let succ = match txn_response {
+                    IdempotentKVTxnResponse::Success(_) => true,
+                    IdempotentKVTxnResponse::AlreadyCommitted => {
+                        info!( "Transaction ID {} exists, the corresponding commit_table_meta transaction has been executed successfully", txn_sender.get_txn_id() );
+                        true
+                    }
+                    IdempotentKVTxnResponse::Failed(_) => false,
+                };
 
                 debug!(
                     name :? =(tenant_dbname_tbname),
@@ -2113,20 +2124,15 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
         &self,
         req: UpdateMultiTableMetaReq,
     ) -> Result<UpdateMultiTableMetaResult, KVAppError> {
-        // Generate a random transaction ID for idempotency
-        let txn_id = Uuid::new_v4().to_string();
-        self.update_multi_table_meta_with_txn_id(req, txn_id).await
+        let kv_txn_sender = IdempotentKVTxnSender::new();
+        self.update_multi_table_meta_with_sender(req, &kv_txn_sender)
+            .await
     }
 
-    /// This function is ONLY for testing purposes.
-    /// In production environment, use `update_multi_table_meta` instead.
-    ///
-    /// `retry_times` is used to simulate the retry of the transaction.
-    /// It is only for test.
-    async fn update_multi_table_meta_with_txn_id(
+    async fn update_multi_table_meta_with_sender(
         &self,
         req: UpdateMultiTableMetaReq,
-        txn_id: String,
+        txn_sender: &IdempotentKVTxnSender,
     ) -> Result<UpdateMultiTableMetaResult, KVAppError> {
         let UpdateMultiTableMetaReq {
             mut update_table_metas,
@@ -2303,47 +2309,27 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
                 .push(build_upsert_table_deduplicated_label(deduplicated_label));
         }
 
-        // Add transaction ID to the transaction with 5-minute expiration
-        let txn_id_key = format!("_txn_id/{}", txn_id);
-
-        // Add condition to check that transaction ID does not exist (empty)
-        txn.condition
-            .push(TxnCondition::eq_seq(txn_id_key.clone(), 0));
-
-        txn.if_then.push(TxnOp::put_with_ttl(
-            txn_id_key.clone(),
-            vec![],
-            Some(Duration::from_secs(300)),
-        ));
-
-        // Add get operation to check if transaction ID exists in else branch
-        // NOTE: Orders of operations matter, please keep this as the last operation of else branch
-        txn.else_then.push(TxnOp::get(txn_id_key.clone()));
+        let txn_response = txn_sender.send_txn(self, txn).await?;
 
-        let (succ, responses) = send_txn(self, txn).await?;
-
-        if succ {
-            return Ok(Ok(UpdateTableMetaReply {}));
-        }
-
-        // Check if transaction ID exists in else branch response (idempotency check)
-        //
-        // Please note that this is a best effort check: the "idempotency check" is only guaranteed
-        // roughly within 300 secs, please DO NOT use it for any safety properties.
-        if let Some(Response::Get(get_resp)) = responses.last().and_then(|r| r.response.as_ref()) {
-            // Defensive check, make sure the get(tx_id_key) is the last operation
-            assert_eq!(get_resp.key, txn_id_key, "Transaction ID key mismatch");
-            if get_resp.value.is_some() {
+        let else_branch_op_responses = match txn_response {
+            IdempotentKVTxnResponse::Success(_) => {
+                return Ok(Ok(UpdateTableMetaReply {}));
+            }
+            IdempotentKVTxnResponse::AlreadyCommitted => {
                 info!(
-                    "Transaction ID {} exists, the corresponding transaction has been executed successfully",
-                    txn_id_key
+                    "Transaction ID {} exists, the corresponding update_multi_table_meta transaction has been executed successfully",
+                    txn_sender.get_txn_id()
                 );
                 return Ok(Ok(UpdateTableMetaReply {}));
             }
-        }
+            IdempotentKVTxnResponse::Failed(op_responses) => op_responses,
+        };
 
         let mut mismatched_tbs = vec![];
-        for (resp, req) in responses.iter().zip(update_table_metas.iter()) {
+        for (resp, req) in else_branch_op_responses
+            .iter()
+            .zip(update_table_metas.iter())
+        {
             let Some(Response::Get(get_resp)) = &resp.response else {
                 unreachable!(
                     "internal error: expect some TxnGetResponseGet, but got {:?}",
@@ -2957,11 +2943,8 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
         let lock_key = &req.lock_key;
         let id_generator = IdGenerator::table_lock_id();
 
-        // ensure idempotent, reference from update_multi_table_meta.
-        let txn_id = Uuid::new_v4().to_string();
-        let txn_id_key = format!("_txn_id/{}", txn_id);
-
         let mut trials = txn_backoff(None, ctx);
+        let txn_sender = IdempotentKVTxnSender::with_ttl(req.ttl);
         loop {
             trials.next().unwrap()?.await;
 
@@ -2982,34 +2965,27 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
                 txn_cond_seq(&id_generator, Eq, current_rev),
                 // assumes lock are absent.
                 txn_cond_seq(&key, Eq, 0),
-                TxnCondition::eq_seq(txn_id_key.clone(), 0),
             ];
             let if_then = vec![
                 txn_op_put(&id_generator, b"".to_vec()),
                 txn_op_put_pb(&key, &lock_meta, Some(req.ttl))?,
-                TxnOp::put_with_ttl(txn_id_key.clone(), vec![], Some(req.ttl)),
             ];
-            let else_then = vec![TxnOp::get(txn_id_key.clone())];
-            let txn_req = TxnRequest::new(condition, if_then).with_else(else_then);
-            let (succ, responses) = send_txn(self, txn_req).await?;
-
-            if succ {
-                return Ok(CreateLockRevReply { revision });
-            }
-
-            // Check if transaction ID exists in else branch response (idempotency check).
-            if let Some(Response::Get(get_resp)) =
-                responses.last().and_then(|r| r.response.as_ref())
-            {
-                // Defensive check, make sure the get(tx_id_key) is the last operation
-                assert_eq!(get_resp.key, txn_id_key, "Transaction ID key mismatch");
-                if get_resp.value.is_some() {
+            let txn_req = TxnRequest::new(condition, if_then);
+            let txn_response = txn_sender.send_txn(self, txn_req).await?;
+            match txn_response {
+                IdempotentKVTxnResponse::Success(_) => {
+                    return Ok(CreateLockRevReply { revision });
+                }
+                IdempotentKVTxnResponse::AlreadyCommitted => {
                     info!(
                         "Transaction ID {} exists, the lock revision has been created successfully",
-                        txn_id_key
+                        txn_sender.get_txn_id()
                     );
                     return Ok(CreateLockRevReply { revision });
                 }
+                _ => {
+                    // continue looping
+                }
             }
         }
     }

src/meta/api/src/schema_api_test_suite.rs

@@ -23,7 +23,6 @@ use std::vec;
 use chrono::DateTime;
 use chrono::Duration;
 use chrono::Utc;
-use databend_common_base::base::uuid::Uuid;
 use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
 use databend_common_exception::ErrorCode;
@@ -155,6 +154,7 @@ use crate::kv_pb_api::UpsertPB;
 use crate::serialize_struct;
 use crate::testing::get_kv_data;
 use crate::testing::get_kv_u64_data;
+use crate::util::IdempotentKVTxnSender;
 use crate::DatamaskApi;
 use crate::RowAccessPolicyApi;
 use crate::SchemaApi;
@@ -2485,6 +2485,7 @@ impl SchemaApiTestSuite {
             }
 
             info!("--- update table meta: simulate kv txn retried after commit");
+            let txn_sender = IdempotentKVTxnSender::new();
             {
                 // 1. Update test table, bump table version, expect success -- just the normal case
                 let table = util.get_table().await.unwrap();
@@ -2498,14 +2499,13 @@ impl SchemaApiTestSuite {
                     new_table_meta: new_table_meta.clone(),
                     base_snapshot_location: None,
                 };
-                let uuid = Uuid::new_v4();
                 let res = mt
-                    .update_multi_table_meta_with_txn_id(
+                    .update_multi_table_meta_with_sender(
                         UpdateMultiTableMetaReq {
                             update_table_metas: vec![(req.clone(), table.as_ref().clone())],
                             ..Default::default()
                         },
-                        uuid.to_string(),
+                        &txn_sender,
                     )
                     .await?;
 
@@ -2530,13 +2530,13 @@ impl SchemaApiTestSuite {
                 // Expects no KV api level error, and no app level error.
                 // For the convenience of reviewing, using explict type signature
                 let _r: databend_common_meta_app::schema::UpdateTableMetaReply = mt
-                    .update_multi_table_meta_with_txn_id(
+                    .update_multi_table_meta_with_sender(
                         UpdateMultiTableMetaReq {
                             update_table_metas: vec![(req, table.as_ref().clone())],
                             ..Default::default()
                         },
-                        // USING THE SAME UUID as kv txn_id
-                        uuid.to_string(),
+                        // USING THE SAME KVTxnSender
+                        &txn_sender,
                     )
                     .await?
                     .unwrap();