[fix][sdk]Fixup get commit ts opportunity

Author: LiuRuoyu01

src/sdk/transaction/txn_impl.cc

@@ -568,52 +568,27 @@ Status TxnImpl::PreCommit2PC() {
   return Status::OK();
 }
 
-Status TxnImpl::ProcessTxnCommitResponse(const TxnCommitResponse* response, bool is_primary) {
-  std::string pk = buffer_->GetPrimaryKey();
-  DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] commit response, pk({}) response({}).", ID(), pk,
-                                  response->ShortDebugString());
-
-  if (!response->has_txn_result()) {
-    return Status::OK();
-  }
-
-  const auto& txn_result = response->txn_result();
-  if (txn_result.has_locked()) {
-    const auto& lock_info = txn_result.locked();
-    DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit lock conflict, is_primary({}) pk({}) response({}).", ID(),
-                                    is_primary, StringToHex(pk), response->ShortDebugString());
-
-  } else if (txn_result.has_txn_not_found()) {
-    DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit not found, is_primary({}) pk({}) response({}).", ID(),
-                                    is_primary, StringToHex(pk), response->ShortDebugString());
-
-  } else if (txn_result.has_write_conflict()) {
-    if (!is_primary) {
-      DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit write conlict, pk({}) response({}).", ID(), StringToHex(pk),
-                                      txn_result.write_conflict().ShortDebugString());
-    }
-    return Status::TxnWriteConflict("txn write conflict");
-
-  } else if (txn_result.has_commit_ts_expired()) {
-    DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] commit ts expired, is_primary({}) pk({}) response({}).", ID(),
-                                      is_primary, StringToHex(pk), txn_result.commit_ts_expired().ShortDebugString());
-    if (is_primary) {
-      int64_t new_commit_ts;
-      auto status = stub_.GetTsoProvider()->GenTs(2, new_commit_ts);
-      commit_ts_.store(new_commit_ts);
-      if (!status.IsOK()) return status;
-      return Status::TxnCommitTsExpired("txn commit ts expired");
-    }
-  }
-
-  return Status::OK();
-}
-
 Status TxnImpl::CommitPrimaryKey() {
   std::vector<std::string> keys = {buffer_->GetPrimaryKey()};
+  Status status;
+  int64_t retry_count = 0;
+  do {
+    TxnCommitTask task(stub_, keys, shared_from_this(), true);
+    status = task.Run();
+    if (status.IsTxnCommitTsExpired()) {
+      int64_t commit_ts;
+      Status s = stub_.GetTsoProvider()->GenTs(2, commit_ts);
+      if (!s.ok()) {
+        DINGO_LOG(ERROR) << fmt::format("[sdk.txn.{}] commit primary key regen ts fail, status({}).", ID(),
+                                        s.ToString());
+        return s;
+      }
+      commit_ts_.store(commit_ts);
+    }
+    retry_count++;
+  } while (status.IsTxnCommitTsExpired() && retry_count < FLAGS_txn_op_max_retry);
 
-  TxnCommitTask task(stub_, keys, shared_from_this(), true);
-  return task.Run();
+  return status;
 }
 
 Status TxnImpl::CommitOrdinaryKey() {

src/sdk/transaction/txn_impl.h

@@ -127,12 +127,11 @@ class TxnImpl : public std::enable_shared_from_this<TxnImpl> {
 
   Status Rollback();
 
-  Status ProcessTxnCommitResponse(const TxnCommitResponse* response, bool is_primary);
-
   bool IsOnePc() const { return is_one_pc_; }
 
   int64_t GetStartTs() const { return start_ts_.load(); }
   int64_t GetCommitTs() const { return commit_ts_.load(); }
+  std::string GetPrimaryKey() const { return buffer_->GetPrimaryKey(); }
   TransactionOptions GetOptions() const { return options_; }
 
   bool CheckFinished() const {

src/sdk/transaction/txn_task/txn_commit_task.cc

@@ -17,11 +17,13 @@
 #include <fmt/format.h>
 #include <glog/logging.h>
 
+#include <cstdint>
 #include <utility>
 
 #include "common/logging.h"
 #include "dingosdk/status.h"
 #include "sdk/common/common.h"
+#include "sdk/common/helper.h"
 #include "sdk/rpc/store_rpc_controller.h"
 #include "sdk/transaction/txn_common.h"
 #include "sdk/utils/rw_lock.h"
@@ -55,7 +57,6 @@ void TxnCommitTask::DoAsync() {
   {
     WriteLockGuard guard(rw_lock_);
     next_batch = next_keys_;
-    need_retry_ = false;
     status_ = Status::OK();
   }
 
@@ -136,37 +137,27 @@ void TxnCommitTask::TxnCommitRpcCallback(const Status& status, TxnCommitRpc* rpc
   DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] rpc: {} request: {} response: {}", txn_impl_->ID(), rpc->Method(),
                                   rpc->Request()->ShortDebugString(), rpc->Response()->ShortDebugString());
   Status s;
-  bool need_retry = false;
   const auto* response = rpc->Response();
   if (!status.ok()) {
     DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] rpc: {} send to region: {} fail: {}", txn_impl_->ID(),
                                       rpc->Method(), rpc->Request()->context().region_id(), status.ToString());
 
     s = status;
   } else {
-    s = txn_impl_->ProcessTxnCommitResponse(response, is_primary_);
+    s = ProcessTxnCommitResponse(response, is_primary_);
     if (!s.ok()) {
       DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] commit fail, region({}) response({}) status({}).",
                                         txn_impl_->ID(), rpc->Request()->context().region_id(),
                                         response->ShortDebugString(), s.ToString());
-      if (s.IsTxnCommitTsExpired() && is_primary_) {
-        need_retry = true;
-        s = Status::OK();
-      }
     }
   }
 
   {
     WriteLockGuard guard(rw_lock_);
     if (s.ok()) {
-      if (!need_retry) {
-        for (const auto& key : rpc->Request()->keys()) {
-          next_keys_.erase(key);
-        }
-      } else {
-        need_retry_ = true;
+      for (const auto& key : rpc->Request()->keys()) {
+        next_keys_.erase(key);
       }
-
     } else {
       if (status_.ok()) {
         // only return first fail status
@@ -181,16 +172,52 @@ void TxnCommitTask::TxnCommitRpcCallback(const Status& status, TxnCommitRpc* rpc
     {
       ReadLockGuard guard(rw_lock_);
       tmp = status_;
-      tmp_need_retry = need_retry_;
-    }
-    if (tmp.ok() && tmp_need_retry) {
-      DoAsyncRetry();
-      return;
     }
     DoAsyncDone(tmp);
   }
 }
 
+Status TxnCommitTask::ProcessTxnCommitResponse(const TxnCommitResponse* response, bool is_primary) {
+  std::string pk = txn_impl_->GetPrimaryKey();
+  int64_t txn_id = txn_impl_->ID();
+  DINGO_LOG(DEBUG) << fmt::format("[sdk.txn.{}] commit response, pk({}) response({}).", txn_id, pk,
+                                  response->ShortDebugString());
+
+  if (!response->has_txn_result()) {
+    return Status::OK();
+  }
+
+  const auto& txn_result = response->txn_result();
+  if (txn_result.has_locked()) {
+    const auto& lock_info = txn_result.locked();
+    DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit lock conflict, is_primary({}) pk({}) response({}).", txn_id,
+                                    is_primary, StringToHex(pk), response->ShortDebugString());
+
+  } else if (txn_result.has_txn_not_found()) {
+    DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit not found, is_primary({}) pk({}) response({}).", txn_id,
+                                    is_primary, StringToHex(pk), response->ShortDebugString());
+
+  } else if (txn_result.has_write_conflict()) {
+    if (!is_primary) {
+      DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit write conlict, pk({}) response({}).", txn_id,
+                                      StringToHex(pk), txn_result.write_conflict().ShortDebugString());
+    }
+    return Status::TxnWriteConflict("txn write conflict");
+
+  } else if (txn_result.has_commit_ts_expired()) {
+    DINGO_LOG(WARNING) << fmt::format("[sdk.txn.{}] commit ts expired, is_primary({}) pk({}) response({}).", txn_id,
+                                      is_primary, StringToHex(pk), txn_result.commit_ts_expired().ShortDebugString());
+    if (is_primary) {
+      return Status::TxnCommitTsExpired("txn commit ts expired");
+    }
+  } else {
+    DINGO_LOG(FATAL) << fmt::format("[sdk.txn.{}] commit unknown txn result, is_primary({}) pk({}) response({}).",
+                                    txn_id, is_primary, StringToHex(pk), response->ShortDebugString());
+  }
+
+  return Status::OK();
+}
+
 }  // namespace sdk
 
 }  // namespace dingodb

src/sdk/transaction/txn_task/txn_commit_task.h

@@ -49,14 +49,15 @@ class TxnCommitTask : public TxnTask {
 
   void TxnCommitRpcCallback(const Status& status, TxnCommitRpc* rpc);
 
+  Status ProcessTxnCommitResponse(const TxnCommitResponse* response, bool is_primary);
+
   const std::vector<std::string> keys_;
   bool is_primary_{false};
   std::shared_ptr<TxnImpl> txn_impl_;
 
   std::vector<StoreRpcController> controllers_;
   std::vector<std::unique_ptr<TxnCommitRpc>> rpcs_;
   std::set<std::string_view> next_keys_;
-  bool need_retry_{false};
 
   std::atomic<int> sub_tasks_count_{0};
   RWLock rw_lock_;

test/unit_test/sdk/transaction/test_txn_impl.cc

@@ -1105,6 +1105,75 @@ TEST_F(SDKTxnImplTest, RollbackSecondKeysFail) {
   EXPECT_EQ(txn->TEST_IsFinishedState(), true);
 }
 
+TEST_F(SDKTxnImplTest, CommitTsExpired) {
+  auto txn = NewTransactionImpl(options);
+
+  EXPECT_EQ(txn->TEST_IsActiveState(), true);
+
+  {
+    txn->Put("a", "a");
+    txn->Delete("a");
+
+    txn->PutIfAbsent("d", "d");
+  }
+  int64_t original_commit_ts;
+
+  EXPECT_CALL(*rpc_client, SendRpc)
+      .WillOnce([&](Rpc& rpc, std::function<void()> cb) {
+        TxnPrewriteRpc* txn_rpc = dynamic_cast<TxnPrewriteRpc*>(&rpc);
+        // precommit
+
+        cb();
+      })
+      .WillOnce([&](Rpc& rpc, std::function<void()> cb) {
+        TxnPrewriteRpc* txn_rpc = dynamic_cast<TxnPrewriteRpc*>(&rpc);
+        // precommit
+
+        cb();
+      })
+      .WillOnce([&](Rpc& rpc, std::function<void()> cb) {
+        TxnCommitRpc* txn_rpc = dynamic_cast<TxnCommitRpc*>(&rpc);
+        // commit primary key
+        CHECK_NOTNULL(txn_rpc);
+        const auto* request = txn_rpc->Request();
+        EXPECT_TRUE(request->has_context());
+        EXPECT_EQ(request->start_ts(), txn->TEST_GetStartTs());
+        EXPECT_EQ(request->commit_ts(), txn->TEST_GetCommitTs());
+        original_commit_ts = request->commit_ts();
+
+        auto* response = txn_rpc->MutableResponse();
+        auto* commit_ts_expired = response->mutable_txn_result()->mutable_commit_ts_expired();
+        commit_ts_expired->mutable_key()->assign(txn->TEST_GetPrimaryKey());
+
+        cb();
+      })
+      .WillRepeatedly([&](Rpc& rpc, std::function<void()> cb) {
+        TxnCommitRpc* txn_rpc = dynamic_cast<TxnCommitRpc*>(&rpc);
+        // commit ordinary key
+        CHECK_NOTNULL(txn_rpc);
+        const auto* request = txn_rpc->Request();
+        EXPECT_TRUE(request->has_context());
+        EXPECT_EQ(request->start_ts(), txn->TEST_GetStartTs());
+        EXPECT_EQ(request->commit_ts(), txn->TEST_GetCommitTs());
+        EXPECT_NE(original_commit_ts, request->commit_ts());
+
+        cb();
+      });
+
+  Status s = txn->PreCommit();
+  EXPECT_TRUE(s.ok());
+  EXPECT_EQ(txn->TEST_IsPreCommittedState(), true);
+
+  s = txn->Commit();
+  EXPECT_TRUE(s.ok());
+  while (true) {
+    if (txn->TEST_IsFinishedState()) {
+      break;
+    }
+    usleep(1000);
+  }
+}
+
 TEST_F(SDKTxnImplTest, LockHeartbeat) {
   auto txn = NewTransactionImpl(options);