[feat][sdk]Support concurrent prewrite

LiuRuoyu01 · chuandew · commit d7b839fa61b6 · 2025-11-28T10:24:03.000Z
diff --git a/include/dingosdk/client.h b/include/dingosdk/client.h
@@ -263,6 +263,8 @@ class Transaction {
   bool IsOnePc() const;
 
   bool IsAsyncCommit() const;
+  
+  bool IsConcurrentPreCommit() const;
 
  private:
   friend class Client;
diff --git a/src/sdk/client.cc b/src/sdk/client.cc
@@ -601,6 +601,8 @@ bool Transaction::IsOnePc() const { return data_->impl->IsOnePc(); }
 
 bool Transaction::IsAsyncCommit() const { return data_->impl->IsAsyncCommit(); }
 
+bool Transaction::IsConcurrentPreCommit() const { return data_->impl->IsConcurrentPreCommit(); }
+
 RegionCreator::RegionCreator(Data* data) : data_(data) {}
 
 RegionCreator::~RegionCreator() { delete data_; }
diff --git a/src/sdk/common/param_config.cc b/src/sdk/common/param_config.cc
@@ -49,14 +49,15 @@ DEFINE_int64(txn_op_max_retry, 20, "txn op max retry times");
 
 DEFINE_int64(txn_prewrite_delay_ms, 500, "txn prewrite delay ms");
 DEFINE_int64(txn_prewrite_max_retry, 300, "txn prewrite max retry");
+DEFINE_bool(enable_txn_concurrent_prewrite, true, "enable txn concurrent prewrite");
 
 DEFINE_int64(raw_kv_delay_ms, 500, "raw kv backoff delay ms");
 DEFINE_int64(raw_kv_max_retry, 10, "raw kv max retry times");
 
 DEFINE_int64(vector_op_delay_ms, 500, "vector task base backoff delay ms");
 DEFINE_int64(vector_op_max_retry, 30, "vector task max retry times");
 
-DEFINE_int64(txn_max_batch_count, 1000, "txn max batch count");
+DEFINE_int64(txn_max_batch_count, 4096, "txn max batch count");
 DEFINE_int64(txn_max_async_commit_count, 256, "txn max async commit count");
 DEFINE_bool(enable_txn_async_commit, true, "enable txn async commit");
 
@@ -65,5 +66,7 @@ DEFINE_bool(log_rpc_time, false, "log rpc time");
 DEFINE_int64(txn_heartbeat_interval_ms, 8000, "txn heartbeat interval time");
 DEFINE_int64(txn_heartbeat_lock_delay_ms, 20000, "txn heartbeat lock delay time");
 
+DEFINE_int64(txn_check_status_interval_ms, 100, "txn check status interval ms");
+
 DEFINE_uint32(stale_period_us, 1000, "stale period us default 1000 us, used for tso provider");
 DEFINE_uint32(tso_batch_size, 256, "tso batch size default 256, used for tso provider");
diff --git a/src/sdk/common/param_config.h b/src/sdk/common/param_config.h
@@ -66,6 +66,7 @@ DECLARE_int64(txn_op_max_retry);
 
 DECLARE_int64(txn_prewrite_delay_ms);
 DECLARE_int64(txn_prewrite_max_retry);
+DECLARE_bool(enable_txn_concurrent_prewrite);
 
 DECLARE_int64(vector_op_delay_ms);
 DECLARE_int64(vector_op_max_retry);
@@ -79,6 +80,8 @@ DECLARE_bool(log_rpc_time);
 DECLARE_int64(txn_heartbeat_interval_ms);
 DECLARE_int64(txn_heartbeat_lock_delay_ms);
 
+DECLARE_int64(txn_check_status_interval_ms);
+
 DECLARE_uint32(stale_period_us);
 DECLARE_uint32(tso_batch_size);
 
diff --git a/src/sdk/transaction/txn_impl.cc b/src/sdk/transaction/txn_impl.cc
@@ -443,9 +443,11 @@ void TxnImpl::ScheduleHeartBeat() {
 void TxnImpl::DoHeartBeat(int64_t start_ts, std::string primary_key) {
   State state = state_.load();
   if (state != kPreCommitted && state != kPreCommitting && state != kCommitting) {
+    DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] heartbeat stop, state({}).", ID(), StateName(state));
     return;
   }
-  if (use_async_commit_.load() && state == kPreCommitted) {
+  if (use_async_commit_.load() && state != kPreCommitting) {
+    DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] heartbeat stop for async commit, state({}).", ID(), StateName(state));
     return;
   }
   std::shared_ptr<TxnHeartbeatTask> heartbeat_task = std::make_shared<TxnHeartbeatTask>(stub_, start_ts, primary_key);
@@ -500,6 +502,8 @@ Status TxnImpl::DoPreCommit() {
 
   use_async_commit_.store(buffer_->MutationsSize() < FLAGS_txn_max_async_commit_count && FLAGS_enable_txn_async_commit);
 
+  use_concurrent_precommit_.store(FLAGS_enable_txn_concurrent_prewrite);
+
   Status status = is_one_pc_.load() ? PreCommit1PC() : PreCommit2PC();
 
   if (!status.ok()) {
@@ -554,6 +558,20 @@ Status TxnImpl::PreCommit1PC() {
 Status TxnImpl::PreCommit2PC() {
   DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] precommit primary key, pk({}).", ID(),
                                   StringToHex(buffer_->GetPrimaryKey()));
+
+  bool is_one_pc = is_one_pc_.load();
+  CHECK(!is_one_pc) << fmt::format("[sdk.txn.{}] precommit 2pc but is_one_pc is true.", ID());
+
+  if (use_concurrent_precommit_.load()) {
+    DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] precommit use concurrent prewrite.", ID());
+    return PreCommit2PCConcurrent();
+  } else {
+    DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] precommit use sequential prewrite.", ID());
+    return PreCommit2PCSequential();
+  }
+}
+
+Status TxnImpl::PreCommit2PCSequential() {
   // primary key map
   std::map<std::string, const TxnMutation*> mutations_map_primary_key;
   mutations_map_primary_key.emplace(
@@ -566,8 +584,10 @@ Status TxnImpl::PreCommit2PC() {
     }
     mutations_map_ordinary_keys.emplace(std::make_pair(key, &mutation));
   }
-  bool is_one_pc = is_one_pc_.load();
-  CHECK(!is_one_pc) << fmt::format("[sdk.txn.{}] precommit 2pc but is_one_pc is true.", ID());
+
+  // precommit primary key
+  DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] precommit primary key.", ID());
+  bool is_one_pc = false;
   bool use_async_commit = use_async_commit_.load();
   uint64_t min_commit_ts = 0;
   TxnPrewriteTask task_primary(stub_, buffer_->GetPrimaryKey(), mutations_map_primary_key, shared_from_this(),
@@ -584,14 +604,9 @@ Status TxnImpl::PreCommit2PC() {
   if (!use_async_commit) {
     use_async_commit_.store(false);
   } else {
-    CHECK(min_commit_ts > 0) << fmt::format("[sdk.txn.{}] min_commit_ts({}) invalid.", ID(), min_commit_ts);
-    min_commit_ts = std::max(min_commit_ts, commit_ts_.load());
-    commit_ts_.store(min_commit_ts);
+    UpdateAsyncCommitTs(min_commit_ts);
   }
 
-  // 2pc need schedule heartbeat to update lock ttl
-  ScheduleHeartBeat();
-
   // precommit ordinary keys
   DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] precommit ordinary keys.", ID());
   min_commit_ts = 0;
@@ -607,14 +622,57 @@ Status TxnImpl::PreCommit2PC() {
   if (!use_async_commit) {
     use_async_commit_.store(false);
   } else {
-    CHECK(min_commit_ts > 0) << fmt::format("[sdk.txn.{}] min_commit_ts({}) invalid.", ID(), min_commit_ts);
-    min_commit_ts = std::max(min_commit_ts, commit_ts_.load());
-    commit_ts_.store(min_commit_ts);
+    UpdateAsyncCommitTs(min_commit_ts);
+  }
+  
+  return Status::OK();
+}
+
+Status TxnImpl::PreCommit2PCConcurrent() {
+  // all keys map
+  std::map<std::string, const TxnMutation*> mutations_map_all_keys;
+  // ordinary keys map for async commit
+  std::map<std::string, const TxnMutation*> mutations_map_ordinary_keys;
+  bool use_async_commit = use_async_commit_.load();
+  for (const auto& [key, mutation] : buffer_->Mutations()) {
+    if (key == buffer_->GetPrimaryKey()) {
+      mutations_map_all_keys.emplace(std::make_pair(key, &mutation));
+      continue;
+    }
+    mutations_map_all_keys.emplace(std::make_pair(key, &mutation));
+    if (use_async_commit) {
+      // for async commit, need to save ordinary keys info
+      mutations_map_ordinary_keys.emplace(std::make_pair(key, &mutation));
+    }
+  }
+
+  bool is_one_pc = false;
+  uint64_t min_commit_ts = 0;
+  TxnPrewriteTask task_primary(stub_, buffer_->GetPrimaryKey(), mutations_map_all_keys, shared_from_this(),
+                               mutations_map_ordinary_keys, is_one_pc, use_async_commit, min_commit_ts);
+
+  Status status = task_primary.Run();
+  if (!status.ok()) {
+    DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] 2pc concurrent precommit keys fail, status({}).", ID(),
+                                      status.ToString());
+    return status;
+  }
+
+  if (!use_async_commit) {
+    use_async_commit_.store(false);
+  } else {
+    UpdateAsyncCommitTs(min_commit_ts);
   }
 
   return Status::OK();
 }
 
+void TxnImpl::UpdateAsyncCommitTs(uint64_t min_commit_ts) {
+  CHECK(min_commit_ts > 0) << fmt::format("[sdk.txn.{}] min_commit_ts({}) invalid.", ID(), min_commit_ts);
+  min_commit_ts = std::max(min_commit_ts, commit_ts_.load());
+  commit_ts_.store(min_commit_ts);
+}
+
 Status TxnImpl::CommitPrimaryKey() {
   std::vector<std::string> keys = {buffer_->GetPrimaryKey()};
   Status status;
diff --git a/src/sdk/transaction/txn_impl.h b/src/sdk/transaction/txn_impl.h
@@ -125,10 +125,14 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
 
   Status Rollback();
 
+  void ScheduleHeartBeat();
+
   bool IsOnePc() const { return is_one_pc_.load(); }
 
   bool IsAsyncCommit() const { return !is_one_pc_.load() && use_async_commit_.load(); }
 
+  bool IsConcurrentPreCommit() const { return use_concurrent_precommit_.load(); }
+
   int64_t GetStartTs() const { return start_ts_.load(); }
   int64_t GetCommitTs() const { return commit_ts_.load(); }
   std::string GetPrimaryKey() const { return buffer_->GetPrimaryKey(); }
@@ -204,6 +208,10 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
   Status DoPreCommit();
   Status PreCommit1PC();
   Status PreCommit2PC();
+  Status PreCommit2PCConcurrent();
+  Status PreCommit2PCSequential();
+
+  void UpdateAsyncCommitTs(uint64_t min_commit_ts);
 
   // txn commit
   Status CommitPrimaryKey();
@@ -219,7 +227,6 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
   Status DoRollback();
 
   void DoHeartBeat(int64_t start_ts, std::string primary_key);
-  void ScheduleHeartBeat();
 
   void CheckStateActive() const;
 
@@ -235,6 +242,7 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
 
   std::atomic<bool> is_one_pc_{false};
   std::atomic<bool> use_async_commit_{false};
+  std::atomic<bool> use_concurrent_precommit_{false};
 
   TxnBufferUPtr buffer_;
 
diff --git a/src/sdk/transaction/txn_lock_resolver.cc b/src/sdk/transaction/txn_lock_resolver.cc
@@ -16,6 +16,7 @@
 
 #include <fmt/format.h>
 #include <fmt/ranges.h>
+#include <unistd.h>
 
 #include <algorithm>
 #include <cstdint>
@@ -27,9 +28,11 @@
 #include "dingosdk/status.h"
 #include "glog/logging.h"
 #include "sdk/client_stub.h"
+#include "sdk/common/param_config.h"
 #include "sdk/transaction/txn_task/txn_check_secondary_locks_task.h"
 #include "sdk/transaction/txn_task/txn_check_status_task.h"
 #include "sdk/transaction/txn_task/txn_resolve_lock_task.h"
+#include "sdk/utils/async_util.h"
 
 namespace dingodb {
 namespace sdk {
@@ -43,18 +46,48 @@ Status TxnLockResolver::ResolveLock(const pb::store::LockInfo& conflict_lock_inf
 
   // check primary key lock status
   TxnStatus txn_status;
-  TxnCheckStatusTask task_check_status(stub_, conflict_lock_info.lock_ts(), conflict_lock_info.primary_lock(), start_ts,
-                                       txn_status, force_sync_commit);
-  Status status = task_check_status.Run();
-  if (!status.ok()) {
-    if (status.IsNotFound()) {
-      DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] not exist txn when check status, status({}) lock({}).", start_ts,
-                                      status.ToString(), conflict_lock_info.ShortDebugString());
-
-      return Status::OK();
+  Status status;
+  int64_t retry_times = 0;
+  int64_t max_retry_times = (FLAGS_txn_heartbeat_lock_delay_ms / FLAGS_txn_check_status_interval_ms) + 1;
+  bool rollback_if_not_exist = false;
+  while (retry_times < max_retry_times) {
+    status = Status();
+    txn_status = TxnStatus();
+    TxnCheckStatusTask task_check_status(stub_, conflict_lock_info.lock_ts(), conflict_lock_info.primary_lock(),
+                                         start_ts, txn_status, force_sync_commit, rollback_if_not_exist);
+    status = task_check_status.Run();
+    if (!status.ok()) {
+      if (status.IsTxnNotFound()) {
+        // check lock ttl expired
+        int64_t current_ts;
+        Status status1 = stub_.GetTsoProvider()->GenPhysicalTs(2, current_ts);
+        if (!status1.ok()) {
+          return status1;
+        }
+        if (conflict_lock_info.lock_ttl() > current_ts) {
+          // lock ttl not expired, wait and retry
+          SleepUs(FLAGS_txn_check_status_interval_ms * 1000);
+        } else {
+          // lock ttl expired, next check will rollback txn if not exist
+          rollback_if_not_exist = true;
+        }
+      } else {
+        return status;
+      }
     } else {
-      return status;
+      break;
     }
+    retry_times++;
+  }
+
+  if (!status.ok()) {
+    DINGO_LOG(WARNING) << fmt::format(
+        "[sdk.txn.{}] check primary lock status fail, retry_times({}), max_retry_times({}), interval_ms({}) "
+        "rollback_if_not_exist({}), "
+        "lock_info({}), status({}).",
+        start_ts, retry_times, max_retry_times, FLAGS_txn_check_status_interval_ms, rollback_if_not_exist,
+        conflict_lock_info.ShortDebugString(), status.ToString());
+    return status;
   }
 
   if (txn_status.primary_lock_info.lock_ts() > 0 && txn_status.primary_lock_info.use_async_commit() &&
diff --git a/src/sdk/transaction/txn_lock_resolver.h b/src/sdk/transaction/txn_lock_resolver.h
@@ -101,12 +101,12 @@ class TxnLockResolver {
   virtual Status ResolveLock(const pb::store::LockInfo& conflict_lock_info, int64_t start_ts,
                              bool force_sync_commit = false);
 
-  virtual Status ResolveLockSecondaryLocks(const pb::store::LockInfo& primary_lock_info, int64_t start_ts,
-                                           const TxnStatus& txn_status, const pb::store::LockInfo& conflict_lock_info);
+ private:
+  Status ResolveLockSecondaryLocks(const pb::store::LockInfo& primary_lock_info, int64_t start_ts,
+                                   const TxnStatus& txn_status, const pb::store::LockInfo& conflict_lock_info);
 
-  virtual Status ResolveNormalLock(const pb::store::LockInfo& lock_info, int64_t start_ts, const TxnStatus& txn_status);
+  Status ResolveNormalLock(const pb::store::LockInfo& lock_info, int64_t start_ts, const TxnStatus& txn_status);
 
- private:
   const ClientStub& stub_;
 };
 
diff --git a/src/sdk/transaction/txn_task/txn_check_status_task.cc b/src/sdk/transaction/txn_task/txn_check_status_task.cc
@@ -56,6 +56,7 @@ void TxnCheckStatusTask::DoAsync() {
   rpc_.MutableRequest()->set_caller_start_ts(start_ts_);
   rpc_.MutableRequest()->set_current_ts(current_ts);
   rpc_.MutableRequest()->set_force_sync_commit(force_sync_commit_);
+  rpc_.MutableRequest()->set_rollback_if_not_exist(rollback_if_not_exist_);
 
   store_rpc_controller_.ResetRegion(region);
   store_rpc_controller_.AsyncCall([this](auto&& s) { TxnCheckStatusRpcCallback(std::forward<decltype(s)>(s)); });
diff --git a/src/sdk/transaction/txn_task/txn_check_status_task.h b/src/sdk/transaction/txn_task/txn_check_status_task.h
@@ -31,13 +31,14 @@ namespace sdk {
 class TxnCheckStatusTask : public TxnTask {
  public:
   TxnCheckStatusTask(const ClientStub& stub, int64_t lock_ts, const std::string& primary_key, int64_t start_ts,
-                     TxnStatus& txn_status, bool force_sync_commit = false)
+                     TxnStatus& txn_status, bool force_sync_commit = false, bool rollback_if_not_exist = false)
       : TxnTask(stub),
         lock_ts_(lock_ts),
         primary_key_(primary_key),
         start_ts_(start_ts),
         txn_status_(txn_status),
         force_sync_commit_(force_sync_commit),
+        rollback_if_not_exist_(rollback_if_not_exist),
         store_rpc_controller_(stub, rpc_) {}
 
   ~TxnCheckStatusTask() override = default;
@@ -57,6 +58,7 @@ class TxnCheckStatusTask : public TxnTask {
   TxnCheckTxnStatusRpc rpc_;
   bool force_sync_commit_{false};
   uint64_t resolved_lock_{0};
+  bool rollback_if_not_exist_{false};
 
   RWLock rw_lock_;
   Status status_;
diff --git a/src/sdk/transaction/txn_task/txn_prewrite_task.cc b/src/sdk/transaction/txn_task/txn_prewrite_task.cc
diff --git a/test/unit_test/sdk/transaction/test_txn_impl.cc b/test/unit_test/sdk/transaction/test_txn_impl.cc
diff --git a/test/unit_test/sdk/transaction/test_txn_lock_resolver.cc b/test/unit_test/sdk/transaction/test_txn_lock_resolver.cc
