impl(spanner): support multiplexed session transaction priority (#15361)

Author: scotthart

google/cloud/spanner/internal/transaction_impl.cc

@@ -21,6 +21,45 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 TransactionImpl::~TransactionImpl() = default;
 
+TransactionImpl::TransactionImpl(
+    google::spanner::v1::TransactionSelector selector, bool route_to_leader,
+    std::string tag)
+    : TransactionImpl(/*session=*/{}, std::move(selector), route_to_leader,
+                      std::move(tag), absl::nullopt) {}
+
+TransactionImpl::TransactionImpl(
+    TransactionImpl const& impl,
+    google::spanner::v1::TransactionSelector selector, bool route_to_leader,
+    std::string tag)
+    : TransactionImpl(impl.session_, std::move(selector), route_to_leader,
+                      std::move(tag),
+                      (impl.session_ && impl.session_->is_multiplexed() &&
+                       impl.selector_->has_id())
+                          ? absl::optional<std::string>(impl.selector_->id())
+                          : absl::nullopt) {}
+
+TransactionImpl::TransactionImpl(
+    SessionHolder session, google::spanner::v1::TransactionSelector selector,
+    bool route_to_leader, std::string tag,
+    absl::optional<std::string> multiplexed_session_previous_transaction_id)
+    : session_(std::move(session)),
+      selector_(std::move(selector)),
+      route_to_leader_(route_to_leader),
+      tag_(std::move(tag)),
+      seqno_(0) {
+  state_ = selector_->has_begin() ? State::kBegin : State::kDone;
+  // If we're attempting to retry an aborted ReadWrite transaction on a
+  // multiplexed session, then propagate the aborted transaction id.
+  if (session_ && session_->is_multiplexed() && selector_.ok() &&
+      selector_->has_begin() && selector_->begin().has_read_write() &&
+      multiplexed_session_previous_transaction_id.has_value()) {
+    selector_->mutable_begin()
+        ->mutable_read_write()
+        ->set_multiplexed_session_previous_transaction_id(
+            *multiplexed_session_previous_transaction_id);
+  }
+}
+
 void TransactionImpl::UpdatePrecommitToken(
     std::unique_lock<std::mutex> const&,
     absl::optional<google::spanner::v1::MultiplexedSessionPrecommitToken>

google/cloud/spanner/internal/transaction_impl.h

@@ -52,26 +52,16 @@ using VisitInvokeResult = ::google::cloud::internal::invoke_result_t<
 class TransactionImpl {
  public:
   TransactionImpl(google::spanner::v1::TransactionSelector selector,
-                  bool route_to_leader, std::string tag)
-      : TransactionImpl(/*session=*/{}, std::move(selector), route_to_leader,
-                        std::move(tag)) {}
+                  bool route_to_leader, std::string tag);
 
   TransactionImpl(TransactionImpl const& impl,
                   google::spanner::v1::TransactionSelector selector,
-                  bool route_to_leader, std::string tag)
-      : TransactionImpl(impl.session_, std::move(selector), route_to_leader,
-                        std::move(tag)) {}
+                  bool route_to_leader, std::string tag);
 
-  TransactionImpl(SessionHolder session,
-                  google::spanner::v1::TransactionSelector selector,
-                  bool route_to_leader, std::string tag)
-      : session_(std::move(session)),
-        selector_(std::move(selector)),
-        route_to_leader_(route_to_leader),
-        tag_(std::move(tag)),
-        seqno_(0) {
-    state_ = selector_->has_begin() ? State::kBegin : State::kDone;
-  }
+  TransactionImpl(
+      SessionHolder session, google::spanner::v1::TransactionSelector selector,
+      bool route_to_leader, std::string tag,
+      absl::optional<std::string> multiplexed_session_previous_transaction_id);
 
   ~TransactionImpl();
 

google/cloud/spanner/transaction.cc

@@ -144,7 +144,8 @@ Transaction::Transaction(std::string session_id, std::string transaction_id,
   selector.set_id(std::move(transaction_id));
   impl_ = std::make_shared<spanner_internal::TransactionImpl>(
       spanner_internal::MakeDissociatedSessionHolder(std::move(session_id)),
-      std::move(selector), route_to_leader, std::move(transaction_tag));
+      std::move(selector), route_to_leader, std::move(transaction_tag),
+      absl::nullopt);
 }
 
 Transaction::~Transaction() = default;

google/cloud/spanner/transaction_test.cc

@@ -126,6 +126,43 @@ TEST(Transaction, SessionAffinity) {
         EXPECT_EQ(a_session, session);  // session affinity
         EXPECT_TRUE(s->has_begin());    // but a new transaction
         EXPECT_EQ(ctx.tag, "app=cart,env=dev");
+        EXPECT_THAT(s->begin()
+                        .read_write()
+                        .multiplexed_session_previous_transaction_id(),
+                    IsEmpty());
+        return 0;
+      });
+}
+
+TEST(Transaction, MultiplexedPreviousTransactionId) {
+  std::string aborted_txn_id = "aborted-txn-id";
+  auto mux_session = spanner_internal::MakeMultiplexedSessionHolder(
+      "multiplexed", std::make_shared<spanner_internal::Session::Clock>());
+  auto opts = Transaction::ReadWriteOptions().WithTag("app=cart,env=dev");
+  Transaction aborted_txn = MakeReadWriteTransaction(opts);
+  spanner_internal::Visit(
+      aborted_txn, [&](spanner_internal::SessionHolder& session,
+                       StatusOr<google::spanner::v1::TransactionSelector>& s,
+                       spanner_internal::TransactionContext const& ctx) {
+        EXPECT_FALSE(session);
+        EXPECT_TRUE(s->has_begin());
+        session = mux_session;
+        s->set_id(aborted_txn_id);
+        EXPECT_EQ(ctx.tag, "app=cart,env=dev");
+        return 0;
+      });
+  Transaction retry_txn = MakeReadWriteTransaction(aborted_txn, opts);
+  spanner_internal::Visit(
+      retry_txn, [&](spanner_internal::SessionHolder& session,
+                     StatusOr<google::spanner::v1::TransactionSelector>& s,
+                     spanner_internal::TransactionContext const& ctx) {
+        EXPECT_EQ(mux_session, session);
+        EXPECT_TRUE(s->has_begin());
+        EXPECT_EQ(s->begin()
+                      .read_write()
+                      .multiplexed_session_previous_transaction_id(),
+                  aborted_txn_id);
+        EXPECT_EQ(ctx.tag, "app=cart,env=dev");
         return 0;
       });
 }