[feat][sdk]Support txn async rollback ordinary keys

Author: LiuRuoyu01

python/src/client_bindings.cc

@@ -247,7 +247,6 @@ void DefineClientBindings(pybind11::module& m) {
              Status status = transaction.Scan(start_key, end_key, limit, kvs);
              return std::make_tuple(status, kvs);
            })
-      .def("PreCommit", &Transaction::PreCommit)
       .def("Commit", &Transaction::Commit)
       .def("Rollback", &Transaction::Rollback);
 

src/sdk/client_stub.cc

@@ -70,6 +70,9 @@ Status ClientStub::Open(const std::vector<EndPoint>& endpoints) {
   actuator_ = std::make_shared<ThreadPoolActuator>();
   actuator_->Start(FLAGS_actuator_thread_num);
 
+  txn_actuator_ = std::make_shared<ThreadPoolActuator>();
+  txn_actuator_->Start(FLAGS_txn_actuator_thread_num);
+
   vector_index_cache_ = std::make_shared<VectorIndexCache>(*this);
 
   document_index_cache_ = std::make_shared<DocumentIndexCache>(*this);
@@ -91,6 +94,9 @@ void ClientStub::Stop() {
   if (actuator_) {
     actuator_->Stop();
   }
+  if (txn_actuator_) {
+    txn_actuator_->Stop();
+  }
 }
 
 }  // namespace sdk

src/sdk/client_stub.h

@@ -95,6 +95,11 @@ class ClientStub {
     return actuator_;
   }
 
+  virtual std::shared_ptr<Actuator> GetTxnActuator() const {
+    DCHECK_NOTNULL(txn_actuator_.get());
+    return txn_actuator_;
+  }
+
   virtual std::shared_ptr<VectorIndexCache> GetVectorIndexCache() const {
     DCHECK_NOTNULL(vector_index_cache_.get());
     return vector_index_cache_;
@@ -132,6 +137,7 @@ class ClientStub {
   std::shared_ptr<AdminTool> admin_tool_;
   std::shared_ptr<TxnLockResolver> txn_lock_resolver_;
   std::shared_ptr<Actuator> actuator_;
+  std::shared_ptr<Actuator> txn_actuator_;
   std::shared_ptr<VectorIndexCache> vector_index_cache_;
   std::shared_ptr<DocumentIndexCache> document_index_cache_;
   std::shared_ptr<AutoIncrementerManager> auto_increment_manager_;

src/sdk/common/param_config.cc

@@ -19,6 +19,7 @@
 
 // sdk config
 DEFINE_int64(actuator_thread_num, 8, "actuator thread num");
+DEFINE_int64(txn_actuator_thread_num, 16, "txn actuator thread num");
 
 // coordinator config
 DEFINE_int64(coordinator_interaction_delay_ms, 500, "coordinator interaction delay ms");
@@ -56,6 +57,7 @@ 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_async_commit_count, 256, "txn max async commit count");
 
 DEFINE_bool(log_rpc_time, false, "log rpc time");
 

src/sdk/common/param_config.h

@@ -26,6 +26,7 @@
 // sdk config
 const int64_t kSdkVlogLevel = 60;
 DECLARE_int64(actuator_thread_num);
+DECLARE_int64(txn_actuator_thread_num);
 
 // coordinator config
 const int64_t kPrefetchRegionCount = 3;
@@ -70,6 +71,7 @@ DECLARE_int64(vector_op_delay_ms);
 DECLARE_int64(vector_op_max_retry);
 
 DECLARE_int64(txn_max_batch_count);
+DECLARE_int64(txn_max_async_commit_count);
 DECLARE_bool(log_rpc_time);
 
 DECLARE_int64(txn_heartbeat_interval_ms);

src/sdk/transaction/txn_impl.cc

@@ -433,7 +433,7 @@ Status TxnImpl::PreWriteAndCommit() {
 }
 
 void TxnImpl::ScheduleHeartBeat() {
-  stub_.GetActuator()->Schedule(
+  stub_.GetTxnActuator()->Schedule(
       [shared_this = shared_from_this(), start_ts = start_ts_.load(), primary_key = buffer_->GetPrimaryKey()] {
         shared_this->DoHeartBeat(start_ts, primary_key);
       },
@@ -608,9 +608,11 @@ Status TxnImpl::CommitOrdinaryKey() {
     }
   }
   // async commit ordinary keys
-  stub_.GetActuator()->Schedule([shared_this = shared_from_this(),
-                                 ordinary_keys = std::move(keys)] { shared_this->DoCommitOrdinaryKey(ordinary_keys); },
-                                0);
+  stub_.GetTxnActuator()->Schedule(
+      [shared_this = shared_from_this(), ordinary_keys = std::move(keys)] {
+        shared_this->DoCommitOrdinaryKey(ordinary_keys);
+      },
+      0);
   return Status::OK();
 }
 
@@ -683,15 +685,28 @@ Status TxnImpl::RollbackPrimaryKey() {
   std::vector<std::string> keys;
   std::string pk = buffer_->GetPrimaryKey();
   keys.push_back(pk);
-  if (is_one_pc_) {
+  bool is_one_pc = is_one_pc_.load();
+  if (is_one_pc) {
     for (const auto& [key, _] : buffer_->Mutations()) {
       if (key != pk) {
         keys.push_back(key);
       }
     }
   }
-  TxnBatchRollbackTask task(stub_, keys, shared_from_this());
-  return task.Run();
+  TxnBatchRollbackTask task(stub_, std::move(keys), shared_from_this(), is_one_pc);
+  Status status = task.Run();
+  if (!status.ok()) {
+    DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] 1pc rollback key fail, status({}).", ID(), status.ToString());
+    if (!is_one_pc) {
+      // retry rollback primary key
+      DINGO_LOG(INFO) << fmt::format("[sdk.txn.{}] retry rollback primary key.", ID());
+      is_one_pc_.store(false);
+      std::vector<std::string> primary_key = {pk};
+      TxnBatchRollbackTask task(stub_, std::move(primary_key), shared_from_this(), is_one_pc);
+      status = task.Run();
+    }
+  }
+  return status;
 }
 
 Status TxnImpl::RollbackOrdinaryKey() {
@@ -704,8 +719,30 @@ Status TxnImpl::RollbackOrdinaryKey() {
     }
     keys.push_back(key);
   }
-  TxnBatchRollbackTask task(stub_, keys, shared_from_this());
-  return task.Run();
+  // async rollback ordinary keys
+  stub_.GetTxnActuator()->Schedule(
+      [shared_this = shared_from_this(), ordinary_keys = std::move(keys)] {
+        shared_this->DoRollbackOrdinaryKey(ordinary_keys);
+      },
+      0);
+  return Status::OK();
+}
+
+void TxnImpl::DoRollbackOrdinaryKey(std::vector<std::string> keys) {
+  CHECK(state_.load() == kRollbacked) << "state is not rollbacked, state:" << StateName(state_.load());
+  bool is_one_pc = is_one_pc_.load();
+  CHECK(!is_one_pc) << fmt::format("[sdk.txn.{}] 1pc txn should not rollback ordinary keys.", ID());
+  std::shared_ptr<TxnBatchRollbackTask> txn_rollback_task =
+      std::make_shared<TxnBatchRollbackTask>(stub_, keys, shared_from_this(), is_one_pc);
+
+  DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}]rollback ordinary keys, size({}).", ID(), keys.size());
+
+  auto status = txn_rollback_task->Run();
+  if (!status.ok()) {
+    DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] rollback ordinary keys fail. status({}).", ID(), status.ToString());
+  }
+  state_.store(kFinshed);
+  Cleanup();
 }
 
 Status TxnImpl::DoRollback() {
@@ -742,8 +779,6 @@ Status TxnImpl::DoRollback() {
   if (!status.IsOK()) {
     DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] rollback ordinary keys fail, status({}).", ID(), status.ToString());
   }
-  state_.store(kFinshed);
-  Cleanup();
 
   return Status::OK();
 }

src/sdk/transaction/txn_impl.h

@@ -212,6 +212,7 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
   // txn rollback
   Status RollbackPrimaryKey();
   Status RollbackOrdinaryKey();
+  void DoRollbackOrdinaryKey(std::vector<std::string> keys);
   Status DoRollback();
 
   void DoHeartBeat(int64_t start_ts, std::string primary_key);

src/sdk/transaction/txn_task/txn_batch_rollback_task.cc

@@ -80,6 +80,7 @@ void TxnBatchRollbackTask::DoAsync() {
       std::string msg = fmt::format("[sdk.txn.{}] one pc rollback only support one region, but got {} regions",
                                     txn_impl_->ID(), region_id_to_region.size());
       DINGO_LOG(ERROR) << msg;
+      is_one_pc_ = false;
       DoAsyncDone(Status::InvalidArgument(msg));
       return;
     }

src/sdk/transaction/txn_task/txn_batch_rollback_task.h

@@ -33,8 +33,8 @@ namespace sdk {
 
 class TxnBatchRollbackTask : public TxnTask {
  public:
-  TxnBatchRollbackTask(const ClientStub& stub, const std::vector<std::string>& keys, std::shared_ptr<TxnImpl> txn_impl,
-                       bool is_one_pc = false)
+  TxnBatchRollbackTask(const ClientStub& stub, const std::vector<std::string> keys, std::shared_ptr<TxnImpl> txn_impl,
+                       bool& is_one_pc)
       : TxnTask(stub), keys_(keys), txn_impl_(txn_impl), is_one_pc_(is_one_pc) {}
 
   ~TxnBatchRollbackTask() override = default;
@@ -48,15 +48,15 @@ class TxnBatchRollbackTask : public TxnTask {
 
   void TxnBatchRollbackRpcCallback(const Status& status, TxnBatchRollbackRpc* rpc);
 
-  const std::vector<std::string>& keys_;
+  const std::vector<std::string> keys_;
   std::shared_ptr<TxnImpl> txn_impl_;
-  bool is_one_pc_;
+  bool& is_one_pc_;
 
   std::vector<StoreRpcController> controllers_;
   std::vector<std::unique_ptr<TxnBatchRollbackRpc>> rpcs_;
   std::set<std::string_view> next_keys_;
   std::atomic<int> sub_tasks_count_{0};
-  
+
   RWLock rw_lock_;
   Status status_;
 };

test/unit_test/sdk/mock_client_stub.h

@@ -35,6 +35,7 @@ class MockClientStub final : public ClientStub {
   MOCK_METHOD(std::shared_ptr<AdminTool>, GetAdminTool, (), (const, override));
   MOCK_METHOD(std::shared_ptr<TxnLockResolver>, GetTxnLockResolver, (), (const, override));
   MOCK_METHOD(std::shared_ptr<Actuator>, GetActuator, (), (const, override));
+  MOCK_METHOD(std::shared_ptr<Actuator>, GetTxnActuator, (), (const, override));
   MOCK_METHOD(std::shared_ptr<VectorIndexCache>, GetVectorIndexCache, (), (const, override));
   MOCK_METHOD(std::shared_ptr<AutoIncrementerManager>, GetAutoIncrementerManager, (), (const, override));
   MOCK_METHOD(std::shared_ptr<TsoProvider>, GetTsoProvider, (), (const, override));