newTopicPartitions = new HashSet<>(topicPartitions);
+ newTopicPartitions.addAll(addedTopicPartitions);
+
+ TransitionData data = new TransitionData(TransactionState.ONGOING);
+ data.topicPartitions = newTopicPartitions;
+ data.txnStartTimestamp = newTxnStartTimestamp;
+ data.txnLastUpdateTimestamp = updateTimestamp;
+ data.clientTransactionVersion = clientTransactionVersion;
+ return prepareTransitionTo(data);
+ }
+
+ public TxnTransitMetadata prepareAbortOrCommit(TransactionState newState,
+ TransactionVersion clientTransactionVersion,
+ long nextProducerId,
+ long updateTimestamp,
+ boolean noPartitionAdded) {
+ TransitionData data = new TransitionData(newState);
+ if (clientTransactionVersion.supportsEpochBump()) {
+ // We already ensured that we do not overflow here. MAX_SHORT is the highest possible value.
+ data.producerEpoch = (short) (producerEpoch + 1);
+ data.lastProducerEpoch = producerEpoch;
+ } else {
+ data.producerEpoch = producerEpoch;
+ data.lastProducerEpoch = lastProducerEpoch;
+ }
+
+ // With transaction V2, it is allowed to abort the transaction without adding any partitions. Then, the transaction
+ // start time is uncertain but it is still required. So we can use the update time as the transaction start time.
+ data.txnStartTimestamp = noPartitionAdded ? updateTimestamp : txnStartTimestamp;
+ data.nextProducerId = nextProducerId;
+ data.txnLastUpdateTimestamp = updateTimestamp;
+ data.clientTransactionVersion = clientTransactionVersion;
+ return prepareTransitionTo(data);
+ }
+
+ public TxnTransitMetadata prepareComplete(long updateTimestamp) {
+ // Since the state change was successfully written to the log, unset the flag for a failed epoch fence
+ hasFailedEpochFence = false;
+
+ TransitionData data = new TransitionData(state == TransactionState.PREPARE_COMMIT ?
+ TransactionState.COMPLETE_COMMIT : TransactionState.COMPLETE_ABORT);
+ // In the prepareComplete transition for the overflow case, the lastProducerEpoch is kept at MAX-1,
+ // which is the last epoch visible to the client.
+ // Internally, however, during the transition between prepareAbort/prepareCommit and prepareComplete, the producer epoch
+ // reaches MAX but the client only sees the transition as MAX-1 followed by 0.
+ // When an epoch overflow occurs, we set the producerId to nextProducerId and reset the epoch to 0,
+ // but lastProducerEpoch remains MAX-1 to maintain consistency with what the client last saw.
+ if (clientTransactionVersion.supportsEpochBump() && nextProducerId != RecordBatch.NO_PRODUCER_ID) {
+ data.producerId = nextProducerId;
+ data.producerEpoch = 0;
+ } else {
+ data.producerId = producerId;
+ data.producerEpoch = producerEpoch;
+ }
+ data.nextProducerId = RecordBatch.NO_PRODUCER_ID;
+ data.topicPartitions = new HashSet<>();
+ data.txnLastUpdateTimestamp = updateTimestamp;
+ return prepareTransitionTo(data);
+ }
+
+ public TxnTransitMetadata prepareDead() {
+ TransitionData data = new TransitionData(TransactionState.DEAD);
+ data.topicPartitions = new HashSet<>();
+ return prepareTransitionTo(data);
+ }
+
+ /**
+ * Check if the epochs have been exhausted for the current producerId. We do not allow the client to use an
+ * epoch equal to Short.MaxValue to ensure that the coordinator will always be able to fence an existing producer.
+ */
+ public boolean isProducerEpochExhausted() {
+ return isEpochExhausted(producerEpoch);
+ }
+
+ /**
+ * Check if this is a distributed two phase commit transaction.
+ * Such transactions have no timeout (identified by maximum value for timeout).
+ */
+ public boolean isDistributedTwoPhaseCommitTxn() {
+ return txnTimeoutMs == Integer.MAX_VALUE;
+ }
+
+ private boolean hasPendingTransaction() {
+ return state == TransactionState.ONGOING ||
+ state == TransactionState.PREPARE_ABORT ||
+ state == TransactionState.PREPARE_COMMIT;
+ }
+
+ private TxnTransitMetadata prepareTransitionTo(TransitionData data) {
+ if (pendingState.isPresent())
+ throw new IllegalStateException("Preparing transaction state transition to " + state +
+ " while it already a pending state " + pendingState.get());
+
+ if (data.producerId < 0)
+ throw new IllegalArgumentException("Illegal new producer id " + data.producerId);
+
+ // The epoch is initialized to NO_PRODUCER_EPOCH when the TransactionMetadata
+ // is created for the first time and it could stay like this until transitioning
+ // to Dead.
+ if (data.state != TransactionState.DEAD && data.producerEpoch < 0)
+ throw new IllegalArgumentException("Illegal new producer epoch " + data.producerEpoch);
+
+ // check that the new state transition is valid and update the pending state if necessary
+ if (data.state.validPreviousStates().contains(this.state)) {
+ TxnTransitMetadata transitMetadata = new TxnTransitMetadata(
+ data.producerId, this.producerId, data.nextProducerId, data.producerEpoch, data.lastProducerEpoch,
+ data.txnTimeoutMs, data.state, data.topicPartitions,
+ data.txnStartTimestamp, data.txnLastUpdateTimestamp, data.clientTransactionVersion
+ );
+
+ LOGGER.debug("TransactionalId {} prepare transition from {} to {}", transactionalId, this.state, data.state);
+ pendingState = Optional.of(data.state);
+ return transitMetadata;
+ }
+ throw new IllegalStateException("Preparing transaction state transition to " + data.state + " failed since the target state " +
+ data.state + " is not a valid previous state of the current state " + this.state);
+ }
+
+ @SuppressWarnings("CyclomaticComplexity")
+ public void completeTransitionTo(TxnTransitMetadata transitMetadata) {
+ // metadata transition is valid only if all the following conditions are met:
+ //
+ // 1. the new state is already indicated in the pending state.
+ // 2. the epoch should be either the same value, the old value + 1, or 0 if we have a new producerId.
+ // 3. the last update time is no smaller than the old value.
+ // 4. the old partitions set is a subset of the new partitions set.
+ //
+ // plus, we should only try to update the metadata after the corresponding log entry has been successfully
+ // written and replicated (see TransactionStateManager#appendTransactionToLog)
+ //
+ // if valid, transition is done via overwriting the whole object to ensure synchronization
+
+ TransactionState toState = pendingState.orElseThrow(() -> {
+ LOGGER.error(MarkerFactory.getMarker(LogLevelConfig.FATAL_LOG_LEVEL),
+ "{}'s transition to {} failed since pendingState is not defined: this should not happen", this, transitMetadata);
+ return new IllegalStateException("TransactionalId " + transactionalId +
+ " completing transaction state transition while it does not have a pending state");
+ });
+
+ if (!toState.equals(transitMetadata.txnState())) throwStateTransitionFailure(transitMetadata);
+
+ switch (toState) {
+ case EMPTY: // from initPid
+ if ((producerEpoch != transitMetadata.producerEpoch() && !validProducerEpochBump(transitMetadata)) ||
+ !transitMetadata.topicPartitions().isEmpty() ||
+ transitMetadata.txnStartTimestamp() != -1) {
+ throwStateTransitionFailure(transitMetadata);
+ }
+ break;
+
+ case ONGOING: // from addPartitions
+ if (!validProducerEpoch(transitMetadata) ||
+ !transitMetadata.topicPartitions().containsAll(topicPartitions) ||
+ txnTimeoutMs != transitMetadata.txnTimeoutMs()) {
+ throwStateTransitionFailure(transitMetadata);
+ }
+ break;
+
+ case PREPARE_ABORT: // from endTxn
+ case PREPARE_COMMIT:
+ // In V2, we allow state transits from Empty, CompleteCommit and CompleteAbort to PrepareAbort. It is possible
+ // their updated start time is not equal to the current start time.
+ boolean allowedEmptyAbort = toState == TransactionState.PREPARE_ABORT && transitMetadata.clientTransactionVersion().supportsEpochBump() &&
+ (state == TransactionState.EMPTY || state == TransactionState.COMPLETE_COMMIT || state == TransactionState.COMPLETE_ABORT);
+ boolean validTimestamp = txnStartTimestamp == transitMetadata.txnStartTimestamp() || allowedEmptyAbort;
+
+ if (!validProducerEpoch(transitMetadata) ||
+ !topicPartitions.equals(transitMetadata.topicPartitions()) ||
+ txnTimeoutMs != transitMetadata.txnTimeoutMs() ||
+ !validTimestamp) {
+ throwStateTransitionFailure(transitMetadata);
+ }
+ break;
+
+ case COMPLETE_ABORT: // from write markers
+ case COMPLETE_COMMIT:
+ if (!validProducerEpoch(transitMetadata) ||
+ txnTimeoutMs != transitMetadata.txnTimeoutMs() ||
+ transitMetadata.txnStartTimestamp() == -1) {
+ throwStateTransitionFailure(transitMetadata);
+ }
+ break;
+
+ case PREPARE_EPOCH_FENCE:
+ // We should never get here, since once we prepare to fence the epoch, we immediately set the pending state
+ // to PrepareAbort, and then consequently to CompleteAbort after the markers are written.. So we should never
+ // ever try to complete a transition to PrepareEpochFence, as it is not a valid previous state for any other state, and hence
+ // can never be transitioned out of.
+ throwStateTransitionFailure(transitMetadata);
+ break;
+
+ case DEAD:
+ // The transactionalId was being expired. The completion of the operation should result in removal of the
+ // the metadata from the cache, so we should never realistically transition to the dead state.
+ throw new IllegalStateException("TransactionalId " + transactionalId + " is trying to complete a transition to " +
+ toState + ". This means that the transactionalId was being expired, and the only acceptable completion of " +
+ "this operation is to remove the transaction metadata from the cache, not to persist the " + toState + " in the log.");
+
+ default:
+ break;
+ }
+
+ LOGGER.debug("TransactionalId {} complete transition from {} to {}", transactionalId, state, transitMetadata);
+ producerId = transitMetadata.producerId();
+ prevProducerId = transitMetadata.prevProducerId();
+ nextProducerId = transitMetadata.nextProducerId();
+ producerEpoch = transitMetadata.producerEpoch();
+ lastProducerEpoch = transitMetadata.lastProducerEpoch();
+ txnTimeoutMs = transitMetadata.txnTimeoutMs();
+ topicPartitions = transitMetadata.topicPartitions();
+ txnStartTimestamp = transitMetadata.txnStartTimestamp();
+ txnLastUpdateTimestamp = transitMetadata.txnLastUpdateTimestamp();
+ clientTransactionVersion = transitMetadata.clientTransactionVersion();
+
+ pendingState = Optional.empty();
+ state = toState;
+ }
+
+ /**
+ * Validates the producer epoch and ID based on transaction state and version.
+ *
+ * Logic:
+ * * 1. **Overflow Case in Transactions V2:**
+ * * - During overflow (epoch reset to 0), we compare both `lastProducerEpoch` values since it
+ * * does not change during completion.
+ * * - For PrepareComplete, the producer ID has been updated. We ensure that the `prevProducerID`
+ * * in the transit metadata matches the current producer ID, confirming the change.
+ * *
+ * * 2. **Epoch Bump Case in Transactions V2:**
+ * * - For PrepareCommit or PrepareAbort, the producer epoch has been bumped. We ensure the `lastProducerEpoch`
+ * * in transit metadata matches the current producer epoch, confirming the bump.
+ * * - We also verify that the producer ID remains the same.
+ * *
+ * * 3. **Other Cases:**
+ * * - For other states and versions, check if the producer epoch and ID match the current values.
+ *
+ * @param transitMetadata The transaction transition metadata containing state, producer epoch, and ID.
+ * @return true if the producer epoch and ID are valid; false otherwise.
+ */
+ private boolean validProducerEpoch(TxnTransitMetadata transitMetadata) {
+ boolean isAtLeastTransactionsV2 = transitMetadata.clientTransactionVersion().supportsEpochBump();
+ TransactionState txnState = transitMetadata.txnState();
+ short transitProducerEpoch = transitMetadata.producerEpoch();
+ long transitProducerId = transitMetadata.producerId();
+ short transitLastProducerEpoch = transitMetadata.lastProducerEpoch();
+
+ if (isAtLeastTransactionsV2 &&
+ (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT) &&
+ transitProducerEpoch == 0) {
+ return transitLastProducerEpoch == lastProducerEpoch && transitMetadata.prevProducerId() == producerId;
+ }
+
+ if (isAtLeastTransactionsV2 &&
+ (txnState == TransactionState.PREPARE_COMMIT || txnState == TransactionState.PREPARE_ABORT)) {
+ return transitLastProducerEpoch == producerEpoch && transitProducerId == producerId;
+ }
+ return transitProducerEpoch == producerEpoch && transitProducerId == producerId;
+ }
+
+ private boolean validProducerEpochBump(TxnTransitMetadata transitMetadata) {
+ short transitEpoch = transitMetadata.producerEpoch();
+ long transitProducerId = transitMetadata.producerId();
+ return transitEpoch == (short) (producerEpoch + 1) || (transitEpoch == 0 && transitProducerId != producerId);
+ }
+
+ private void throwStateTransitionFailure(TxnTransitMetadata txnTransitMetadata) {
+ LOGGER.error(MarkerFactory.getMarker(LogLevelConfig.FATAL_LOG_LEVEL),
+ "{}'s transition to {} failed: this should not happen", this, txnTransitMetadata);
+
+ throw new IllegalStateException("TransactionalId " + transactionalId + " failed transition to state " + txnTransitMetadata +
+ " due to unexpected metadata");
+ }
+
+ public boolean pendingTransitionInProgress() {
+ return pendingState.isPresent();
+ }
+
+ public String transactionalId() {
+ return transactionalId;
+ }
+
+ public void setProducerId(long producerId) {
+ this.producerId = producerId;
+ }
+ public long producerId() {
+ return producerId;
+ }
+
+ public void setPrevProducerId(long prevProducerId) {
+ this.prevProducerId = prevProducerId;
+ }
+ public long prevProducerId() {
+ return prevProducerId;
+ }
+
+ public void setProducerEpoch(short producerEpoch) {
+ this.producerEpoch = producerEpoch;
+ }
+
+ public short producerEpoch() {
+ return producerEpoch;
+ }
+
+ public void setLastProducerEpoch(short lastProducerEpoch) {
+ this.lastProducerEpoch = lastProducerEpoch;
+ }
+
+ public short lastProducerEpoch() {
+ return lastProducerEpoch;
+ }
+
+ public int txnTimeoutMs() {
+ return txnTimeoutMs;
+ }
+
+ public void state(TransactionState state) {
+ this.state = state;
+ }
+
+ public TransactionState state() {
+ return state;
+ }
+
+ public Set topicPartitions() {
+ return topicPartitions;
+ }
+
+ public long txnStartTimestamp() {
+ return txnStartTimestamp;
+ }
+
+ public void txnLastUpdateTimestamp(long txnLastUpdateTimestamp) {
+ this.txnLastUpdateTimestamp = txnLastUpdateTimestamp;
+ }
+
+ public long txnLastUpdateTimestamp() {
+ return txnLastUpdateTimestamp;
+ }
+
+ public void clientTransactionVersion(TransactionVersion clientTransactionVersion) {
+ this.clientTransactionVersion = clientTransactionVersion;
+ }
+
+ public TransactionVersion clientTransactionVersion() {
+ return clientTransactionVersion;
+ }
+
+ public void pendingState(Optional pendingState) {
+ this.pendingState = pendingState;
+ }
+
+ public Optional pendingState() {
+ return pendingState;
+ }
+
+ public void hasFailedEpochFence(boolean hasFailedEpochFence) {
+ this.hasFailedEpochFence = hasFailedEpochFence;
+ }
+
+ public boolean hasFailedEpochFence() {
+ return hasFailedEpochFence;
+ }
+
+ @Override
+ public String toString() {
+ return "TransactionMetadata(" +
+ "transactionalId=" + transactionalId +
+ ", producerId=" + producerId +
+ ", prevProducerId=" + prevProducerId +
+ ", nextProducerId=" + nextProducerId +
+ ", producerEpoch=" + producerEpoch +
+ ", lastProducerEpoch=" + lastProducerEpoch +
+ ", txnTimeoutMs=" + txnTimeoutMs +
+ ", state=" + state +
+ ", pendingState=" + pendingState +
+ ", topicPartitions=" + topicPartitions +
+ ", txnStartTimestamp=" + txnStartTimestamp +
+ ", txnLastUpdateTimestamp=" + txnLastUpdateTimestamp +
+ ", clientTransactionVersion=" + clientTransactionVersion +
+ ")";
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) return true;
+ if (obj == null || getClass() != obj.getClass()) return false;
+
+ TransactionMetadata other = (TransactionMetadata) obj;
+ return transactionalId.equals(other.transactionalId) &&
+ producerId == other.producerId &&
+ prevProducerId == other.prevProducerId &&
+ nextProducerId == other.nextProducerId &&
+ producerEpoch == other.producerEpoch &&
+ lastProducerEpoch == other.lastProducerEpoch &&
+ txnTimeoutMs == other.txnTimeoutMs &&
+ state.equals(other.state) &&
+ topicPartitions.equals(other.topicPartitions) &&
+ txnStartTimestamp == other.txnStartTimestamp &&
+ txnLastUpdateTimestamp == other.txnLastUpdateTimestamp &&
+ clientTransactionVersion.equals(other.clientTransactionVersion);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ transactionalId,
+ producerId,
+ prevProducerId,
+ nextProducerId,
+ producerEpoch,
+ lastProducerEpoch,
+ txnTimeoutMs,
+ state,
+ topicPartitions,
+ txnStartTimestamp,
+ txnLastUpdateTimestamp,
+ clientTransactionVersion
+ );
+ }
+
+ /**
+ * This class is used to hold the data that is needed to transition the transaction metadata to a new state.
+ * The data is copied from the current transaction metadata to avoid a lot of duplicated code in the prepare methods.
+ */
+ private class TransitionData {
+ final TransactionState state;
+ long producerId = TransactionMetadata.this.producerId;
+ long nextProducerId = TransactionMetadata.this.nextProducerId;
+ short producerEpoch = TransactionMetadata.this.producerEpoch;
+ short lastProducerEpoch = TransactionMetadata.this.lastProducerEpoch;
+ int txnTimeoutMs = TransactionMetadata.this.txnTimeoutMs;
+ HashSet topicPartitions = TransactionMetadata.this.topicPartitions;
+ long txnStartTimestamp = TransactionMetadata.this.txnStartTimestamp;
+ long txnLastUpdateTimestamp = TransactionMetadata.this.txnLastUpdateTimestamp;
+ TransactionVersion clientTransactionVersion = TransactionMetadata.this.clientTransactionVersion;
+
+ private TransitionData(TransactionState state) {
+ this.state = state;
+ }
+ }
+}
\ No newline at end of file
diff --git a/transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TxnTransitMetadata.java b/transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TxnTransitMetadata.java
index 59a05227d7da7..452c168687e56 100644
--- a/transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TxnTransitMetadata.java
+++ b/transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TxnTransitMetadata.java
@@ -19,7 +19,7 @@
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.server.common.TransactionVersion;
-import java.util.Set;
+import java.util.HashSet;
/**
* Immutable object representing the target transition of the transaction metadata
@@ -32,7 +32,9 @@ public record TxnTransitMetadata(
short lastProducerEpoch,
int txnTimeoutMs,
TransactionState txnState,
- Set topicPartitions,
+ // The TransactionMetadata#topicPartitions field is mutable.
+ // To avoid deep copy when assigning value from TxnTransitMetadata to TransactionMetadata, use HashSet here.
+ HashSet topicPartitions,
long txnStartTimestamp,
long txnLastUpdateTimestamp,
TransactionVersion clientTransactionVersion