Backport to branch(3) : Decouple CoordinatorCommitHandler from TransactionContext and WriteSetEncoder (#3675)

Co-authored-by: Toshihiro Suzuki <brfrn169@gmail.com>

Author: feeblefakie

core/src/main/java/com/scalar/db/transaction/consensuscommit/CommitHandler.java

@@ -13,6 +13,7 @@
 import com.scalar.db.exception.transaction.UnknownTransactionStatusException;
 import com.scalar.db.exception.transaction.ValidationConflictException;
 import com.scalar.db.exception.transaction.ValidationException;
+import com.scalar.db.transaction.consensuscommit.proto.v1.WriteSet;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.Optional;
 import java.util.concurrent.Future;
@@ -39,19 +40,22 @@
  * (ConsensusCommit / TwoPhaseConsensusCommit / ConsensusCommitManager) can depend on the
  * orchestrator alone.
  *
- * <p>Public methods like {@code commitState}, {@code abortState}, {@code prepareRecords}, {@code
- * commitRecords}, {@code rollbackRecords}, {@code commitStateWithoutWriteSet}, and {@code
- * abortStateWithoutWriteSet} are exposed on this class as thin pass-throughs to the specialized
- * handlers so direct callers (notably {@link TwoPhaseConsensusCommit}) can drive individual commit
- * phases without depending on the handlers directly.
+ * <p>Several public methods are exposed on this class so direct callers (notably {@link
+ * TwoPhaseConsensusCommit}) can drive individual commit phases without depending on the specialized
+ * handlers directly. {@code prepareRecords}, {@code commitRecords}, {@code rollbackRecords}, {@code
+ * commitStateWithoutWriteSet}, and {@code abortStateWithoutWriteSet} are thin pass-throughs. {@code
+ * commitState} and {@code abortState} are not: they encode the transaction's write set via the
+ * orchestrator-owned {@link WriteSetEncoder} before delegating to the Coordinator-side handler.
  */
 @ThreadSafe
 public class CommitHandler {
   private static final Logger logger = LoggerFactory.getLogger(CommitHandler.class);
 
   private final CoordinatorCommitHandler coordinatorCommitHandler;
   private final ParticipantCommitHandler participantCommitHandler;
+  protected final WriteSetEncoder writeSetEncoder;
   protected final boolean coordinatorWriteOmissionOnReadOnlyEnabled;
+  protected final boolean coordinatorWriteSetLoggingEnabled;
 
   @LazyInit @Nullable private BeforePreparationHook beforePreparationHook;
 
@@ -66,11 +70,9 @@ public CommitHandler(
       boolean coordinatorWriteSetLoggingEnabled,
       boolean onePhaseCommitEnabled) {
     this.coordinatorWriteOmissionOnReadOnlyEnabled = coordinatorWriteOmissionOnReadOnlyEnabled;
-    this.coordinatorCommitHandler =
-        new CoordinatorCommitHandler(
-            coordinator,
-            new WriteSetEncoder(tableMetadataManager),
-            coordinatorWriteSetLoggingEnabled);
+    this.coordinatorWriteSetLoggingEnabled = coordinatorWriteSetLoggingEnabled;
+    this.writeSetEncoder = new WriteSetEncoder(tableMetadataManager);
+    this.coordinatorCommitHandler = new CoordinatorCommitHandler(coordinator);
     this.participantCommitHandler =
         new ParticipantCommitHandler(
             storage,
@@ -87,9 +89,13 @@ public CommitHandler(
   @SuppressFBWarnings("EI_EXPOSE_REP2")
   protected CommitHandler(
       boolean coordinatorWriteOmissionOnReadOnlyEnabled,
+      boolean coordinatorWriteSetLoggingEnabled,
+      WriteSetEncoder writeSetEncoder,
       CoordinatorCommitHandler coordinatorCommitHandler,
       ParticipantCommitHandler participantCommitHandler) {
     this.coordinatorWriteOmissionOnReadOnlyEnabled = coordinatorWriteOmissionOnReadOnlyEnabled;
+    this.coordinatorWriteSetLoggingEnabled = coordinatorWriteSetLoggingEnabled;
+    this.writeSetEncoder = checkNotNull(writeSetEncoder);
     this.coordinatorCommitHandler = checkNotNull(coordinatorCommitHandler);
     this.participantCommitHandler = checkNotNull(participantCommitHandler);
   }
@@ -296,22 +302,44 @@ void onePhaseCommitRecords(TransactionContext context)
 
   public long commitState(TransactionContext context)
       throws CommitConflictException, UnknownTransactionStatusException {
-    return coordinatorCommitHandler.commitState(context);
+    return coordinatorCommitHandler.commitState(
+        context.transactionId, encodeWriteSetIfLoggingEnabled(context));
+  }
+
+  // The tx_write_set column is part of the Coordinator schema only when the opt-in
+  // `coordinator.write_set_logging.enabled` config is on. When it is off, skip encoding entirely so
+  // no WriteSet is persisted (the column is not part of the schema in that case).
+  @Nullable
+  private WriteSet encodeWriteSetIfLoggingEnabled(TransactionContext context) {
+    return coordinatorWriteSetLoggingEnabled
+        ? writeSetEncoder.encodeSingleGroupWriteSet(context, false)
+        : null;
   }
 
+  // 2PC-only. Delegates to commitState(id, null) / abortState(id, null), which the group-commit
+  // coordinator handler overrides to route through the group committer. That null never reaches the
+  // group committer here because 2PC forbids group commit
+  // (TwoPhaseConsensusCommitManager#throwIfGroupCommitIsEnabled) and never builds a
+  // CommitHandlerWithGroupCommit.
+  //
+  // TODO: revisit this if/when the Two-phase Commit I/F is removed.
   public long commitStateWithoutWriteSet(TransactionContext context)
       throws CommitConflictException, UnknownTransactionStatusException {
-    return coordinatorCommitHandler.commitStateWithoutWriteSet(context);
+    return coordinatorCommitHandler.commitState(context.transactionId, null);
   }
 
   public TransactionState abortState(TransactionContext context)
       throws UnknownTransactionStatusException {
-    return coordinatorCommitHandler.abortState(context);
+    return coordinatorCommitHandler.abortState(
+        context.transactionId, encodeWriteSetIfLoggingEnabled(context));
   }
 
+  // 2PC-only; see the null-write-set note on commitStateWithoutWriteSet above.
+  //
+  // TODO: revisit this if/when the Two-phase Commit I/F is removed.
   public TransactionState abortStateWithoutWriteSet(String id)
       throws UnknownTransactionStatusException {
-    return coordinatorCommitHandler.abortStateWithoutWriteSet(id);
+    return coordinatorCommitHandler.abortState(id, null);
   }
 
   public TransactionState forceAbortState(String id) throws UnknownTransactionStatusException {

core/src/main/java/com/scalar/db/transaction/consensuscommit/CommitHandlerWithGroupCommit.java

@@ -3,6 +3,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.scalar.db.api.DistributedStorage;
+import com.scalar.db.api.TransactionState;
 import com.scalar.db.exception.transaction.CommitConflictException;
 import com.scalar.db.exception.transaction.CommitException;
 import com.scalar.db.exception.transaction.UnknownTransactionStatusException;
@@ -59,11 +60,10 @@ private CommitHandlerWithGroupCommit(
       CoordinatorGroupCommitter groupCommitter) {
     this(
         coordinatorWriteOmissionOnReadOnlyEnabled,
+        coordinatorWriteSetLoggingEnabled,
+        new WriteSetEncoder(tableMetadataManager),
         new CoordinatorCommitHandlerWithGroupCommit(
-            coordinator,
-            new WriteSetEncoder(tableMetadataManager),
-            groupCommitter,
-            coordinatorWriteSetLoggingEnabled),
+            coordinator, groupCommitter, coordinatorWriteSetLoggingEnabled),
         participantCommitHandler);
   }
 
@@ -72,9 +72,16 @@ private CommitHandlerWithGroupCommit(
   @SuppressFBWarnings("EI_EXPOSE_REP2")
   CommitHandlerWithGroupCommit(
       boolean coordinatorWriteOmissionOnReadOnlyEnabled,
+      boolean coordinatorWriteSetLoggingEnabled,
+      WriteSetEncoder writeSetEncoder,
       CoordinatorCommitHandlerWithGroupCommit coordinatorHandler,
       ParticipantCommitHandler participantCommitHandler) {
-    super(coordinatorWriteOmissionOnReadOnlyEnabled, coordinatorHandler, participantCommitHandler);
+    super(
+        coordinatorWriteOmissionOnReadOnlyEnabled,
+        coordinatorWriteSetLoggingEnabled,
+        writeSetEncoder,
+        coordinatorHandler,
+        participantCommitHandler);
     this.coordinatorHandler = coordinatorHandler;
   }
 
@@ -87,45 +94,68 @@ public void commit(TransactionContext context)
 
   /**
    * Releases the group-commit slot reservation when the orchestrator will not write a Coordinator
-   * state row through the normal 2-phase commit path. That happens when the transaction was begun
-   * as non-read-only (and therefore reserved a slot) but turned out to have no writes/deletes,
-   * combined with coordinator-write-omission being enabled — in which case {@link
-   * CommitHandler#commit} skips {@code commitState} entirely and the reserved slot would otherwise
-   * sit unused for the rest of the commit flow.
-   *
-   * <p>The other places that call {@link
-   * CoordinatorCommitHandlerWithGroupCommit#cancelGroupCommitIfNeeded} (one-phase commit,
-   * abort/conflict paths, before-commit failure cleanup) cannot share this predicate: at those
-   * sites the transaction has writes/deletes, so the predicate would block the cancellation that
-   * those paths actually need.
+   * state row through the normal 2-phase commit path: a non-read-only transaction (which reserved a
+   * slot) that turns out to have no writes/deletes while coordinator-write-omission is enabled, so
+   * {@link CommitHandler#commit} skips {@code commitState} and the slot would otherwise sit unused.
    */
   private void cancelGroupCommitIfCoordinatorStateOmitted(TransactionContext context) {
     if (!context.readOnly
         && !context.snapshot.hasWritesOrDeletes()
         && coordinatorWriteOmissionOnReadOnlyEnabled) {
-      coordinatorHandler.cancelGroupCommitIfNeeded(context);
+      cancelGroupCommitIfSlotReserved(context);
+    }
+  }
+
+  private void cancelGroupCommitIfSlotReserved(TransactionContext context) {
+    // A transaction only holds a group commit slot when one was reserved at begin time (e.g., a
+    // read-only transaction does not reserve a slot when coordinator write omission is enabled).
+    // When no slot was reserved there is nothing to release.
+    if (!context.groupCommitSlotReserved) {
+      return;
     }
+
+    // FIXME: When a slot was reserved (so the early return above did not fire), this can run more
+    // than once on a single failure path, because several cleanup callbacks each release the slot:
+    //   - onFailureBeforeCommit() runs on the pre-commit-state failure paths (via
+    //     safelyCallOnFailureBeforeCommit);
+    //   - abortState() additionally runs from abortStateAndRollbackRecordsIfNeeded() when the
+    //     transaction has writes/deletes, or coordinator write omission on read-only is disabled;
+    //   - onePhaseCommitRecords() pre-cancels the slot before its body, so a one-phase-commit
+    //     failure followed by onFailureBeforeCommit() also cancels twice.
+    // It is currently safe only because groupCommitter.remove() is idempotent. The cleanup paths
+    // should be reworked so the slot is released exactly once instead of relying on that
+    // idempotency.
+    coordinatorHandler.cancelGroupCommit(context.transactionId);
   }
 
   @Override
   boolean canOnePhaseCommit(TransactionContext context) throws CommitException {
     try {
       return super.canOnePhaseCommit(context);
     } catch (CommitException e) {
-      coordinatorHandler.cancelGroupCommitIfNeeded(context);
+      cancelGroupCommitIfSlotReserved(context);
       throw e;
     }
   }
 
   @Override
   void onePhaseCommitRecords(TransactionContext context)
       throws CommitConflictException, UnknownTransactionStatusException {
-    coordinatorHandler.cancelGroupCommitIfNeeded(context);
+    cancelGroupCommitIfSlotReserved(context);
     super.onePhaseCommitRecords(context);
   }
 
   @Override
   protected void onFailureBeforeCommit(TransactionContext context) {
-    coordinatorHandler.cancelGroupCommitIfNeeded(context);
+    cancelGroupCommitIfSlotReserved(context);
+  }
+
+  @Override
+  public TransactionState abortState(TransactionContext context)
+      throws UnknownTransactionStatusException {
+    // Release the group-commit slot before writing the ABORTED state so it does not sit in the
+    // group buffer waiting for a sibling.
+    cancelGroupCommitIfSlotReserved(context);
+    return super.abortState(context);
   }
 }

core/src/main/java/com/scalar/db/transaction/consensuscommit/CoordinatorCommitHandler.java

@@ -18,64 +18,35 @@
  * Handles the Coordinator-side (state-table) operations of the Consensus Commit protocol: writing
  * the COMMITTED / ABORTED records to the Coordinator state table and resolving putState conflicts.
  *
- * <p>Methods here only touch the Coordinator state table via {@link Coordinator}; they never touch
- * user data tables. When {@code commitState} loses a putState race and the persisted state turns
- * out to be ABORTED (or absent), this handler reports a {@link CommitConflictException} and leaves
- * the rollback of the transaction's prepared records to the caller. This keeps the handler free of
- * any participant-side dependency.
+ * <p>Operates on a transaction id and a pre-encoded {@link WriteSet} (plus, for the commit path,
+ * the {@code committedAt} to stamp); callers encode the write set before calling.
+ *
+ * <p>When {@code commitState} loses a putState race and the persisted state turns out to be ABORTED
+ * (or absent), this handler reports a {@link CommitConflictException} and leaves the rollback of
+ * the transaction's prepared records to the caller.
  */
 @ThreadSafe
 class CoordinatorCommitHandler {
   private static final Logger logger = LoggerFactory.getLogger(CoordinatorCommitHandler.class);
 
   private final Coordinator coordinator;
-  private final WriteSetEncoder writeSetEncoder;
-  private final boolean coordinatorWriteSetLoggingEnabled;
 
   @SuppressFBWarnings("EI_EXPOSE_REP2")
-  CoordinatorCommitHandler(
-      Coordinator coordinator,
-      WriteSetEncoder writeSetEncoder,
-      boolean coordinatorWriteSetLoggingEnabled) {
+  CoordinatorCommitHandler(Coordinator coordinator) {
     this.coordinator = checkNotNull(coordinator);
-    this.writeSetEncoder = checkNotNull(writeSetEncoder);
-    this.coordinatorWriteSetLoggingEnabled = coordinatorWriteSetLoggingEnabled;
-  }
-
-  /**
-   * Writes the COMMITTED state. When coordinator write-set logging is enabled, the {@code
-   * tx_write_set} column is populated from the given context's snapshot; otherwise the column is
-   * omitted.
-   *
-   * @return the {@code committedAt} timestamp written to the COMMITTED Coordinator state row. The
-   *     caller stamps the committed data records with the same value so the row and the records
-   *     share one timestamp.
-   */
-  long commitState(TransactionContext context)
-      throws CommitConflictException, UnknownTransactionStatusException {
-    // The tx_write_set column is added only when the opt-in `coordinator.write_set_logging.enabled`
-    // config is enabled. When it's disabled the column is not part of the Coordinator schema, so we
-    // must skip encoding/persisting the WriteSet entirely.
-    WriteSet writeSet =
-        coordinatorWriteSetLoggingEnabled
-            ? writeSetEncoder.encodeSingleGroupWriteSet(context, false)
-            : null;
-    return commitStateInternal(context, writeSet);
   }
 
   /**
-   * Writes the COMMITTED state without persisting a {@code tx_write_set}.
+   * Writes the COMMITTED state for the transaction identified by {@code id}, persisting the given
+   * pre-encoded {@code writeSet} (or none when {@code null}).
    *
-   * @return the {@code committedAt} timestamp written to the COMMITTED Coordinator state row.
+   * @return the {@code committedAt} stamped on the COMMITTED Coordinator state row — the value this
+   *     call generates, or, when this commit lost a putState race to an already-COMMITTED row, that
+   *     row's committedAt. The caller stamps the committed data records with the returned value so
+   *     the row and the records share one timestamp.
    */
-  long commitStateWithoutWriteSet(TransactionContext context)
+  long commitState(String id, @Nullable WriteSet writeSet)
       throws CommitConflictException, UnknownTransactionStatusException {
-    return commitStateInternal(context, null);
-  }
-
-  private long commitStateInternal(TransactionContext context, @Nullable WriteSet writeSet)
-      throws CommitConflictException, UnknownTransactionStatusException {
-    String id = context.transactionId;
     try {
       long committedAt = System.currentTimeMillis();
       Coordinator.State state =
@@ -84,7 +55,7 @@ private long commitStateInternal(TransactionContext context, @Nullable WriteSet
       logger.debug("Transaction {} is committed successfully at {}", id, committedAt);
       return committedAt;
     } catch (CoordinatorConflictException e) {
-      return handleCommitConflict(context, e);
+      return handleCommitConflict(id, e);
     } catch (CoordinatorException e) {
       throw new UnknownTransactionStatusException(
           CoreError.CONSENSUS_COMMIT_UNKNOWN_COORDINATOR_STATUS.buildMessage(e.getMessage()),
@@ -95,21 +66,21 @@ private long commitStateInternal(TransactionContext context, @Nullable WriteSet
 
   /**
    * Resolves a putState conflict. Returns the {@code committedAt} of the already-persisted
-   * COMMITTED state when this transaction turns out to be already committed; otherwise rolls back
-   * the records and throws.
+   * COMMITTED state when this transaction turns out to be already committed; otherwise reports the
+   * conflict (the caller rolls the records back) by throwing {@link CommitConflictException}.
    */
-  long handleCommitConflict(TransactionContext context, Exception cause)
+  long handleCommitConflict(String id, Exception cause)
       throws CommitConflictException, UnknownTransactionStatusException {
     try {
-      Optional<Coordinator.State> s = coordinator.getState(context.transactionId);
+      Optional<Coordinator.State> s = coordinator.getState(id);
       if (s.isPresent()) {
         Coordinator.State persisted = s.get();
         if (persisted.getState() == TransactionState.ABORTED) {
           throw new CommitConflictException(
               CoreError.CONSENSUS_COMMIT_CONFLICT_OCCURRED_WHEN_COMMITTING_STATE.buildMessage(
                   cause.getMessage()),
               cause,
-              context.transactionId);
+              id);
         }
         // Otherwise the coordinator state is present and COMMITTED, which means this transaction
         // has already committed. Only Two-phase Commit I/F reaches this branch: there the same
@@ -163,37 +134,21 @@ long handleCommitConflict(TransactionContext context, Exception cause)
             CoreError.CONSENSUS_COMMIT_CONFLICT_OCCURRED_WHEN_COMMITTING_STATE.buildMessage(
                 cause.getMessage()),
             cause,
-            context.transactionId);
+            id);
       }
     } catch (CoordinatorException ex) {
       throw new UnknownTransactionStatusException(
           CoreError.CONSENSUS_COMMIT_CANNOT_GET_COORDINATOR_STATUS.buildMessage(ex.getMessage()),
           ex,
-          context.transactionId);
+          id);
     }
   }
 
   /**
-   * Writes the ABORTED state. When coordinator write-set logging is enabled, the {@code
-   * tx_write_set} column is populated from the given context's snapshot; otherwise the column is
-   * omitted.
+   * Writes the ABORTED state for the transaction identified by {@code id}, persisting the given
+   * pre-encoded {@code writeSet} (or none when {@code null}).
    */
-  TransactionState abortState(TransactionContext context) throws UnknownTransactionStatusException {
-    // Same opt-in gating as commitState: skip WriteSet encoding when write-set logging is
-    // disabled, since the Coordinator schema does not include the column in that case.
-    WriteSet writeSet =
-        coordinatorWriteSetLoggingEnabled
-            ? writeSetEncoder.encodeSingleGroupWriteSet(context, false)
-            : null;
-    return abortStateInternal(context.transactionId, writeSet);
-  }
-
-  /** Writes the ABORTED state without persisting a {@code tx_write_set} via a single putState. */
-  TransactionState abortStateWithoutWriteSet(String id) throws UnknownTransactionStatusException {
-    return abortStateInternal(id, null);
-  }
-
-  private TransactionState abortStateInternal(String id, @Nullable WriteSet writeSet)
+  TransactionState abortState(String id, @Nullable WriteSet writeSet)
       throws UnknownTransactionStatusException {
     try {
       Coordinator.State state =