[yugabyte#9593] docdb - Move client-side subtransaction state to YBTr…

…ansaction

Summary:
We previously had SubTransactionMetadata living in PgTxnManager in order to allow tracking changes
to this state before a transaction was taken. This revision modifies the behavior of PgTxnManager to
take a transaction and initialize it's txn_ instance as soon as any SubTransactionMetadata state is
interacted with. By doing this, we can move SubTransactionMetadata into YBTransaction.

Test Plan:
ybd --java-test 'org.yb.pgsql.TestPgSavepoints'
ybd --java-test 'org.yb.pgsql.TestPgRegressTransactionSavepoints'

Reviewers: dmitry, sergei

Reviewed By: sergei

Subscribers: mbautin, bogdan, ybase

Differential Revision: https://phabricator.dev.yugabyte.com/D12679

src/yb/client/transaction.cc

@@ -104,6 +104,27 @@ Result<ChildTransactionData> ChildTransactionData::FromPB(const ChildTransaction
 
 YB_DEFINE_ENUM(MetadataState, (kMissing)(kMaybePresent)(kPresent));
 
+YBSubTransaction::YBSubTransaction() {}
+
+void YBSubTransaction::SetActiveSubTransaction(SubTransactionId id) {
+  sub_txn_.subtransaction_id = id;
+  highest_subtransaction_id_ = std::max(highest_subtransaction_id_, id);
+}
+
+CHECKED_STATUS YBSubTransaction::RollbackSubTransaction(SubTransactionId id) {
+  // We should abort the range [id, sub_txn_.highest_subtransaction_id]. It's possible that we
+  // have created and released savepoints, such that there have been writes with a
+  // subtransaction_id greater than sub_txn_.subtransaction_id, and those should be aborted as
+  // well.
+  SCHECK_GE(
+    highest_subtransaction_id_, id,
+    InternalError,
+    "Attempted to rollback to non-existent savepoint.");
+  return sub_txn_.aborted.SetRange(id, highest_subtransaction_id_);
+}
+
+const SubTransactionMetadata& YBSubTransaction::get() { return sub_txn_; }
+
 class YBTransaction::Impl final {
  public:
   Impl(TransactionManager* manager, YBTransaction* transaction)
@@ -295,8 +316,10 @@ class YBTransaction::Impl final {
     }
 
     ops_info->metadata = {
-      metadata_,
-      subtransaction_ ? boost::make_optional(*subtransaction_) : boost::none,
+      .transaction = metadata_,
+      .subtransaction = subtransaction_opt_ != boost::none
+          ? boost::make_optional(subtransaction_opt_->get())
+          : boost::none,
     };
 
     return true;
@@ -624,8 +647,22 @@ class YBTransaction::Impl final {
     RequestStatusTablet(TransactionRpcDeadline());
   }
 
-  void SetSubTransactionMetadata(const SubTransactionMetadata& subtransaction) {
-    subtransaction_ = &subtransaction;
+  void SetActiveSubTransaction(SubTransactionId id) {
+    if (subtransaction_opt_ == boost::none) {
+      subtransaction_opt_ = boost::make_optional(YBSubTransaction());
+    }
+    return subtransaction_opt_->SetActiveSubTransaction(id);
+  }
+
+  CHECKED_STATUS RollbackSubTransaction(SubTransactionId id) {
+    SCHECK(
+        subtransaction_opt_ != boost::none, InternalError,
+        "Attempted to rollback to savepoint before creating any savepoints.");
+    return subtransaction_opt_->RollbackSubTransaction(id);
+  }
+
+  bool HasSubTransactionState() {
+    return subtransaction_opt_ != boost::none;
   }
 
  private:
@@ -710,8 +747,8 @@ class YBTransaction::Impl final {
       return;
     }
 
-    if (subtransaction_) {
-      subtransaction_->aborted.ToPB(state.mutable_aborted()->mutable_set());
+    if (subtransaction_opt_ != boost::none) {
+      subtransaction_opt_->get().aborted.ToPB(state.mutable_aborted()->mutable_set());
     }
 
     manager_->rpcs().RegisterAndStart(
@@ -1120,9 +1157,8 @@ class YBTransaction::Impl final {
   TransactionMetadata metadata_;
   ConsistentReadPoint read_point_;
 
-  // Pointer to SubTransactionMetadata tracking savepoint-related state for the scope of this
-  // transaction.
-  const SubTransactionMetadata* subtransaction_ = nullptr;
+  // Metadata tracking savepoint-related state for the scope of this transaction.
+  boost::optional<YBSubTransaction> subtransaction_opt_ = boost::none;
 
   std::atomic<bool> requested_status_tablet_{false};
   internal::RemoteTabletPtr status_tablet_ GUARDED_BY(mutex_);
@@ -1321,8 +1357,16 @@ YBTransactionPtr YBTransaction::Take(
   return result;
 }
 
-void YBTransaction::SetSubTransactionMetadata(const SubTransactionMetadata& subtransaction) {
-  return impl_->SetSubTransactionMetadata(subtransaction);
+void YBTransaction::SetActiveSubTransaction(SubTransactionId id) {
+  return impl_->SetActiveSubTransaction(id);
+}
+
+Status YBTransaction::RollbackSubTransaction(SubTransactionId id) {
+  return impl_->RollbackSubTransaction(id);
+}
+
+bool YBTransaction::HasSubTransactionState() {
+  return impl_->HasSubTransactionState();
 }
 
 } // namespace client

src/yb/client/transaction.h

@@ -193,13 +193,35 @@ class YBTransaction : public std::enable_shared_from_this<YBTransaction> {
   // between application instances.
   static YBTransactionPtr Take(TransactionManager* manager, const TransactionMetadata& metadata);
 
-  void SetSubTransactionMetadata(const SubTransactionMetadata& subtransaction);
+  void SetActiveSubTransaction(SubTransactionId id);
+
+  CHECKED_STATUS RollbackSubTransaction(SubTransactionId id);
+
+  bool HasSubTransactionState();
 
  private:
   class Impl;
   std::unique_ptr<Impl> impl_;
 };
 
+class YBSubTransaction {
+ public:
+  YBSubTransaction();
+
+  void SetActiveSubTransaction(SubTransactionId id);
+
+  CHECKED_STATUS RollbackSubTransaction(SubTransactionId id);
+
+  const SubTransactionMetadata& get();
+
+ private:
+  SubTransactionMetadata sub_txn_;
+
+  // Tracks the highest observed subtransaction_id. Used during "ROLLBACK TO s" to abort from s to
+  // the highest live subtransaction_id.
+  SubTransactionId highest_subtransaction_id_ = kMinSubTransactionId;
+};
+
 } // namespace client
 } // namespace yb
 

src/yb/common/transaction.h

@@ -216,22 +216,14 @@ using AbortedSubTransactionSet = UnsignedIntSet<SubTransactionId>;
 struct SubTransactionMetadata {
   SubTransactionId subtransaction_id = kMinSubTransactionId;
   AbortedSubTransactionSet aborted;
-  // Tracks the highest observed subtransaction_id. Used during "ROLLBACK TO s" to abort from s to
-  // the highest live subtransaction_id.
-  SubTransactionId highest_subtransaction_id = kMinSubTransactionId;
-
-  // This will lose highest_subtransaction_id, so SubTransactionMetadata::FromPB(stm.ToPB) is
-  // not always equal to stm for `SubTransactionMetadata stm`.
-  // TODO: refactor this to something like
-  // `SubTransactionMetadataWithHighest : public SubTransactionMetadata`.
-  // See https://github.com/yugabyte/yugabyte-db/issues/9593.
+
   void ToPB(SubTransactionMetadataPB* dest) const;
 
   static Result<SubTransactionMetadata> FromPB(
       const SubTransactionMetadataPB& source);
 
   std::string ToString() const {
-    return YB_STRUCT_TO_STRING(subtransaction_id, highest_subtransaction_id, aborted);
+    return YB_STRUCT_TO_STRING(subtransaction_id, aborted);
   }
 
   // Returns true if this is the default state, i.e. default subtransaction_id. This indicates

src/yb/yql/pggate/pg_session.cc

@@ -1258,8 +1258,9 @@ Status PgSession::SetActiveSubTransaction(SubTransactionId id) {
   // already queued operations may incorrectly use this newly modified SubTransactionMetadata when
   // they are eventually sent to DocDB.
   RETURN_NOT_OK(FlushBufferedOperations());
-  pg_txn_manager_->SetActiveSubTransaction(id);
-  return Status::OK();
+  RETURN_NOT_OK(pg_txn_manager_->BeginWriteTransactionIfNecessary(
+      IsReadOnlyOperation::kFalse, IsPessimisticLockRequired::kFalse));
+  return pg_txn_manager_->SetActiveSubTransaction(id);
 }
 
 Status PgSession::RollbackSubTransaction(SubTransactionId id) {

src/yb/yql/pggate/pg_txn_manager.cc

@@ -288,7 +288,6 @@ Status PgTxnManager::BeginWriteTransactionIfNecessary(bool read_only_op,
       DCHECK_EQ(docdb_isolation, IsolationLevel::SERIALIZABLE_ISOLATION);
       RETURN_NOT_OK(txn_->Init(docdb_isolation));
     }
-    txn_->SetSubTransactionMetadata(sub_txn_);
     session_->SetTransaction(txn_);
 
     VLOG_TXN_STATE(2) << "effective isolation level: "
@@ -298,26 +297,22 @@ Status PgTxnManager::BeginWriteTransactionIfNecessary(bool read_only_op,
   return Status::OK();
 }
 
-void PgTxnManager::SetActiveSubTransaction(SubTransactionId id) {
-  sub_txn_.subtransaction_id = id;
-  sub_txn_.highest_subtransaction_id = std::max(sub_txn_.highest_subtransaction_id, id);
+Status PgTxnManager::SetActiveSubTransaction(SubTransactionId id) {
+  SCHECK(
+      txn_, InternalError, "Attempted to set active subtransaction on uninitialized transaciton.");
+  txn_->SetActiveSubTransaction(id);
+  return Status::OK();
 }
 
 Status PgTxnManager::RollbackSubTransaction(SubTransactionId id) {
-  // We should abort the range [id, sub_txn_.highest_subtransaction_id]. It's possible that we have
-  // created and released savepoints, such that there have been writes with a subtransaction_id
-  // greater than sub_txn_.subtransaction_id, and those should be aborted as well.
-  SCHECK_GE(
-    sub_txn_.highest_subtransaction_id, id,
-    InternalError,
-    "Attempted to rollback to non-existent savepoint.");
-  return sub_txn_.aborted.SetRange(id, sub_txn_.highest_subtransaction_id);
+  SCHECK(txn_, InternalError, "Attempted to rollback on uninitialized transaciton.");
+  return txn_->RollbackSubTransaction(id);
 }
 
 Status PgTxnManager::RestartTransaction() {
-  if (!sub_txn_.IsDefaultState()) {
-    return STATUS(IllegalState, "Attempted to restart when session has established savepoints");
-  }
+  SCHECK(
+    !txn_ || !txn_->HasSubTransactionState(), IllegalState,
+    "Attempted to restart when session has established savepoints");
   if (!txn_in_progress_ || !txn_) {
     CHECK_NOTNULL(session_);
     if (!session_->IsRestartRequired()) {
@@ -330,7 +325,6 @@ Status PgTxnManager::RestartTransaction() {
     return STATUS(IllegalState, "Attempted to restart when transaction does not require restart");
   }
   txn_ = VERIFY_RESULT(txn_->CreateRestartedTransaction());
-  txn_->SetSubTransactionMetadata(sub_txn_);
   session_->SetTransaction(txn_);
 
   DCHECK(can_restart_.load(std::memory_order_acquire));
@@ -417,7 +411,6 @@ void PgTxnManager::ResetTxnAndSession() {
   txn_in_progress_ = false;
   session_ = nullptr;
   txn_ = nullptr;
-  sub_txn_ = SubTransactionMetadata();
   can_restart_.store(true, std::memory_order_release);
 }
 

src/yb/yql/pggate/pg_txn_manager.h

@@ -80,7 +80,7 @@ class PgTxnManager : public RefCountedThreadSafe<PgTxnManager> {
   CHECKED_STATUS BeginWriteTransactionIfNecessary(bool read_only_op,
                                                   bool needs_pessimistic_locking = false);
 
-  void SetActiveSubTransaction(SubTransactionId id);
+  CHECKED_STATUS SetActiveSubTransaction(SubTransactionId id);
 
   CHECKED_STATUS RollbackSubTransaction(SubTransactionId id);