You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "artemlivshits (via GitHub)" <gi...@apache.org> on 2026/10/12 03:44:55 UTC
Re: [PR] KAFKA-15370: Support Participation in 2PC (KIP-939) (K/N) [kafka]
artemlivshits commented on code in PR #21735:
URL: https://github.com/apache/kafka/pull/21735#discussion_r4244496323
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionLog.java:
##########
@@ -84,18 +84,25 @@ public static byte[] valueToBytes(TxnTransitMetadata txnMetadata,
.setPartitionIds(entry.getValue().stream().map(TopicPartition::partition).toList())).toList();
}
- return MessageUtil.toVersionPrefixedBytes(
- transactionVersionLevel.transactionLogValueVersion(),
- new TransactionLogValue()
- .setProducerId(txnMetadata.producerId())
- .setProducerEpoch(txnMetadata.producerEpoch())
- .setTransactionTimeoutMs(txnMetadata.txnTimeoutMs())
- .setTransactionStatus(txnMetadata.txnState().id())
- .setTransactionLastUpdateTimestampMs(txnMetadata.txnLastUpdateTimestamp())
- .setTransactionStartTimestampMs(txnMetadata.txnStartTimestamp())
- .setTransactionPartitions(transactionPartitions)
- .setClientTransactionVersion(txnMetadata.clientTransactionVersion().featureLevel())
- );
+ short logVersion = transactionVersionLevel.transactionLogValueVersion();
+ TransactionLogValue value = new TransactionLogValue()
+ .setProducerId(txnMetadata.producerId())
+ .setProducerEpoch(txnMetadata.producerEpoch())
+ .setTransactionTimeoutMs(txnMetadata.txnTimeoutMs())
+ .setTransactionStatus(txnMetadata.txnState().id())
+ .setTransactionLastUpdateTimestampMs(txnMetadata.txnLastUpdateTimestamp())
+ .setTransactionStartTimestampMs(txnMetadata.txnStartTimestamp())
+ .setTransactionPartitions(transactionPartitions)
+ .setClientTransactionVersion(txnMetadata.clientTransactionVersion().featureLevel());
+
+ // Tagged fields are only supported in log version 1+ (flexible versions)
+ if (logVersion >= 1) {
+ value.setPreviousProducerId(txnMetadata.prevProducerId());
+ value.setNextProducerId(txnMetadata.nextProducerId());
+ value.setNextProducerEpoch(txnMetadata.nextProducerEpoch());
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
Review Comment:
I don't think we should add this condition. prevProducerId is overwritten on every transition with the pre-transition producerId (prepareTransitionTo passes this.producerId), so it doesn't identify the client's producer id. With keepPreparedTxn=true it breaks valid retries: in the example above (ongoing {42, 10}, client-facing {73, 85}), the commit completes with {pid, epoch} = {73, 86} and prevProducerId = 42 (the ongoing txn pid), and the client retries with {73, 85} -- `prevProducerId == producerId` would be 42 == 73 and reject the retry. It would also reject retries after failover onto records written before KAFKA-20310, where prevProducerId is loaded as -1. The producer id itself is already validated by the `clientProducerId != producerId && !retryOnOverflow` check below (and in COMPLETE_* clientProducerId == producerId).
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -927,9 +987,14 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
val txnMetadata = epochAndMetadata.transactionMetadata
txnMetadata.inLock(() => {
- if (txnMetadata.producerId != producerId)
+ // Note that if we did a few InitProducerId(keepPreparedTxn) the client's producer id
+ // may be different from the producerId (we track it in the nextProducerId), but at
+ // this point even nextProducerId could be different from the client's producer id
+ // if we generated a new nextProducerId due to nextProducerEpoch overflow. So we
+ // check with the copy of the txnMetadata.producerId. Ditto for the epoch.
+ if (txnMetadata.producerId != producerIdCopy)
Left(Errors.INVALID_PRODUCER_ID_MAPPING)
- else if (txnMetadata.producerEpoch != producerEpoch && txnMetadata.producerEpoch != producerEpoch + 1)
+ else if (txnMetadata.producerEpoch != producerEpochCopy && txnMetadata.producerEpoch != producerEpochCopy + 1)
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
+ case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
+ retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ }
+ } else {
+ // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
+ (!isFromClient || producerEpoch == txnMetadata.clientProducerEpoch) && producerEpoch >= txnMetadata.clientProducerEpoch
}
- } else {
- // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
- (!isFromClient || producerEpoch == txnMetadata.producerEpoch) && producerEpoch >= txnMetadata.producerEpoch
}
- }
- val isRetry = retryOnEpochBump || retryOnOverflow
-
- def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
- // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
- // This must also apply to the epoch-fence path.
- val nextProducerIdOrErrors =
- if (txnMetadata.isProducerEpochExhausted) {
- try {
- Right(producerIdManager.generateProducerId())
- } catch {
- case e: Exception => Left(Errors.forException(e))
+ val isRetry = retryOnEpochBump || retryOnOverflow
+
+ def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
+ // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
+ // This must also apply to the epoch-fence path.
+ val nextProducerIdOrErrors =
+ if (txnMetadata.isProducerEpochExhausted) {
+ try {
+ Right(producerIdManager.generateProducerId())
+ } catch {
+ case e: Exception => Left(Errors.forException(e))
+ }
+ } else {
+ Right(txnMetadata.nextProducerId)
}
+
+ // If the next producer epoch is set (which can happen if we used InitProducerId(keepPreparedTxn))
+ // then we also need to bump this epoch -- we want to fence requests that were issued with that
+ // epoch after the transaction is complete. Consider this example:
+ // 1. (initial state): Ongoing, producerId = 42, producerEpoch = 100, nextProducerId = -1, nextProducerEpoch = -1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, producerId = 42, producerEpoch = 100, nextProducerId = 42, nextProducerEpoch = 101
+ // 3. CommitTxn: timed out (but actually delayed)
+ // 4. CommitTxn + complete: CompleteCommit, producerId = 42, producerEpoch = 101, nextProducerId = -1, nextProducerEpoch = -1
+ // 5. Start new transaction, do some actions
+ // 6. Delayed CommitTxn from step 3 arrives and commits partially completed transaction.
+ // In order to prevent this problem, we need to bump the epoch at step 4 (so it becomes 102) even though we
+ // already bumped it at step 2.
+ // Note that we can arrive at this scenario when a client did a few keep-prepared calls and then
+ // decided to do one without keep-prepared (e.g. to force-terminate), so this logic applies
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -298,7 +354,7 @@ public TxnTransitMetadata prepareDead() {
* 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);
+ return isEpochExhausted(clientProducerEpoch());
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -247,7 +299,7 @@ public TxnTransitMetadata prepareAbortOrCommit(TransactionState newState,
// We already ensured that we do not overflow here. MAX_SHORT is the highest possible value.
data.producerEpoch = (short) (producerEpoch + 1);
}
- data.lastProducerEpoch = producerEpoch;
+ data.lastProducerEpoch = clientProducerEpoch();
Review Comment:
This is addressed on trunk by KAFKA-20357 (#22096), which persists LastProducerEpoch in the transaction log (tagged field 4) and reads it back on load, so it now survives coordinator failover. I've merged trunk into this branch; TransactionLog now writes/reads both lastProducerEpoch and nextProducerEpoch, and TransactionLogTest.testAllPersistedFieldsRoundTrip verifies lastProducerEpoch round-trips.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -1056,6 +935,548 @@ class TransactionsTest extends IntegrationTestHarness {
producer.abortTransaction()
}
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testFencingWithDualIdentity(groupProtocol: String): Unit = {
+ // Test that zombie producer is properly fenced when dual identity is established
+ // after crash recovery with 2PC keepPreparedTxn.
+ val transactionalId = "test-fencing-dual-identity"
+ val testTopic = s"test-fencing-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and send records (simulating crash before commit)
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+ // Simulate crash - don't commit (leave transaction prepared)
+
+ // Create new producer and recover with initTransactions(keepPreparedTxn=true)
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // This establishes dual identity
+
+ // Old (zombie) producer tries to abort - should get ProducerFencedException
+ assertThrows(classOf[ProducerFencedException], () => {
+ producer1.abortTransaction()
+ }, "Zombie producer should be fenced when trying to abort")
+
+ // New producer commits the prepared transaction
+ producer2.commitTransaction()
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records from the first producer
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records from first producer after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionCommit(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for commit path.
+ // This tests the client-side 2PC API where the producer explicitly calls prepareTransaction()
+ // and then completeTransaction() to commit.
+ val transactionalId = "test-prepare-complete-commit"
+ val testTopic = s"test-prepare-complete-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and produce records
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Prepare the transaction (moves to PREPARED state)
+ val preparedState = producer1.prepareTransaction()
+ assertNotNull(preparedState, "prepareTransaction should return prepared state")
+
+ // Simulate crash - don't call commit/abort, just leave in prepared state
+ // (In real scenario, producer would crash here)
+
+ // New producer recovers and completes the transaction
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete the transaction with the prepared state (should commit)
+ producer2.completeTransaction(preparedState)
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionAbort(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for abort path.
+ // This tests that completeTransaction() with a non-matching state triggers abort.
+ val transactionalId = "test-prepare-complete-abort"
+ val testTopic = s"test-prepare-complete-abort-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+
+ // First transaction - prepare but then abort
+ producer1.beginTransaction()
+ for (i <- 0 until 3) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key1-$i".getBytes, s"value1-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for first transaction
+ val marker1 = producer1.prepareTransaction()
+ assertNotNull(marker1, "prepareTransaction should return prepared state")
+
+ // Abort the first prepared transaction
+ producer1.abortTransaction()
+
+ // Second transaction - prepare and leave in prepared state
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key2-$i".getBytes, s"value2-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for second transaction
+ val marker2 = producer1.prepareTransaction()
+ assertNotNull(marker2, "prepareTransaction should return prepared state")
+
+ // Simulate crash - leave second transaction in prepared state
+
+ // New producer recovers
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete with marker1 (old, aborted transaction) - should trigger abort
+ // because it doesn't match the current prepared state (marker2)
+ producer2.completeTransaction(marker1)
+
+ // Wait for transaction to reach COMPLETE_ABORT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_ABORT)
+
+ // Verify consumer sees NO records (both transactions aborted)
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(0, consumedRecords.size,
+ "Consumer should see no records after abort")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
Review Comment:
Will do.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -1056,6 +935,548 @@ class TransactionsTest extends IntegrationTestHarness {
producer.abortTransaction()
}
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testFencingWithDualIdentity(groupProtocol: String): Unit = {
+ // Test that zombie producer is properly fenced when dual identity is established
+ // after crash recovery with 2PC keepPreparedTxn.
+ val transactionalId = "test-fencing-dual-identity"
+ val testTopic = s"test-fencing-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and send records (simulating crash before commit)
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+ // Simulate crash - don't commit (leave transaction prepared)
+
+ // Create new producer and recover with initTransactions(keepPreparedTxn=true)
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // This establishes dual identity
+
+ // Old (zombie) producer tries to abort - should get ProducerFencedException
+ assertThrows(classOf[ProducerFencedException], () => {
+ producer1.abortTransaction()
+ }, "Zombie producer should be fenced when trying to abort")
+
+ // New producer commits the prepared transaction
+ producer2.commitTransaction()
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records from the first producer
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records from first producer after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionCommit(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for commit path.
+ // This tests the client-side 2PC API where the producer explicitly calls prepareTransaction()
+ // and then completeTransaction() to commit.
+ val transactionalId = "test-prepare-complete-commit"
+ val testTopic = s"test-prepare-complete-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and produce records
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Prepare the transaction (moves to PREPARED state)
+ val preparedState = producer1.prepareTransaction()
+ assertNotNull(preparedState, "prepareTransaction should return prepared state")
+
+ // Simulate crash - don't call commit/abort, just leave in prepared state
+ // (In real scenario, producer would crash here)
+
+ // New producer recovers and completes the transaction
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete the transaction with the prepared state (should commit)
+ producer2.completeTransaction(preparedState)
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionAbort(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for abort path.
+ // This tests that completeTransaction() with a non-matching state triggers abort.
+ val transactionalId = "test-prepare-complete-abort"
+ val testTopic = s"test-prepare-complete-abort-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+
+ // First transaction - prepare but then abort
+ producer1.beginTransaction()
+ for (i <- 0 until 3) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key1-$i".getBytes, s"value1-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for first transaction
+ val marker1 = producer1.prepareTransaction()
+ assertNotNull(marker1, "prepareTransaction should return prepared state")
+
+ // Abort the first prepared transaction
+ producer1.abortTransaction()
+
+ // Second transaction - prepare and leave in prepared state
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key2-$i".getBytes, s"value2-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for second transaction
+ val marker2 = producer1.prepareTransaction()
+ assertNotNull(marker2, "prepareTransaction should return prepared state")
+
+ // Simulate crash - leave second transaction in prepared state
+
+ // New producer recovers
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete with marker1 (old, aborted transaction) - should trigger abort
+ // because it doesn't match the current prepared state (marker2)
+ producer2.completeTransaction(marker1)
+
+ // Wait for transaction to reach COMPLETE_ABORT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_ABORT)
+
+ // Verify consumer sees NO records (both transactions aborted)
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(0, consumedRecords.size,
+ "Consumer should see no records after abort")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_ABORT, txnDescription.state(),
+ "Transaction should be in COMPLETE_ABORT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersAll"))
+ def testProducerCrashAndRecoverWith2PC(groupProtocol: String): Unit = {
+ // Test producer crash and recovery with 2PC keepPrepared transaction flow.
+ // Note: This uses a standard transactional producer. 2PC is enabled server-side
+ // and triggered by calling initTransactions(keepPreparedTxn=true) after crash.
+
+ def test2PCRecovery(numCrashes: Int, shouldCommit: Boolean): Unit = {
+ val transactionalId = s"test-2pc-recovery-${System.nanoTime()}"
+ val testTopic = s"test-2pc-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100)) // Trigger assignment
+
+ // Create producer and send records in a transaction
+ var producer = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer.initTransactions()
+ producer.beginTransaction()
+
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes,
+ s"value-$i".getBytes))
+ }
+ producer.flush()
+
+ // Verify records not visible to read_committed consumer
+ consumer.poll(Duration.ofMillis(100))
+ assertEquals(0, consumer.assignment().asScala.map(tp =>
+ consumer.position(tp)).sum, "Records should not be visible before commit")
+
+ // Simulate crash by closing without committing
+ producer.close(Duration.ZERO)
+
+ val baseEpoch: Short = 0
+
+ // Crash and recover numCrashes times
+ for (crashNum <- 1 to numCrashes) {
+ val recoveredProducer = createTransactionalProducer(transactionalId, enable2PC = true)
+ // Use keepPreparedTxn=true to preserve in-flight transaction
+ recoveredProducer.initTransactions(true)
+
+ // Crash the transaction coordinator to test resilience during 2PC recovery
+ crashAndRestartTransactionCoordinator(transactionalId)
+
+ // Verify dual identity after recovery
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+
+ // After crash recovery, server should set up dual identity
+ assertTrue(txnDescription.producerEpoch() >= baseEpoch,
+ s"Crash $crashNum: client epoch should be >= base epoch")
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
Review Comment:
Will do.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -1056,6 +935,548 @@ class TransactionsTest extends IntegrationTestHarness {
producer.abortTransaction()
}
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testFencingWithDualIdentity(groupProtocol: String): Unit = {
+ // Test that zombie producer is properly fenced when dual identity is established
+ // after crash recovery with 2PC keepPreparedTxn.
+ val transactionalId = "test-fencing-dual-identity"
+ val testTopic = s"test-fencing-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and send records (simulating crash before commit)
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+ // Simulate crash - don't commit (leave transaction prepared)
+
+ // Create new producer and recover with initTransactions(keepPreparedTxn=true)
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // This establishes dual identity
+
+ // Old (zombie) producer tries to abort - should get ProducerFencedException
+ assertThrows(classOf[ProducerFencedException], () => {
+ producer1.abortTransaction()
+ }, "Zombie producer should be fenced when trying to abort")
+
+ // New producer commits the prepared transaction
+ producer2.commitTransaction()
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records from the first producer
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records from first producer after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionCommit(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for commit path.
+ // This tests the client-side 2PC API where the producer explicitly calls prepareTransaction()
+ // and then completeTransaction() to commit.
+ val transactionalId = "test-prepare-complete-commit"
+ val testTopic = s"test-prepare-complete-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and produce records
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Prepare the transaction (moves to PREPARED state)
+ val preparedState = producer1.prepareTransaction()
+ assertNotNull(preparedState, "prepareTransaction should return prepared state")
+
+ // Simulate crash - don't call commit/abort, just leave in prepared state
+ // (In real scenario, producer would crash here)
+
+ // New producer recovers and completes the transaction
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete the transaction with the prepared state (should commit)
+ producer2.completeTransaction(preparedState)
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionAbort(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for abort path.
+ // This tests that completeTransaction() with a non-matching state triggers abort.
+ val transactionalId = "test-prepare-complete-abort"
+ val testTopic = s"test-prepare-complete-abort-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+
+ // First transaction - prepare but then abort
+ producer1.beginTransaction()
+ for (i <- 0 until 3) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key1-$i".getBytes, s"value1-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for first transaction
+ val marker1 = producer1.prepareTransaction()
+ assertNotNull(marker1, "prepareTransaction should return prepared state")
+
+ // Abort the first prepared transaction
+ producer1.abortTransaction()
+
+ // Second transaction - prepare and leave in prepared state
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key2-$i".getBytes, s"value2-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for second transaction
+ val marker2 = producer1.prepareTransaction()
+ assertNotNull(marker2, "prepareTransaction should return prepared state")
+
+ // Simulate crash - leave second transaction in prepared state
+
+ // New producer recovers
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete with marker1 (old, aborted transaction) - should trigger abort
+ // because it doesn't match the current prepared state (marker2)
+ producer2.completeTransaction(marker1)
+
+ // Wait for transaction to reach COMPLETE_ABORT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_ABORT)
+
+ // Verify consumer sees NO records (both transactions aborted)
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(0, consumedRecords.size,
+ "Consumer should see no records after abort")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_ABORT, txnDescription.state(),
+ "Transaction should be in COMPLETE_ABORT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersAll"))
+ def testProducerCrashAndRecoverWith2PC(groupProtocol: String): Unit = {
+ // Test producer crash and recovery with 2PC keepPrepared transaction flow.
+ // Note: This uses a standard transactional producer. 2PC is enabled server-side
+ // and triggered by calling initTransactions(keepPreparedTxn=true) after crash.
+
+ def test2PCRecovery(numCrashes: Int, shouldCommit: Boolean): Unit = {
+ val transactionalId = s"test-2pc-recovery-${System.nanoTime()}"
+ val testTopic = s"test-2pc-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100)) // Trigger assignment
+
+ // Create producer and send records in a transaction
+ var producer = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer.initTransactions()
+ producer.beginTransaction()
+
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes,
+ s"value-$i".getBytes))
+ }
+ producer.flush()
+
+ // Verify records not visible to read_committed consumer
+ consumer.poll(Duration.ofMillis(100))
+ assertEquals(0, consumer.assignment().asScala.map(tp =>
+ consumer.position(tp)).sum, "Records should not be visible before commit")
+
+ // Simulate crash by closing without committing
+ producer.close(Duration.ZERO)
+
+ val baseEpoch: Short = 0
+
+ // Crash and recover numCrashes times
+ for (crashNum <- 1 to numCrashes) {
+ val recoveredProducer = createTransactionalProducer(transactionalId, enable2PC = true)
+ // Use keepPreparedTxn=true to preserve in-flight transaction
+ recoveredProducer.initTransactions(true)
+
+ // Crash the transaction coordinator to test resilience during 2PC recovery
+ crashAndRestartTransactionCoordinator(transactionalId)
+
+ // Verify dual identity after recovery
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+
+ // After crash recovery, server should set up dual identity
+ assertTrue(txnDescription.producerEpoch() >= baseEpoch,
+ s"Crash $crashNum: client epoch should be >= base epoch")
+
+ // For simplicity in this test, we just verify epoch progression.
+ // Detailed dual identity verification is in unit tests.
+
+ if (crashNum < numCrashes) {
+ // Simulate another crash
+ recoveredProducer.close(Duration.ZERO)
+ } else {
+ // Last recovery - complete the transaction
+ producer = recoveredProducer
+ }
+ }
+
+ // Complete prepared transaction directly. Cannot call beginTransaction()
+ // in prepared state - must call commitTransaction() or abortTransaction() directly.
+ if (shouldCommit) {
+ producer.commitTransaction()
+ } else {
+ producer.abortTransaction()
+ }
+
+ // Wait for the transaction to fully complete on the server before proceeding.
+ // The client-side commitTransaction() returns after receiving the EndTxn response,
+ // but the server might still be writing markers and transitioning to COMPLETE_* state.
+ val expectedState = if (shouldCommit) TransactionState.COMPLETE_COMMIT else TransactionState.COMPLETE_ABORT
+ waitForTransactionState(transactionalId, expectedState)
+
+ // Verify consumer sees correct records
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+
+ if (shouldCommit) {
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records after commit")
+ } else {
+ assertEquals(0, consumedRecords.size,
+ "Consumer should see no records after abort")
+ }
+
+ // Verify fresh transaction works with bumped epoch
+ producer.beginTransaction()
+ producer.send(new ProducerRecord(testTopic, 0, "fresh-key".getBytes,
+ "fresh-value".getBytes))
+ producer.commitTransaction()
+
+ // Seek to beginning to read all records (both 2PC and fresh)
+ consumer.seekToBeginning(consumer.assignment())
+ val allRecords = consumeRecordsFor(consumer)
+ val expectedTotal = if (shouldCommit) numRecords + 1 else 1
+ assertEquals(expectedTotal, allRecords.size,
+ s"Should have $expectedTotal total records after fresh transaction")
+
+ producer.close()
+
+ // Unsubscribe consumer for next test iteration
+ consumer.unsubscribe()
+ }
+
+ // Test single crash with commit
+ test2PCRecovery(numCrashes = 1, shouldCommit = true)
+
+ // Test single crash with abort
+ test2PCRecovery(numCrashes = 1, shouldCommit = false)
+
+ // Test multiple crashes before final commit (validates epoch progression)
+ test2PCRecovery(numCrashes = 3, shouldCommit = true)
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testProducerIdRotationWithEpochExhaustion(groupProtocol: String): Unit = {
+ // Test producer ID rotation when client epoch is exhausted. This tests both scenarios:
+ // 1. Rotation during initTransactions(keepPreparedTxn=true)
+ // 2. Rotation during commitTransaction()
+
+ def testRotation(startEpoch: Short, doubleRotation: Boolean = false): Unit = {
+ val transactionalId = s"test-rotation-${System.nanoTime()}"
+ val testTopic = s"test-rotation-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Establish transactional ID
+ var producer = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer.initTransactions()
+ producer.close()
+
+ // Set epoch to trigger rotation at desired point
+ setProducerEpoch(transactionalId, startEpoch)
+
+ // Create producer and start a prepared transaction
+ producer = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer.initTransactions()
+ producer.beginTransaction()
+ val numRecords = 3
+ for (i <- 0 until numRecords) {
+ producer.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer.flush()
+ // Don't commit - leave transaction prepared (simulates crash)
+
+ val originalProducerId = getProducerId(transactionalId)
+
+ if (doubleRotation) {
+ // First rotation: loop calling initTransactions(true) until rotation occurs
+ var rotationCount = 0
+ var currentClientId = originalProducerId
+ var iteration = 0
+ while (rotationCount == 0 && iteration < 20) {
+ iteration += 1
+ val recoveryProducer = createTransactionalProducer(transactionalId, enable2PC = true)
+ recoveryProducer.initTransactions(true) // keepPreparedTxn
+
+ val clientId = getClientProducerId(transactionalId)
+ if (clientId != currentClientId) {
+ rotationCount = 1
+ currentClientId = clientId
+ }
+ recoveryProducer.close(Duration.ZERO)
+ }
+
+ assertTrue(rotationCount >= 1, s"First rotation should have occurred after $iteration iterations")
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
+ case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
+ retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ }
+ } else {
+ // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
+ (!isFromClient || producerEpoch == txnMetadata.clientProducerEpoch) && producerEpoch >= txnMetadata.clientProducerEpoch
}
- } else {
- // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
- (!isFromClient || producerEpoch == txnMetadata.producerEpoch) && producerEpoch >= txnMetadata.producerEpoch
}
- }
- val isRetry = retryOnEpochBump || retryOnOverflow
-
- def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
- // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
- // This must also apply to the epoch-fence path.
- val nextProducerIdOrErrors =
- if (txnMetadata.isProducerEpochExhausted) {
- try {
- Right(producerIdManager.generateProducerId())
- } catch {
- case e: Exception => Left(Errors.forException(e))
+ val isRetry = retryOnEpochBump || retryOnOverflow
+
+ def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
+ // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
+ // This must also apply to the epoch-fence path.
+ val nextProducerIdOrErrors =
+ if (txnMetadata.isProducerEpochExhausted) {
+ try {
+ Right(producerIdManager.generateProducerId())
+ } catch {
+ case e: Exception => Left(Errors.forException(e))
+ }
+ } else {
+ Right(txnMetadata.nextProducerId)
}
+
+ // If the next producer epoch is set (which can happen if we used InitProducerId(keepPreparedTxn))
+ // then we also need to bump this epoch -- we want to fence requests that were issued with that
+ // epoch after the transaction is complete. Consider this example:
+ // 1. (initial state): Ongoing, producerId = 42, producerEpoch = 100, nextProducerId = -1, nextProducerEpoch = -1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, producerId = 42, producerEpoch = 100, nextProducerId = 42, nextProducerEpoch = 101
+ // 3. CommitTxn: timed out (but actually delayed)
+ // 4. CommitTxn + complete: CompleteCommit, producerId = 42, producerEpoch = 101, nextProducerId = -1, nextProducerEpoch = -1
+ // 5. Start new transaction, do some actions
+ // 6. Delayed CommitTxn from step 3 arrives and commits partially completed transaction.
+ // In order to prevent this problem, we need to bump the epoch at step 4 (so it becomes 102) even though we
+ // already bumped it at step 2.
+ // Note that we can arrive at this scenario when a client did a few keep-prepared calls and then
+ // decided to do one without keep-prepared (e.g. to force-terminate), so this logic applies
+ // to isEpochFence as well.
+ val nextProducerEpoch = if (txnMetadata.hasNextProducerEpoch) {
+ if (txnMetadata.isProducerEpochExhausted) 0.toShort else (txnMetadata.nextProducerEpoch + 1).toShort
} else {
- Right(RecordBatch.NO_PRODUCER_ID)
+ RecordBatch.NO_PRODUCER_EPOCH
}
- if (nextState == TransactionState.PREPARE_ABORT && isEpochFence) {
- // We should clear the pending state to make way for the transition to PrepareAbort
- txnMetadata.pendingState(util.Optional.empty())
- // For TV2+, don't manually set the epoch - let prepareAbortOrCommit handle it naturally.
- }
+ if (nextState == TransactionState.PREPARE_ABORT && isEpochFence) {
+ // We should clear the pending state to make way for the transition to PrepareAbort
+ txnMetadata.pendingState(util.Optional.empty())
+ // For TV2+, don't manually set the epoch - let prepareAbortOrCommit handle it naturally.
+ }
- nextProducerIdOrErrors.flatMap {
- nextProducerId =>
- Right(coordinatorEpoch, txnMetadata.prepareAbortOrCommit(nextState, clientTransactionVersion, nextProducerId.asInstanceOf[Long], time.milliseconds(), noPartitionAdded))
+ nextProducerIdOrErrors.flatMap {
+ nextProducerId =>
+ Right(coordinatorEpoch, txnMetadata.prepareAbortOrCommit(nextState, clientTransactionVersion, nextProducerId, nextProducerEpoch, time.milliseconds(), noPartitionAdded))
+ }
}
- }
-
- if (txnMetadata.pendingTransitionInProgress && txnMetadata.pendingState.get != TransactionState.PREPARE_EPOCH_FENCE) {
- // This check is performed first so that the pending transition can complete before the next checks.
- // With TV2, we may be transitioning over a producer epoch overflow, and the producer may be using the
- // new producer ID that is still only in pending state.
- Left(Errors.CONCURRENT_TRANSACTIONS)
- } else if (txnMetadata.producerId != producerId && !retryOnOverflow)
- Left(Errors.INVALID_PRODUCER_ID_MAPPING)
- else if (!isValidEpoch)
- Left(Errors.PRODUCER_FENCED)
- else txnMetadata.state match {
- case TransactionState.ONGOING =>
- val nextState = if (txnMarkerResult == TransactionResult.COMMIT)
- TransactionState.PREPARE_COMMIT
- else
- TransactionState.PREPARE_ABORT
- generateTxnTransitMetadataForTxnCompletion(nextState, false)
- case TransactionState.COMPLETE_COMMIT =>
- if (txnMarkerResult == TransactionResult.COMMIT) {
- if (isRetry)
- Left(Errors.NONE)
+ if (txnMetadata.clientProducerId != producerId && !retryOnOverflow)
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
+ case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
+ retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
Review Comment:
I don't think we should add this. prevProducerId is set on every transition (to the producerId before the transition), so in COMPLETE_COMMIT/COMPLETE_ABORT it is effectively never -1, and the extra condition would make the 'Current' column of the TV2 table always fail. E.g. CompleteCommit + Abort at the current epoch (an empty abort after a completed transaction) must bump the epoch (EB), and CompleteCommit + Commit at the current epoch must return INVALID_TXN_STATE; with the change both would return PRODUCER_FENCED. `producerEpoch == txnMetadata.producerEpoch` is exactly the 'Current' case, and the state match below then decides between EB and ITS.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
+ case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
+ retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ }
+ } else {
+ // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
+ (!isFromClient || producerEpoch == txnMetadata.clientProducerEpoch) && producerEpoch >= txnMetadata.clientProducerEpoch
}
- } else {
- // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
- (!isFromClient || producerEpoch == txnMetadata.producerEpoch) && producerEpoch >= txnMetadata.producerEpoch
}
- }
- val isRetry = retryOnEpochBump || retryOnOverflow
-
- def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
- // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
- // This must also apply to the epoch-fence path.
- val nextProducerIdOrErrors =
- if (txnMetadata.isProducerEpochExhausted) {
- try {
- Right(producerIdManager.generateProducerId())
- } catch {
- case e: Exception => Left(Errors.forException(e))
+ val isRetry = retryOnEpochBump || retryOnOverflow
+
+ def generateTxnTransitMetadataForTxnCompletion(nextState: TransactionState, noPartitionAdded: Boolean): ApiResult[(Int, TxnTransitMetadata)] = {
+ // EndTxn completion on TV2 bumps epoch, so rotate producer ID whenever the current epoch is exhausted.
+ // This must also apply to the epoch-fence path.
+ val nextProducerIdOrErrors =
+ if (txnMetadata.isProducerEpochExhausted) {
+ try {
+ Right(producerIdManager.generateProducerId())
+ } catch {
+ case e: Exception => Left(Errors.forException(e))
+ }
+ } else {
+ Right(txnMetadata.nextProducerId)
}
+
+ // If the next producer epoch is set (which can happen if we used InitProducerId(keepPreparedTxn))
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -927,9 +987,14 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
val txnMetadata = epochAndMetadata.transactionMetadata
txnMetadata.inLock(() => {
- if (txnMetadata.producerId != producerId)
+ // Note that if we did a few InitProducerId(keepPreparedTxn) the client's producer id
+ // may be different from the producerId (we track it in the nextProducerId), but at
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
+ producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
+ // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
+ val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
+
+ val isValidEpoch = {
+ if (!isEpochFence) {
+ // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
+ // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
+ // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
+ // Use the following criteria to determine if a v2 retry is valid:
+ txnMetadata.state match {
+ case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
+ producerEpoch == txnMetadata.clientProducerEpoch
+ case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
+ retryOnEpochBump
+ case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
+ retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ }
+ } else {
+ // If the epoch is going to be fenced, it bumps the epoch differently with TV2.
+ (!isFromClient || producerEpoch == txnMetadata.clientProducerEpoch) && producerEpoch >= txnMetadata.clientProducerEpoch
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -779,134 +808,165 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
val coordinatorEpoch = epochAndTxnMetadata.coordinatorEpoch
txnMetadata.inLock(() => {
- producerIdCopy = txnMetadata.producerId
- producerEpochCopy = txnMetadata.producerEpoch
- // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
- // Note that, it can only happen when the current state is Ongoing.
- isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
- // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
- val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
- producerEpoch == Short.MaxValue - 1 && txnMetadata.producerEpoch == 0
- // True if the client retried an endTxn request, and the bumped producer epoch is stored in the txnMetadata.
- val retryOnEpochBump = !isEpochFence && txnMetadata.producerEpoch == producerEpoch + 1
-
- val isValidEpoch = {
- if (!isEpochFence) {
- // With transactions V2, state + same epoch is not sufficient to determine if a retry transition is valid. If the epoch is the
- // same it actually indicates the next endTransaction call. Instead, we want to check the epoch matches with the epoch in the retry conditions.
- // Return producer fenced even in the cases where the epoch is higher and could indicate an invalid state transition.
- // Use the following criteria to determine if a v2 retry is valid:
- txnMetadata.state match {
- case TransactionState.ONGOING | TransactionState.EMPTY | TransactionState.DEAD | TransactionState.PREPARE_EPOCH_FENCE =>
- producerEpoch == txnMetadata.producerEpoch
- case TransactionState.PREPARE_COMMIT | TransactionState.PREPARE_ABORT =>
- retryOnEpochBump
- case TransactionState.COMPLETE_COMMIT | TransactionState.COMPLETE_ABORT =>
- retryOnEpochBump || retryOnOverflow || producerEpoch == txnMetadata.producerEpoch
+ // These checks are performed first so that we don't have to deal with intermediate states
+ // (including PrepareCommit and PrepareAbort) in already complex state machine.
+ if (txnMetadata.pendingTransitionInProgress && !txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else if (txnMetadata.state == TransactionState.PREPARE_COMMIT || txnMetadata.state == TransactionState.PREPARE_ABORT) {
+ Left(Errors.CONCURRENT_TRANSACTIONS)
+ } else {
+ // Copy the data under lock, this is going to be used if we want to return success
+ // on retrying commit / abort.
+ producerIdCopy = txnMetadata.producerId
+ producerEpochCopy = txnMetadata.producerEpoch
+
+ // PrepareEpochFence has slightly different epoch bumping logic so don't include it here.
+ // Note that, it can only happen when the current state is Ongoing.
+ isEpochFence = txnMetadata.pendingState.filter(s => s == TransactionState.PREPARE_EPOCH_FENCE).isPresent
+
+ // Calculate retry conditions, note that they are going to be used only
+ // in CompleteCommit / CompleteAbort states, so we can use the semantics
+ // of the producerId / epoch that's applicable in those states.
+ // Also note that the retry-on-overflow condition happens only when
+ // overflow happens during last commit / abort, overflows during
+ // InitProducerId(keepPreparedTxn=true) go through retry-on-epoch-bump
+ // path. Example:
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 85} -- the client-facing epoch has overflown
+ // 2. Commit transitions to {pid, epoch} = {73, 86} and the client retries with {73, 85} -- retryOnEpochBump
+ // Another example where we overflow during commit after going through overflow with keepPreparedTxn = true
+ // 1. State ONGOING, {pid, epoch} = {42, 10}, {nextPid, nextEpoch} = {73, 32766}
+ // 2. Commit transitions to {pid, epoch} = {85, 0}, prevPid = 73 and the client retries with {73, 32766} -- retryOnOverflow
+
+ // True if the client retried a request that had overflowed the epoch, and a new producer ID is stored in the txnMetadata
+ val retryOnOverflow = !isEpochFence && txnMetadata.prevProducerId == producerId &&
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -211,6 +237,31 @@ public TxnTransitMetadata prepareProducerIdRotation(long newProducerId,
return prepareTransitionTo(data);
}
+ public TxnTransitMetadata prepareIncrementProducerEpochForOngoing(long updateTimestamp) {
+ if (isProducerEpochExhausted())
Review Comment:
Will do.
##########
core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala:
##########
@@ -927,9 +987,14 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
val txnMetadata = epochAndMetadata.transactionMetadata
txnMetadata.inLock(() => {
- if (txnMetadata.producerId != producerId)
+ // Note that if we did a few InitProducerId(keepPreparedTxn) the client's producer id
+ // may be different from the producerId (we track it in the nextProducerId), but at
+ // this point even nextProducerId could be different from the client's producer id
+ // if we generated a new nextProducerId due to nextProducerEpoch overflow. So we
+ // check with the copy of the txnMetadata.producerId. Ditto for the epoch.
+ if (txnMetadata.producerId != producerIdCopy)
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -474,14 +531,24 @@ private boolean validProducerEpoch(TxnTransitMetadata transitMetadata) {
short transitLastProducerEpoch = transitMetadata.lastProducerEpoch();
if (isAtLeastTransactionsV2 &&
- (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT) &&
- transitProducerEpoch == 0) {
- return transitLastProducerEpoch == lastProducerEpoch && transitMetadata.prevProducerId() == producerId;
+ (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT)) {
+ if (transitProducerId != producerId) {
+ // Producer ID rotation case: validate that the transition metadata has the correct prevProducerId
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -474,14 +531,24 @@ private boolean validProducerEpoch(TxnTransitMetadata transitMetadata) {
short transitLastProducerEpoch = transitMetadata.lastProducerEpoch();
if (isAtLeastTransactionsV2 &&
- (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT) &&
- transitProducerEpoch == 0) {
- return transitLastProducerEpoch == lastProducerEpoch && transitMetadata.prevProducerId() == producerId;
+ (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT)) {
Review Comment:
Will do.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -1056,6 +935,548 @@ class TransactionsTest extends IntegrationTestHarness {
producer.abortTransaction()
}
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testFencingWithDualIdentity(groupProtocol: String): Unit = {
+ // Test that zombie producer is properly fenced when dual identity is established
+ // after crash recovery with 2PC keepPreparedTxn.
+ val transactionalId = "test-fencing-dual-identity"
+ val testTopic = s"test-fencing-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and send records (simulating crash before commit)
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+ // Simulate crash - don't commit (leave transaction prepared)
+
+ // Create new producer and recover with initTransactions(keepPreparedTxn=true)
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // This establishes dual identity
+
+ // Old (zombie) producer tries to abort - should get ProducerFencedException
+ assertThrows(classOf[ProducerFencedException], () => {
+ producer1.abortTransaction()
+ }, "Zombie producer should be fenced when trying to abort")
+
+ // New producer commits the prepared transaction
+ producer2.commitTransaction()
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records from the first producer
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records from first producer after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionCommit(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for commit path.
+ // This tests the client-side 2PC API where the producer explicitly calls prepareTransaction()
+ // and then completeTransaction() to commit.
+ val transactionalId = "test-prepare-complete-commit"
+ val testTopic = s"test-prepare-complete-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ // Start transaction and produce records
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes, s"value-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Prepare the transaction (moves to PREPARED state)
+ val preparedState = producer1.prepareTransaction()
+ assertNotNull(preparedState, "prepareTransaction should return prepared state")
+
+ // Simulate crash - don't call commit/abort, just leave in prepared state
+ // (In real scenario, producer would crash here)
+
+ // New producer recovers and completes the transaction
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete the transaction with the prepared state (should commit)
+ producer2.completeTransaction(preparedState)
+
+ // Wait for transaction to reach COMPLETE_COMMIT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_COMMIT)
+
+ // Verify consumer sees the records
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(numRecords, consumedRecords.size,
+ "Consumer should see all records after commit")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_COMMIT, txnDescription.state(),
+ "Transaction should be in COMPLETE_COMMIT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly"))
+ def testPrepareAndCompleteTransactionAbort(groupProtocol: String): Unit = {
+ // Test the 2PC API: prepareTransaction() and completeTransaction() for abort path.
+ // This tests that completeTransaction() with a non-matching state triggers abort.
+ val transactionalId = "test-prepare-complete-abort"
+ val testTopic = s"test-prepare-complete-abort-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100))
+
+ val producer1 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer1.initTransactions()
+
+ // First transaction - prepare but then abort
+ producer1.beginTransaction()
+ for (i <- 0 until 3) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key1-$i".getBytes, s"value1-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for first transaction
+ val marker1 = producer1.prepareTransaction()
+ assertNotNull(marker1, "prepareTransaction should return prepared state")
+
+ // Abort the first prepared transaction
+ producer1.abortTransaction()
+
+ // Second transaction - prepare and leave in prepared state
+ producer1.beginTransaction()
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer1.send(new ProducerRecord(testTopic, 0, s"key2-$i".getBytes, s"value2-$i".getBytes))
+ }
+ producer1.flush()
+
+ // Get marker for second transaction
+ val marker2 = producer1.prepareTransaction()
+ assertNotNull(marker2, "prepareTransaction should return prepared state")
+
+ // Simulate crash - leave second transaction in prepared state
+
+ // New producer recovers
+ val producer2 = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer2.initTransactions(true) // keepPreparedTxn=true
+
+ // Complete with marker1 (old, aborted transaction) - should trigger abort
+ // because it doesn't match the current prepared state (marker2)
+ producer2.completeTransaction(marker1)
+
+ // Wait for transaction to reach COMPLETE_ABORT state
+ waitForTransactionState(transactionalId, TransactionState.COMPLETE_ABORT)
+
+ // Verify consumer sees NO records (both transactions aborted)
+ consumer.seekToBeginning(consumer.assignment())
+ val consumedRecords = consumeRecordsFor(consumer)
+ assertEquals(0, consumedRecords.size,
+ "Consumer should see no records after abort")
+
+ // Verify transaction state
+ val txnDescription = adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get()
+ assertEquals(TransactionState.COMPLETE_ABORT, txnDescription.state(),
+ "Transaction should be in COMPLETE_ABORT state")
+
+ producer1.close()
+ producer2.close()
+ consumer.unsubscribe()
+ }
+
+ @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames)
+ @MethodSource(Array("getTestGroupProtocolParametersAll"))
+ def testProducerCrashAndRecoverWith2PC(groupProtocol: String): Unit = {
+ // Test producer crash and recovery with 2PC keepPrepared transaction flow.
+ // Note: This uses a standard transactional producer. 2PC is enabled server-side
+ // and triggered by calling initTransactions(keepPreparedTxn=true) after crash.
+
+ def test2PCRecovery(numCrashes: Int, shouldCommit: Boolean): Unit = {
+ val transactionalId = s"test-2pc-recovery-${System.nanoTime()}"
+ val testTopic = s"test-2pc-topic-${System.nanoTime()}"
+ createTopic(testTopic, 1, brokerCount, topicConfig())
+
+ val consumer = transactionalConsumers.head
+ consumer.subscribe(Seq(testTopic).asJava)
+ consumer.poll(Duration.ofMillis(100)) // Trigger assignment
+
+ // Create producer and send records in a transaction
+ var producer = createTransactionalProducer(transactionalId, enable2PC = true)
+ producer.initTransactions()
+ producer.beginTransaction()
+
+ val numRecords = 5
+ for (i <- 0 until numRecords) {
+ producer.send(new ProducerRecord(testTopic, 0, s"key-$i".getBytes,
+ s"value-$i".getBytes))
+ }
+ producer.flush()
+
+ // Verify records not visible to read_committed consumer
+ consumer.poll(Duration.ofMillis(100))
+ assertEquals(0, consumer.assignment().asScala.map(tp =>
+ consumer.position(tp)).sum, "Records should not be visible before commit")
+
+ // Simulate crash by closing without committing
+ producer.close(Duration.ZERO)
+
+ val baseEpoch: Short = 0
+
+ // Crash and recover numCrashes times
Review Comment:
Will do.
##########
transaction-coordinator/src/main/java/org/apache/kafka/coordinator/transaction/TransactionMetadata.java:
##########
@@ -474,14 +531,24 @@ private boolean validProducerEpoch(TxnTransitMetadata transitMetadata) {
short transitLastProducerEpoch = transitMetadata.lastProducerEpoch();
if (isAtLeastTransactionsV2 &&
- (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT) &&
- transitProducerEpoch == 0) {
- return transitLastProducerEpoch == lastProducerEpoch && transitMetadata.prevProducerId() == producerId;
+ (txnState == TransactionState.COMPLETE_COMMIT || txnState == TransactionState.COMPLETE_ABORT)) {
+ if (transitProducerId != producerId) {
+ // Producer ID rotation case: validate that the transition metadata has the correct prevProducerId
+ // and that the transitProducerEpoch matches the nextProducerEpoch (which may have been bumped during EndTransaction).
+ // If nextProducerEpoch is not set (NO_PRODUCER_EPOCH), then transitProducerEpoch should be 0 (standard rotation).
+ boolean epochMatches = hasNextProducerEpoch() ?
+ transitProducerEpoch == nextProducerEpoch : transitProducerEpoch == 0;
+ return transitLastProducerEpoch == lastProducerEpoch &&
+ transitMetadata.prevProducerId() == producerId &&
+ epochMatches;
+ } else {
+ return transitProducerEpoch == clientProducerEpoch() && transitProducerId == producerId;
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
Review Comment:
Will do.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -71,6 +72,11 @@ class TransactionsTest extends IntegrationTestHarness {
props.put(ReplicationConfigs.AUTO_LEADER_REBALANCE_ENABLE_CONFIG, false.toString)
props.put(GroupCoordinatorConfig.GROUP_INITIAL_REBALANCE_DELAY_MS_CONFIG, "0")
props.put(TransactionStateManagerConfig.TRANSACTIONS_ABORT_TIMED_OUT_TRANSACTION_CLEANUP_INTERVAL_MS_CONFIG, "200")
+ // Enable unstable API versions to support KIP-939 2PC (InitProducerId v6 with keepPreparedTxn)
+ props.put(ServerConfigs.UNSTABLE_API_VERSIONS_ENABLE_CONFIG, "true")
+ props.put(ServerConfigs.UNSTABLE_FEATURE_VERSIONS_ENABLE_CONFIG, "true")
+ // Enable 2PC support on the broker side
+ props.put(TransactionStateManagerConfig.TRANSACTIONS_2PC_ENABLED_CONFIG, "true")
Review Comment:
I don't think we should make it internal in this PR. transaction.two.phase.commit.enable is already a public broker config: it shipped in 4.1.0 and 4.2.0 and is documented in KIP-939. In practice 2PC is gated by the protocol: Enable2Pc and KeepPreparedTxn only exist in InitProducerId v6, and InitProducerIdRequest.json has latestVersionUnstable: true, so brokers don't advertise v6 unless unstable.api.versions.enable is set, and a client can't use 2PC regardless of this config. Also, after merging trunk, the 2PC integration tests moved to a separate TransactionsWith2PCTest class (TransactionsTest was rewritten in Java on trunk), so the unstable API/feature versions are no longer enabled for the rest of the transaction tests.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
Review Comment:
Will do.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java:
##########
@@ -1528,6 +1531,9 @@ public void handleResponse(AbstractResponse response) {
initProducerIdResponse.data().ongoingTxnProducerId(),
initProducerIdResponse.data().ongoingTxnProducerEpoch()
);
+ // Set transactionStarted to true because there is an ongoing transaction
+ // on the server that needs to be completed with EndTxn.
+ transactionStarted = true;
Review Comment:
I don't think that sequence can happen within a TransactionManager. initializeTransactions() transitions to INITIALIZING, which is only valid from UNINITIALIZED, COMMITTING_TRANSACTION or ABORTING_TRANSACTION (State.isTransitionValid), not from PREPARED_TRANSACTION. So after initTransactions(keepPreparedTxn=true) returns an ongoing transaction, the producer must commit or abort before anything else; the internal paths from COMMITTING/ABORTING back to INITIALIZING go through resetTransactionState(), which already sets transactionStarted = false. A new producer instance starts with transactionStarted = false. So the else branch doesn't need to reset it.
##########
core/src/test/scala/integration/kafka/api/TransactionsTest.scala:
##########
@@ -1143,4 +1566,167 @@ class TransactionsTest extends IntegrationTestHarness {
def maybeVerifyLocalLogStartOffsets(partitionStartOffsets: Map[TopicPartition, JLong]): Unit = {
// Non-tiered storage topic partition doesn't have local log start offset
}
+
+ /**
+ * Helper method to wait until a transaction reaches a specific state.
+ * Polls listTransactions API until the transaction with given ID reaches the expected state.
+ */
+ private def waitForTransactionState(transactionalId: String, expectedState: TransactionState): Unit = {
+ waitUntilTrue(() => {
+ val listResult = adminClient.listTransactions()
+ val txns = listResult.all().get().asScala
+ txns.exists(txn =>
+ txn.transactionalId() == transactionalId &&
+ txn.state() == expectedState
+ )
+ }, s"Transaction did not reach $expectedState state")
+ }
+
+ /**
+ * Helper method to access transaction metadata for a given transactional ID.
+ * Finds the transaction coordinator and executes the provided function on the transaction metadata.
+ */
+ private def withTransactionMetadata[T](transactionalId: String)(f: org.apache.kafka.coordinator.transaction.TransactionMetadata => T): T = {
+ val coordinatorId = findTransactionCoordinatorId(transactionalId)
+ val coordinatorBroker = brokers.find(_.config.brokerId == coordinatorId).get
+ val txnCoordinator = coordinatorBroker.asInstanceOf[kafka.server.BrokerServer].transactionCoordinator
+
+ var result: Option[T] = None
+ txnCoordinator.transactionManager.getTransactionState(transactionalId).foreach { txnMetadataOpt =>
+ txnMetadataOpt.foreach { epochAndMetadata =>
+ result = Some(f(epochAndMetadata.transactionMetadata))
+ }
+ }
+ result.getOrElse(throw new IllegalStateException(s"Transaction metadata not found for $transactionalId"))
+ }
+
+ /**
+ * Helper method to manually set producer epoch to a high value for testing epoch exhaustion.
+ */
+ private def setProducerEpoch(transactionalId: String, epoch: Short): Unit = {
+ withTransactionMetadata(transactionalId) { metadata =>
+ metadata.inLock(() => {
+ metadata.setProducerEpoch(epoch)
+ null // inLock expects a Supplier that returns a value
+ })
+ }
+ }
+
+ /**
+ * Helper method to manually set client producer epoch (nextProducerEpoch) for testing 2PC epoch exhaustion.
+ * This is used when testing epoch exhaustion during prepared transactions where client credentials
+ * are stored in nextProducerId/nextProducerEpoch.
+ */
+ private def setClientProducerEpoch(transactionalId: String, epoch: Short): Unit = {
+ withTransactionMetadata(transactionalId) { metadata =>
+ metadata.inLock(() => {
+ metadata.setNextProducerEpoch(epoch)
+ null // inLock expects a Supplier that returns a value
+ })
+ }
+ }
+
+ /**
+ * Helper method to get current producer ID from DescribeTransactions API.
+ */
+ private def getProducerId(transactionalId: String): Long = {
+ adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get().producerId()
+ }
+
+ /**
+ * Helper method to get current producer epoch from DescribeTransactions API.
+ */
+ private def getProducerEpoch(transactionalId: String): Short = {
+ adminClient.describeTransactions(java.util.List.of(transactionalId))
+ .description(transactionalId).get().producerEpoch().toShort
+ }
+
+ /**
+ * Helper method to get client-facing producer ID (for 2PC dual identity scenarios).
+ * When a transaction is prepared, the ongoing transaction keeps its original ID/epoch,
+ * but the client gets new credentials (nextProducerId/nextProducerEpoch).
+ */
+ private def getClientProducerId(transactionalId: String): Long = {
+ withTransactionMetadata(transactionalId)(_.clientProducerId())
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -87,6 +87,76 @@ class TransactionCoordinatorTest {
.thenReturn(true)
}
+ /**
+ * Sets up an ONGOING 2PC transaction, configures mocks,
+ * and calls InitProducerId(keepPreparedTxn=true) to bump client epoch.
+ * Returns the transit metadata captor for verifying subsequent operations.
+ */
+ private def setupPrepared2PcTxnWithBumpedClientEpoch(initialEpoch: Short = producerEpoch, testProducerId: Long = producerId): ArgumentCaptor[TxnTransitMetadata] = {
+ // Setup: ONGOING transaction with no next producer epoch set
+ // Use Integer.MAX_VALUE for timeout to indicate this is a distributed 2PC transaction
+ val txnMetadata = new TransactionMetadata(
+ transactionalId,
+ testProducerId, testProducerId, RecordBatch.NO_PRODUCER_ID,
+ initialEpoch, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_PRODUCER_EPOCH,
+ Integer.MAX_VALUE /*2PC*/, TransactionState.ONGOING, partitions,
+ time.milliseconds(), time.milliseconds(), TV_2
+ )
+ // Verify this is correctly identified as a distributed 2PC transaction
+ assertTrue(txnMetadata.isDistributedTwoPhaseCommitTxn)
+
+ // Add partitions so the transaction has data
+ txnMetadata.addPartitions(partitions)
+
+ // Configure mocks
+ when(transactionManager.validateTransactionTimeoutMs(anyBoolean(), anyInt()))
+ .thenReturn(true)
+ when(transactionManager.isTransaction2pcEnabled())
+ .thenReturn(true)
+ when(transactionManager.getTransactionState(ArgumentMatchers.eq(transactionalId)))
+ .thenReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+ when(transactionManager.transactionVersionLevel()).thenReturn(TV_2)
+
+ val capturedTransitMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ when(transactionManager.appendTransactionToLog(
+ ArgumentMatchers.eq(transactionalId),
+ ArgumentMatchers.eq(coordinatorEpoch),
+ capturedTransitMetadata.capture(),
+ capturedErrorsCallback.capture(),
+ any(),
+ any()
+ )).thenAnswer(invocation => {
+ val metadata = invocation.getArgument[TxnTransitMetadata](2)
+ txnMetadata.completeTransitionTo(metadata)
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
+ // This represents the COMPLETE_* transition metadata created by the internal prepareComplete() call
+ val capturedMarkerMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ verify(transactionMarkerChannelManager).addTxnMarkersToSend(
+ ArgumentMatchers.eq(coordinatorEpoch),
+ ArgumentMatchers.eq(txnResult),
+ any[TransactionMetadata],
+ capturedMarkerMetadata.capture()
+ )
+
+ // Verify successful completion
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to expected state.
+ // Index 0: helper's InitProducerId
+ // Index 1: test's InitProducerId (rotation)
+ // Index 2: EndTransaction
+ val prepareMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, prepareMetadata.txnState)
+
+ // Verify that transaction markers are sent with original (not rotated)
+ // producerId and startEpoch + 1. The ongoing transaction maintains the
+ // original producerId/epoch, so markers use those credentials.
+ assertEquals(testProducerId, prepareMetadata.producerId,
+ "Markers should use original producerId (ongoing transaction)")
+ assertEquals((startEpoch + 1).toShort, prepareMetadata.producerEpoch,
+ "Markers should use startEpoch + 1 (bumped from ongoing transaction epoch)")
+
+ // Simulate markers being sent and transaction completing to
+ // COMPLETE_COMMIT/COMPLETE_ABORT. At this point, handleEndTransaction
+ // has already internally called prepareComplete() in the
+ // sendTxnMarkersCallback, so the metadata already has
+ // pendingState=COMPLETE_COMMIT/COMPLETE_ABORT. We just need to write
+ // the transition to the log to complete it.
+
+ // Use the TxnTransitMetadata created by the internal prepareComplete()
+ val txnTransitMetadata = capturedMarkerMetadata.getValue
+
+ // Append COMPLETE_* transition to log (mock calls completeTransitionTo)
+ transactionManager.appendTransactionToLog(
+ transactionalId,
+ coordinatorEpoch,
+ txnTransitMetadata,
+ _ => {}, // No-op callback since we're just completing the transition
+ _ => false, // Don't retry on error
+ null
+ )
+
+ // Verify that rotation happened (regardless of whether it occurred during
+ // InitProducerId or EndTransaction). After transaction completes, the final
+ // state should have:
+ // - A new producerId (different from the original testProducerId)
+ // - Epoch determined by when rotation occurred (see scenarios below)
+ // - prevProducerId set to the original testProducerId for retry detection
+ assertNotEquals(testProducerId, txnTransitMetadata.producerId,
+ "Producer ID should have rotated by transaction completion")
+
+ // Verify epoch overflow occurred and total epoch increments.
+ // We do 3 total epoch increments in this test:
+ // 1. In setupPrepared2PcTxnWithBumpedClientEpoch: startEpoch → startEpoch + 1
+ // 2. Second InitProducerId call: startEpoch + 1 → startEpoch + 2 (or rotation)
+ // 3. EndTransaction: bumps nextProducerEpoch again (post-rotation or triggers rotation)
+ // Since rotation happens when trying to exceed MaxValue - 1 (32766), the final
+ // epoch wraps around.
+
+ // Assert 1: Final epoch is less than startEpoch (overflow occurred)
+ assertTrue(txnTransitMetadata.producerEpoch < startEpoch,
+ s"Overflow should have occurred: finalEpoch=${txnTransitMetadata.producerEpoch} < startEpoch=$startEpoch")
+
+ // Assert 2: Total epoch increments = 3 (accounting for overflow)
+ // Formula: finalEpoch + (MaxValue - startEpoch) = total increments
+ // This accounts for increments from startEpoch to MaxValue boundary, then
+ // wrap to 0 and increments to finalEpoch.
+ val totalIncrements = txnTransitMetadata.producerEpoch + Short.MaxValue - startEpoch
+ assertEquals(3, totalIncrements,
+ s"Should have 3 total epoch increments: finalEpoch=${txnTransitMetadata.producerEpoch} + " +
+ s"(MaxValue=${Short.MaxValue} - startEpoch=$startEpoch) = $totalIncrements")
+
+ assertEquals(testProducerId, txnTransitMetadata.prevProducerId,
+ "prevProducerId should track original ID for retry detection")
+
+ // Retry EndTransaction with same client credentials - should succeed as a retry
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use same client producer ID from InitProducerId result
+ clientProducerEpoch, // Use same client epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify retry succeeds with NONE
+ assertEquals(Errors.NONE, error, "Retry with client credentials should succeed")
+
+ // Reset mock to clear call history for next invocation
+ reset(transactionMarkerChannelManager)
+ }
+
+ // Test both commit and abort with rotation - use different producer IDs
+ // to avoid conflicts
+ mockPidGenerator() // Initialize PID generator
+
+ // Scenario 1: Rotation happens during InitProducerId call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=100, epoch=32765, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=100, epoch=32765,
+ // nextPid=100, nextEpoch=32766
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32766 triggers
+ // isEpochExhausted() → rotation. Ongoing, pid=100, epoch=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 4. CommitTxn/AbortTxn: bumps nextEpoch 0→1 to fence delayed requests
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=1, prevPid=100
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch, 100L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch, 200L)
+
+ // Scenario 2: Rotation happens during EndTransaction call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=300, epoch=32764, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32765
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32765 < 32766
+ // (not exhausted) → just epoch bump. Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32766
+ // 4. CommitTxn/AbortTxn: prepareAbortOrCommit bumps epoch 32764→32765,
+ // then generateTxnTransitMetadataForTxnCompletion sees
+ // isProducerEpochExhausted() → rotation. PrepareCommit, pid=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=0, prevPid=300
+ val startEpoch2 = (Short.MaxValue - 3).toShort // 32764
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch2, 300L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch2, 400L)
+ }
+
+ @Test
+ def shouldHandleProducerIdRotationPersistence(): Unit = {
+ // Test that producer ID rotation metadata is correctly persisted through state transitions.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction and make first InitProducerId call
+ // Helper will bump: 32765 → 32766, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Call InitProducerId again - this will exhaust at 32767 and rotate
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify rotation occurred
+ val rotatedProducerId = result.producerId
+ val rotatedEpoch = result.producerEpoch
+ assertNotEquals(producerId, rotatedProducerId, "Producer ID should have rotated")
+ assertEquals(0, rotatedEpoch, "Client epoch should reset to 0")
+
+ // Verify: Check captured metadata shows rotation with prevProducerId set
+ val rotationMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID should be stored in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch should be 0")
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId should be set to original ID for retry detection")
+
+ // Verify ongoing identity still uses original producer ID
+ assertEquals(producerId, rotationMetadata.producerId, "Ongoing transaction producerId unchanged")
+ assertEquals(startEpoch, rotationMetadata.producerEpoch, "Ongoing transaction epoch unchanged")
+
+ // Complete the transaction (commit)
+ coordinator.handleEndTransaction(
+ transactionalId,
+ rotatedProducerId, // Use client's producer ID (rotated)
+ rotatedEpoch, // Use client's epoch (0 after rotation)
+ TransactionResult.COMMIT,
+ TV_2,
+ endTxnCallback
+ )
+
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to PREPARE_COMMIT
+ val prepareCommitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState)
+
+ // Verify: The key aspect of persistence is that prevProducerId was
+ // correctly set during rotation. This ensures that after the transaction
+ // completes to EMPTY, retry detection will work correctly by recognizing
+ // requests from the old producer ID.
+ assertEquals(producerId, rotationMetadata.prevProducerId,
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
+ // This represents the COMPLETE_* transition metadata created by the internal prepareComplete() call
+ val capturedMarkerMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ verify(transactionMarkerChannelManager).addTxnMarkersToSend(
+ ArgumentMatchers.eq(coordinatorEpoch),
+ ArgumentMatchers.eq(txnResult),
+ any[TransactionMetadata],
+ capturedMarkerMetadata.capture()
+ )
+
+ // Verify successful completion
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to expected state.
+ // Index 0: helper's InitProducerId
+ // Index 1: test's InitProducerId (rotation)
+ // Index 2: EndTransaction
+ val prepareMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, prepareMetadata.txnState)
+
+ // Verify that transaction markers are sent with original (not rotated)
+ // producerId and startEpoch + 1. The ongoing transaction maintains the
+ // original producerId/epoch, so markers use those credentials.
+ assertEquals(testProducerId, prepareMetadata.producerId,
+ "Markers should use original producerId (ongoing transaction)")
+ assertEquals((startEpoch + 1).toShort, prepareMetadata.producerEpoch,
+ "Markers should use startEpoch + 1 (bumped from ongoing transaction epoch)")
+
+ // Simulate markers being sent and transaction completing to
+ // COMPLETE_COMMIT/COMPLETE_ABORT. At this point, handleEndTransaction
+ // has already internally called prepareComplete() in the
+ // sendTxnMarkersCallback, so the metadata already has
+ // pendingState=COMPLETE_COMMIT/COMPLETE_ABORT. We just need to write
+ // the transition to the log to complete it.
+
+ // Use the TxnTransitMetadata created by the internal prepareComplete()
+ val txnTransitMetadata = capturedMarkerMetadata.getValue
+
+ // Append COMPLETE_* transition to log (mock calls completeTransitionTo)
+ transactionManager.appendTransactionToLog(
+ transactionalId,
+ coordinatorEpoch,
+ txnTransitMetadata,
+ _ => {}, // No-op callback since we're just completing the transition
+ _ => false, // Don't retry on error
+ null
+ )
+
+ // Verify that rotation happened (regardless of whether it occurred during
+ // InitProducerId or EndTransaction). After transaction completes, the final
+ // state should have:
+ // - A new producerId (different from the original testProducerId)
+ // - Epoch determined by when rotation occurred (see scenarios below)
+ // - prevProducerId set to the original testProducerId for retry detection
+ assertNotEquals(testProducerId, txnTransitMetadata.producerId,
+ "Producer ID should have rotated by transaction completion")
+
+ // Verify epoch overflow occurred and total epoch increments.
+ // We do 3 total epoch increments in this test:
+ // 1. In setupPrepared2PcTxnWithBumpedClientEpoch: startEpoch → startEpoch + 1
+ // 2. Second InitProducerId call: startEpoch + 1 → startEpoch + 2 (or rotation)
+ // 3. EndTransaction: bumps nextProducerEpoch again (post-rotation or triggers rotation)
+ // Since rotation happens when trying to exceed MaxValue - 1 (32766), the final
+ // epoch wraps around.
+
+ // Assert 1: Final epoch is less than startEpoch (overflow occurred)
+ assertTrue(txnTransitMetadata.producerEpoch < startEpoch,
+ s"Overflow should have occurred: finalEpoch=${txnTransitMetadata.producerEpoch} < startEpoch=$startEpoch")
+
+ // Assert 2: Total epoch increments = 3 (accounting for overflow)
+ // Formula: finalEpoch + (MaxValue - startEpoch) = total increments
+ // This accounts for increments from startEpoch to MaxValue boundary, then
+ // wrap to 0 and increments to finalEpoch.
+ val totalIncrements = txnTransitMetadata.producerEpoch + Short.MaxValue - startEpoch
+ assertEquals(3, totalIncrements,
+ s"Should have 3 total epoch increments: finalEpoch=${txnTransitMetadata.producerEpoch} + " +
+ s"(MaxValue=${Short.MaxValue} - startEpoch=$startEpoch) = $totalIncrements")
+
+ assertEquals(testProducerId, txnTransitMetadata.prevProducerId,
+ "prevProducerId should track original ID for retry detection")
+
+ // Retry EndTransaction with same client credentials - should succeed as a retry
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use same client producer ID from InitProducerId result
+ clientProducerEpoch, // Use same client epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify retry succeeds with NONE
+ assertEquals(Errors.NONE, error, "Retry with client credentials should succeed")
+
+ // Reset mock to clear call history for next invocation
+ reset(transactionMarkerChannelManager)
+ }
+
+ // Test both commit and abort with rotation - use different producer IDs
+ // to avoid conflicts
+ mockPidGenerator() // Initialize PID generator
+
+ // Scenario 1: Rotation happens during InitProducerId call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=100, epoch=32765, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=100, epoch=32765,
+ // nextPid=100, nextEpoch=32766
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32766 triggers
+ // isEpochExhausted() → rotation. Ongoing, pid=100, epoch=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 4. CommitTxn/AbortTxn: bumps nextEpoch 0→1 to fence delayed requests
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=1, prevPid=100
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch, 100L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch, 200L)
+
+ // Scenario 2: Rotation happens during EndTransaction call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=300, epoch=32764, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32765
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32765 < 32766
+ // (not exhausted) → just epoch bump. Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32766
+ // 4. CommitTxn/AbortTxn: prepareAbortOrCommit bumps epoch 32764→32765,
+ // then generateTxnTransitMetadataForTxnCompletion sees
+ // isProducerEpochExhausted() → rotation. PrepareCommit, pid=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=0, prevPid=300
+ val startEpoch2 = (Short.MaxValue - 3).toShort // 32764
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch2, 300L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch2, 400L)
+ }
+
+ @Test
+ def shouldHandleProducerIdRotationPersistence(): Unit = {
+ // Test that producer ID rotation metadata is correctly persisted through state transitions.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction and make first InitProducerId call
+ // Helper will bump: 32765 → 32766, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Call InitProducerId again - this will exhaust at 32767 and rotate
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify rotation occurred
+ val rotatedProducerId = result.producerId
+ val rotatedEpoch = result.producerEpoch
+ assertNotEquals(producerId, rotatedProducerId, "Producer ID should have rotated")
+ assertEquals(0, rotatedEpoch, "Client epoch should reset to 0")
+
+ // Verify: Check captured metadata shows rotation with prevProducerId set
+ val rotationMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID should be stored in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch should be 0")
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId should be set to original ID for retry detection")
+
+ // Verify ongoing identity still uses original producer ID
+ assertEquals(producerId, rotationMetadata.producerId, "Ongoing transaction producerId unchanged")
+ assertEquals(startEpoch, rotationMetadata.producerEpoch, "Ongoing transaction epoch unchanged")
+
+ // Complete the transaction (commit)
+ coordinator.handleEndTransaction(
+ transactionalId,
+ rotatedProducerId, // Use client's producer ID (rotated)
+ rotatedEpoch, // Use client's epoch (0 after rotation)
+ TransactionResult.COMMIT,
+ TV_2,
+ endTxnCallback
+ )
+
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to PREPARE_COMMIT
+ val prepareCommitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState)
+
+ // Verify: The key aspect of persistence is that prevProducerId was
+ // correctly set during rotation. This ensures that after the transaction
+ // completes to EMPTY, retry detection will work correctly by recognizing
+ // requests from the old producer ID.
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId persisted for retry detection")
+
+ // Verify: The rotation information is persisted in the metadata. After
+ // this transaction completes, the next transaction will start with
+ // rotatedProducerId.
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID persisted in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch persisted as 0")
+
+ // Verify: The state transitions are captured correctly in metadata
+ assertEquals(3, capturedTransitMetadata.getAllValues.size(),
+ "Three state transitions captured")
+ assertEquals(TransactionState.ONGOING,
+ capturedTransitMetadata.getAllValues.get(0).txnState,
+ "First: InitProducerId → ONGOING")
+ assertEquals(TransactionState.ONGOING, rotationMetadata.txnState,
+ "Second: InitProducerId with rotation → ONGOING")
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState,
+ "Third: EndTransaction → PREPARE_COMMIT")
+ }
+
+ @Test
+ def shouldHandleConcurrentTransactionsInPrepareStatesTV2(): Unit = {
+ // Test that attempting EndTransaction while already in a PREPARE state
+ // returns CONCURRENT_TRANSACTIONS error in TV2. This error is retriable -
+ // the client should wait for the current transaction to complete and then
+ // retry.
+
+ def testConcurrentInPrepareState(prepareState: TransactionState, attemptedResult: TransactionResult): Unit = {
+ // Setup ONGOING transaction with 2PC enabled
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Determine transaction result to reach desired PREPARE state
+ val txnResult = if (prepareState == TransactionState.PREPARE_COMMIT) {
+ TransactionResult.COMMIT
+ } else {
+ TransactionResult.ABORT
+ }
+
+ // Call EndTransaction to move to PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Use client producer ID
+ result.producerEpoch, // Use client epoch
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify we're in the expected PREPARE state
+ assertEquals(Errors.NONE, error, "EndTransaction should succeed")
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(prepareState, transitMetadata.txnState,
+ s"Transaction should be in $prepareState state")
+
+ val stateBefore = transitMetadata.txnState
+
+ // Attempt another EndTransaction while still in PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Same client producer ID
+ result.producerEpoch, // Same client epoch
+ attemptedResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: TV2 returns CONCURRENT_TRANSACTIONS (retriable error)
+ assertEquals(Errors.CONCURRENT_TRANSACTIONS, error,
+ s"Should return CONCURRENT_TRANSACTIONS when attempting $attemptedResult while in $prepareState")
+
+ // Verify: State unchanged after rejection
+ assertEquals(stateBefore, capturedTransitMetadata.getAllValues.get(1).txnState,
+ "State should remain unchanged after CONCURRENT_TRANSACTIONS error")
+
+ // Verify: No additional state transitions captured (request was rejected)
+ assertEquals(2, capturedTransitMetadata.getAllValues.size(),
+ "No new state transition should be captured for rejected request")
+ }
+
+ // Test Matrix: All 4 combinations of PREPARE state × attempted operation.
+ // All should return CONCURRENT_TRANSACTIONS regardless of operation type.
+ testConcurrentInPrepareState(TransactionState.PREPARE_COMMIT, TransactionResult.COMMIT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_COMMIT, TransactionResult.ABORT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_ABORT, TransactionResult.ABORT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_ABORT, TransactionResult.COMMIT)
+ }
+
+ @Test
+ def shouldRejectKeepPreparedTxnWithExpectedEpoch(): Unit = {
+ // Test that keepPreparedTxn=true with expectedEpoch returns INVALID_REQUEST.
+ // This is defense-in-depth validation: properly implemented clients should
+ // never send both flags together, as keepPreparedTxn is for initial calls
+ // and expectedEpoch is for retries.
+
+ // Setup EMPTY transaction and enable 2PC
+ when(transactionManager.getTransactionState(ArgumentMatchers.eq(transactionalId)))
+ .thenReturn(Right(None))
+ when(transactionManager.isTransaction2pcEnabled())
+ .thenReturn(true)
+
+ // Attempt InitProducerId with both keepPreparedTxn and expectedEpoch
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ expectedProducerIdAndEpoch = Some(new ProducerIdAndEpoch(producerId, producerEpoch)),
+ initProducerIdMockCallback
+ )
+
+ // Verify: Returns INVALID_REQUEST
+ assertEquals(Errors.INVALID_REQUEST, result.error,
+ "Should return INVALID_REQUEST when both keepPreparedTxn and expectedEpoch are set")
+ }
+
+ @Test
+ def shouldRejectKeepPreparedTxnWhenTransactionVersionTooLow(): Unit = {
+ // Test that keepPreparedTxn=true is rejected when transaction.version < TV_2.
+ // keepPreparedTxn requires TV_2 to ensure the 2PC metadata (nextProducerEpoch,
+ // prevProducerId, nextProducerId) is properly persisted to the transaction log.
+ // Without TV_2, these fields would be omitted during serialization, causing
+ // silent data corruption on coordinator recovery.
+
+ // Setup: ONGOING transaction at TV_0
+ val txnMetadata = new TransactionMetadata(
+ transactionalId,
+ producerId,
+ RecordBatch.NO_PRODUCER_ID,
+ RecordBatch.NO_PRODUCER_ID,
+ producerEpoch,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ Integer.MAX_VALUE, // 2PC transaction timeout
+ TransactionState.ONGOING,
+ partitions,
+ time.milliseconds(),
+ time.milliseconds(),
+ TV_0 // Transaction version 0 - doesn't support epoch bump
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
+ // This represents the COMPLETE_* transition metadata created by the internal prepareComplete() call
+ val capturedMarkerMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ verify(transactionMarkerChannelManager).addTxnMarkersToSend(
+ ArgumentMatchers.eq(coordinatorEpoch),
+ ArgumentMatchers.eq(txnResult),
+ any[TransactionMetadata],
+ capturedMarkerMetadata.capture()
+ )
+
+ // Verify successful completion
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to expected state.
+ // Index 0: helper's InitProducerId
+ // Index 1: test's InitProducerId (rotation)
+ // Index 2: EndTransaction
+ val prepareMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, prepareMetadata.txnState)
+
+ // Verify that transaction markers are sent with original (not rotated)
+ // producerId and startEpoch + 1. The ongoing transaction maintains the
+ // original producerId/epoch, so markers use those credentials.
+ assertEquals(testProducerId, prepareMetadata.producerId,
+ "Markers should use original producerId (ongoing transaction)")
+ assertEquals((startEpoch + 1).toShort, prepareMetadata.producerEpoch,
+ "Markers should use startEpoch + 1 (bumped from ongoing transaction epoch)")
+
+ // Simulate markers being sent and transaction completing to
+ // COMPLETE_COMMIT/COMPLETE_ABORT. At this point, handleEndTransaction
+ // has already internally called prepareComplete() in the
+ // sendTxnMarkersCallback, so the metadata already has
+ // pendingState=COMPLETE_COMMIT/COMPLETE_ABORT. We just need to write
+ // the transition to the log to complete it.
+
+ // Use the TxnTransitMetadata created by the internal prepareComplete()
+ val txnTransitMetadata = capturedMarkerMetadata.getValue
+
+ // Append COMPLETE_* transition to log (mock calls completeTransitionTo)
+ transactionManager.appendTransactionToLog(
+ transactionalId,
+ coordinatorEpoch,
+ txnTransitMetadata,
+ _ => {}, // No-op callback since we're just completing the transition
+ _ => false, // Don't retry on error
+ null
+ )
+
+ // Verify that rotation happened (regardless of whether it occurred during
+ // InitProducerId or EndTransaction). After transaction completes, the final
+ // state should have:
+ // - A new producerId (different from the original testProducerId)
+ // - Epoch determined by when rotation occurred (see scenarios below)
+ // - prevProducerId set to the original testProducerId for retry detection
+ assertNotEquals(testProducerId, txnTransitMetadata.producerId,
+ "Producer ID should have rotated by transaction completion")
+
+ // Verify epoch overflow occurred and total epoch increments.
+ // We do 3 total epoch increments in this test:
+ // 1. In setupPrepared2PcTxnWithBumpedClientEpoch: startEpoch → startEpoch + 1
+ // 2. Second InitProducerId call: startEpoch + 1 → startEpoch + 2 (or rotation)
+ // 3. EndTransaction: bumps nextProducerEpoch again (post-rotation or triggers rotation)
+ // Since rotation happens when trying to exceed MaxValue - 1 (32766), the final
+ // epoch wraps around.
+
+ // Assert 1: Final epoch is less than startEpoch (overflow occurred)
+ assertTrue(txnTransitMetadata.producerEpoch < startEpoch,
+ s"Overflow should have occurred: finalEpoch=${txnTransitMetadata.producerEpoch} < startEpoch=$startEpoch")
+
+ // Assert 2: Total epoch increments = 3 (accounting for overflow)
+ // Formula: finalEpoch + (MaxValue - startEpoch) = total increments
+ // This accounts for increments from startEpoch to MaxValue boundary, then
+ // wrap to 0 and increments to finalEpoch.
+ val totalIncrements = txnTransitMetadata.producerEpoch + Short.MaxValue - startEpoch
+ assertEquals(3, totalIncrements,
+ s"Should have 3 total epoch increments: finalEpoch=${txnTransitMetadata.producerEpoch} + " +
+ s"(MaxValue=${Short.MaxValue} - startEpoch=$startEpoch) = $totalIncrements")
+
+ assertEquals(testProducerId, txnTransitMetadata.prevProducerId,
+ "prevProducerId should track original ID for retry detection")
+
+ // Retry EndTransaction with same client credentials - should succeed as a retry
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use same client producer ID from InitProducerId result
+ clientProducerEpoch, // Use same client epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify retry succeeds with NONE
+ assertEquals(Errors.NONE, error, "Retry with client credentials should succeed")
+
+ // Reset mock to clear call history for next invocation
+ reset(transactionMarkerChannelManager)
+ }
+
+ // Test both commit and abort with rotation - use different producer IDs
+ // to avoid conflicts
+ mockPidGenerator() // Initialize PID generator
+
+ // Scenario 1: Rotation happens during InitProducerId call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=100, epoch=32765, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=100, epoch=32765,
+ // nextPid=100, nextEpoch=32766
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32766 triggers
+ // isEpochExhausted() → rotation. Ongoing, pid=100, epoch=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 4. CommitTxn/AbortTxn: bumps nextEpoch 0→1 to fence delayed requests
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=1, prevPid=100
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch, 100L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch, 200L)
+
+ // Scenario 2: Rotation happens during EndTransaction call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=300, epoch=32764, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32765
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32765 < 32766
+ // (not exhausted) → just epoch bump. Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32766
+ // 4. CommitTxn/AbortTxn: prepareAbortOrCommit bumps epoch 32764→32765,
+ // then generateTxnTransitMetadataForTxnCompletion sees
+ // isProducerEpochExhausted() → rotation. PrepareCommit, pid=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=0, prevPid=300
+ val startEpoch2 = (Short.MaxValue - 3).toShort // 32764
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch2, 300L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch2, 400L)
+ }
+
+ @Test
+ def shouldHandleProducerIdRotationPersistence(): Unit = {
+ // Test that producer ID rotation metadata is correctly persisted through state transitions.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction and make first InitProducerId call
+ // Helper will bump: 32765 → 32766, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Call InitProducerId again - this will exhaust at 32767 and rotate
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify rotation occurred
+ val rotatedProducerId = result.producerId
+ val rotatedEpoch = result.producerEpoch
+ assertNotEquals(producerId, rotatedProducerId, "Producer ID should have rotated")
+ assertEquals(0, rotatedEpoch, "Client epoch should reset to 0")
+
+ // Verify: Check captured metadata shows rotation with prevProducerId set
+ val rotationMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID should be stored in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch should be 0")
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId should be set to original ID for retry detection")
+
+ // Verify ongoing identity still uses original producer ID
+ assertEquals(producerId, rotationMetadata.producerId, "Ongoing transaction producerId unchanged")
+ assertEquals(startEpoch, rotationMetadata.producerEpoch, "Ongoing transaction epoch unchanged")
+
+ // Complete the transaction (commit)
+ coordinator.handleEndTransaction(
+ transactionalId,
+ rotatedProducerId, // Use client's producer ID (rotated)
+ rotatedEpoch, // Use client's epoch (0 after rotation)
+ TransactionResult.COMMIT,
+ TV_2,
+ endTxnCallback
+ )
+
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to PREPARE_COMMIT
+ val prepareCommitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState)
+
+ // Verify: The key aspect of persistence is that prevProducerId was
+ // correctly set during rotation. This ensures that after the transaction
+ // completes to EMPTY, retry detection will work correctly by recognizing
+ // requests from the old producer ID.
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId persisted for retry detection")
+
+ // Verify: The rotation information is persisted in the metadata. After
+ // this transaction completes, the next transaction will start with
+ // rotatedProducerId.
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID persisted in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch persisted as 0")
+
+ // Verify: The state transitions are captured correctly in metadata
+ assertEquals(3, capturedTransitMetadata.getAllValues.size(),
+ "Three state transitions captured")
+ assertEquals(TransactionState.ONGOING,
+ capturedTransitMetadata.getAllValues.get(0).txnState,
+ "First: InitProducerId → ONGOING")
+ assertEquals(TransactionState.ONGOING, rotationMetadata.txnState,
+ "Second: InitProducerId with rotation → ONGOING")
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState,
+ "Third: EndTransaction → PREPARE_COMMIT")
+ }
+
+ @Test
+ def shouldHandleConcurrentTransactionsInPrepareStatesTV2(): Unit = {
+ // Test that attempting EndTransaction while already in a PREPARE state
+ // returns CONCURRENT_TRANSACTIONS error in TV2. This error is retriable -
+ // the client should wait for the current transaction to complete and then
+ // retry.
+
+ def testConcurrentInPrepareState(prepareState: TransactionState, attemptedResult: TransactionResult): Unit = {
+ // Setup ONGOING transaction with 2PC enabled
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Determine transaction result to reach desired PREPARE state
+ val txnResult = if (prepareState == TransactionState.PREPARE_COMMIT) {
+ TransactionResult.COMMIT
+ } else {
+ TransactionResult.ABORT
+ }
+
+ // Call EndTransaction to move to PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Use client producer ID
+ result.producerEpoch, // Use client epoch
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify we're in the expected PREPARE state
+ assertEquals(Errors.NONE, error, "EndTransaction should succeed")
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(prepareState, transitMetadata.txnState,
+ s"Transaction should be in $prepareState state")
+
+ val stateBefore = transitMetadata.txnState
+
+ // Attempt another EndTransaction while still in PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Same client producer ID
+ result.producerEpoch, // Same client epoch
+ attemptedResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: TV2 returns CONCURRENT_TRANSACTIONS (retriable error)
+ assertEquals(Errors.CONCURRENT_TRANSACTIONS, error,
+ s"Should return CONCURRENT_TRANSACTIONS when attempting $attemptedResult while in $prepareState")
+
+ // Verify: State unchanged after rejection
+ assertEquals(stateBefore, capturedTransitMetadata.getAllValues.get(1).txnState,
+ "State should remain unchanged after CONCURRENT_TRANSACTIONS error")
+
+ // Verify: No additional state transitions captured (request was rejected)
+ assertEquals(2, capturedTransitMetadata.getAllValues.size(),
+ "No new state transition should be captured for rejected request")
+ }
+
+ // Test Matrix: All 4 combinations of PREPARE state × attempted operation.
+ // All should return CONCURRENT_TRANSACTIONS regardless of operation type.
+ testConcurrentInPrepareState(TransactionState.PREPARE_COMMIT, TransactionResult.COMMIT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_COMMIT, TransactionResult.ABORT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_ABORT, TransactionResult.ABORT)
+ testConcurrentInPrepareState(TransactionState.PREPARE_ABORT, TransactionResult.COMMIT)
+ }
+
+ @Test
+ def shouldRejectKeepPreparedTxnWithExpectedEpoch(): Unit = {
+ // Test that keepPreparedTxn=true with expectedEpoch returns INVALID_REQUEST.
+ // This is defense-in-depth validation: properly implemented clients should
+ // never send both flags together, as keepPreparedTxn is for initial calls
+ // and expectedEpoch is for retries.
+
+ // Setup EMPTY transaction and enable 2PC
+ when(transactionManager.getTransactionState(ArgumentMatchers.eq(transactionalId)))
+ .thenReturn(Right(None))
+ when(transactionManager.isTransaction2pcEnabled())
+ .thenReturn(true)
+
+ // Attempt InitProducerId with both keepPreparedTxn and expectedEpoch
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ expectedProducerIdAndEpoch = Some(new ProducerIdAndEpoch(producerId, producerEpoch)),
+ initProducerIdMockCallback
+ )
+
+ // Verify: Returns INVALID_REQUEST
+ assertEquals(Errors.INVALID_REQUEST, result.error,
+ "Should return INVALID_REQUEST when both keepPreparedTxn and expectedEpoch are set")
+ }
+
+ @Test
+ def shouldRejectKeepPreparedTxnWhenTransactionVersionTooLow(): Unit = {
+ // Test that keepPreparedTxn=true is rejected when transaction.version < TV_2.
+ // keepPreparedTxn requires TV_2 to ensure the 2PC metadata (nextProducerEpoch,
+ // prevProducerId, nextProducerId) is properly persisted to the transaction log.
+ // Without TV_2, these fields would be omitted during serialization, causing
+ // silent data corruption on coordinator recovery.
+
+ // Setup: ONGOING transaction at TV_0
+ val txnMetadata = new TransactionMetadata(
+ transactionalId,
+ producerId,
+ RecordBatch.NO_PRODUCER_ID,
+ RecordBatch.NO_PRODUCER_ID,
+ producerEpoch,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ Integer.MAX_VALUE, // 2PC transaction timeout
+ TransactionState.ONGOING,
+ partitions,
+ time.milliseconds(),
+ time.milliseconds(),
+ TV_0 // Transaction version 0 - doesn't support epoch bump
+ )
+
+ when(transactionManager.getTransactionState(ArgumentMatchers.eq(transactionalId)))
+ .thenReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+ when(transactionManager.isTransaction2pcEnabled())
+ .thenReturn(true)
+ when(transactionManager.transactionVersionLevel())
+ .thenReturn(TV_0) // Cluster finalized feature is TV_0
+ when(transactionManager.validateTransactionTimeoutMs(anyBoolean(), anyInt()))
+ .thenReturn(true)
+
+ // Attempt InitProducerId with keepPreparedTxn when TV < TV_2
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Returns UNSUPPORTED_VERSION
+ assertEquals(Errors.UNSUPPORTED_VERSION, result.error,
+ "Should return UNSUPPORTED_VERSION when keepPreparedTxn=true but transaction.version < TV_2")
+ }
+
+ @Test
+ def shouldRejectKeepPreparedTxnWithoutOngoingTransaction(): Unit = {
+ // Test that keepPreparedTxn=true without an ONGOING transaction returns an
+ // appropriate error. Properly implemented clients call keepPreparedTxn only
+ // after a crash with a transaction in flight, but server validates this.
+
+ def testKeepPreparedInState(state: TransactionState): Unit = {
+ val txnMetadata = new TransactionMetadata(
+ transactionalId,
+ producerId,
+ RecordBatch.NO_PRODUCER_ID,
+ RecordBatch.NO_PRODUCER_ID,
+ producerEpoch,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ RecordBatch.NO_PRODUCER_EPOCH,
+ txnTimeoutMs,
+ state,
+ new util.HashSet[TopicPartition](),
+ time.milliseconds(),
+ time.milliseconds(),
+ TV_0
+ )
+
+ when(transactionManager.getTransactionState(ArgumentMatchers.eq(transactionalId)))
+ .thenReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+
+ // Attempt InitProducerId with keepPreparedTxn in non-ONGOING state
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Returns error (no ONGOING transaction to prepare)
+ assertNotEquals(Errors.NONE, result.error,
+ s"Should return error when keepPreparedTxn=true in $state state")
+ }
+
+ // Test keepPreparedTxn in states where there's no ongoing transaction
+ testKeepPreparedInState(TransactionState.EMPTY)
+ testKeepPreparedInState(TransactionState.COMPLETE_COMMIT)
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
+ // This represents the COMPLETE_* transition metadata created by the internal prepareComplete() call
+ val capturedMarkerMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ verify(transactionMarkerChannelManager).addTxnMarkersToSend(
+ ArgumentMatchers.eq(coordinatorEpoch),
+ ArgumentMatchers.eq(txnResult),
+ any[TransactionMetadata],
+ capturedMarkerMetadata.capture()
+ )
+
+ // Verify successful completion
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to expected state.
+ // Index 0: helper's InitProducerId
+ // Index 1: test's InitProducerId (rotation)
+ // Index 2: EndTransaction
+ val prepareMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, prepareMetadata.txnState)
+
+ // Verify that transaction markers are sent with original (not rotated)
+ // producerId and startEpoch + 1. The ongoing transaction maintains the
+ // original producerId/epoch, so markers use those credentials.
+ assertEquals(testProducerId, prepareMetadata.producerId,
+ "Markers should use original producerId (ongoing transaction)")
+ assertEquals((startEpoch + 1).toShort, prepareMetadata.producerEpoch,
+ "Markers should use startEpoch + 1 (bumped from ongoing transaction epoch)")
+
+ // Simulate markers being sent and transaction completing to
+ // COMPLETE_COMMIT/COMPLETE_ABORT. At this point, handleEndTransaction
+ // has already internally called prepareComplete() in the
+ // sendTxnMarkersCallback, so the metadata already has
+ // pendingState=COMPLETE_COMMIT/COMPLETE_ABORT. We just need to write
+ // the transition to the log to complete it.
+
+ // Use the TxnTransitMetadata created by the internal prepareComplete()
+ val txnTransitMetadata = capturedMarkerMetadata.getValue
+
+ // Append COMPLETE_* transition to log (mock calls completeTransitionTo)
+ transactionManager.appendTransactionToLog(
+ transactionalId,
+ coordinatorEpoch,
+ txnTransitMetadata,
+ _ => {}, // No-op callback since we're just completing the transition
+ _ => false, // Don't retry on error
+ null
+ )
+
+ // Verify that rotation happened (regardless of whether it occurred during
+ // InitProducerId or EndTransaction). After transaction completes, the final
+ // state should have:
+ // - A new producerId (different from the original testProducerId)
+ // - Epoch determined by when rotation occurred (see scenarios below)
+ // - prevProducerId set to the original testProducerId for retry detection
+ assertNotEquals(testProducerId, txnTransitMetadata.producerId,
+ "Producer ID should have rotated by transaction completion")
+
+ // Verify epoch overflow occurred and total epoch increments.
+ // We do 3 total epoch increments in this test:
+ // 1. In setupPrepared2PcTxnWithBumpedClientEpoch: startEpoch → startEpoch + 1
+ // 2. Second InitProducerId call: startEpoch + 1 → startEpoch + 2 (or rotation)
+ // 3. EndTransaction: bumps nextProducerEpoch again (post-rotation or triggers rotation)
+ // Since rotation happens when trying to exceed MaxValue - 1 (32766), the final
+ // epoch wraps around.
+
+ // Assert 1: Final epoch is less than startEpoch (overflow occurred)
+ assertTrue(txnTransitMetadata.producerEpoch < startEpoch,
+ s"Overflow should have occurred: finalEpoch=${txnTransitMetadata.producerEpoch} < startEpoch=$startEpoch")
+
+ // Assert 2: Total epoch increments = 3 (accounting for overflow)
+ // Formula: finalEpoch + (MaxValue - startEpoch) = total increments
+ // This accounts for increments from startEpoch to MaxValue boundary, then
+ // wrap to 0 and increments to finalEpoch.
+ val totalIncrements = txnTransitMetadata.producerEpoch + Short.MaxValue - startEpoch
+ assertEquals(3, totalIncrements,
+ s"Should have 3 total epoch increments: finalEpoch=${txnTransitMetadata.producerEpoch} + " +
+ s"(MaxValue=${Short.MaxValue} - startEpoch=$startEpoch) = $totalIncrements")
+
+ assertEquals(testProducerId, txnTransitMetadata.prevProducerId,
+ "prevProducerId should track original ID for retry detection")
+
+ // Retry EndTransaction with same client credentials - should succeed as a retry
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use same client producer ID from InitProducerId result
+ clientProducerEpoch, // Use same client epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify retry succeeds with NONE
+ assertEquals(Errors.NONE, error, "Retry with client credentials should succeed")
+
+ // Reset mock to clear call history for next invocation
+ reset(transactionMarkerChannelManager)
+ }
+
+ // Test both commit and abort with rotation - use different producer IDs
+ // to avoid conflicts
+ mockPidGenerator() // Initialize PID generator
+
+ // Scenario 1: Rotation happens during InitProducerId call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=100, epoch=32765, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=100, epoch=32765,
+ // nextPid=100, nextEpoch=32766
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32766 triggers
+ // isEpochExhausted() → rotation. Ongoing, pid=100, epoch=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 4. CommitTxn/AbortTxn: bumps nextEpoch 0→1 to fence delayed requests
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=1, prevPid=100
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch, 100L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch, 200L)
+
+ // Scenario 2: Rotation happens during EndTransaction call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=300, epoch=32764, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32765
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32765 < 32766
+ // (not exhausted) → just epoch bump. Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32766
+ // 4. CommitTxn/AbortTxn: prepareAbortOrCommit bumps epoch 32764→32765,
+ // then generateTxnTransitMetadataForTxnCompletion sees
+ // isProducerEpochExhausted() → rotation. PrepareCommit, pid=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=0, prevPid=300
+ val startEpoch2 = (Short.MaxValue - 3).toShort // 32764
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch2, 300L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch2, 400L)
+ }
+
+ @Test
+ def shouldHandleProducerIdRotationPersistence(): Unit = {
+ // Test that producer ID rotation metadata is correctly persisted through state transitions.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction and make first InitProducerId call
+ // Helper will bump: 32765 → 32766, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Call InitProducerId again - this will exhaust at 32767 and rotate
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify rotation occurred
+ val rotatedProducerId = result.producerId
+ val rotatedEpoch = result.producerEpoch
+ assertNotEquals(producerId, rotatedProducerId, "Producer ID should have rotated")
+ assertEquals(0, rotatedEpoch, "Client epoch should reset to 0")
+
+ // Verify: Check captured metadata shows rotation with prevProducerId set
+ val rotationMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID should be stored in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch should be 0")
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId should be set to original ID for retry detection")
+
+ // Verify ongoing identity still uses original producer ID
+ assertEquals(producerId, rotationMetadata.producerId, "Ongoing transaction producerId unchanged")
+ assertEquals(startEpoch, rotationMetadata.producerEpoch, "Ongoing transaction epoch unchanged")
+
+ // Complete the transaction (commit)
+ coordinator.handleEndTransaction(
+ transactionalId,
+ rotatedProducerId, // Use client's producer ID (rotated)
+ rotatedEpoch, // Use client's epoch (0 after rotation)
+ TransactionResult.COMMIT,
+ TV_2,
+ endTxnCallback
+ )
+
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to PREPARE_COMMIT
+ val prepareCommitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState)
+
+ // Verify: The key aspect of persistence is that prevProducerId was
+ // correctly set during rotation. This ensures that after the transaction
+ // completes to EMPTY, retry detection will work correctly by recognizing
+ // requests from the old producer ID.
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId persisted for retry detection")
+
+ // Verify: The rotation information is persisted in the metadata. After
+ // this transaction completes, the next transaction will start with
+ // rotatedProducerId.
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID persisted in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch persisted as 0")
+
+ // Verify: The state transitions are captured correctly in metadata
+ assertEquals(3, capturedTransitMetadata.getAllValues.size(),
+ "Three state transitions captured")
+ assertEquals(TransactionState.ONGOING,
+ capturedTransitMetadata.getAllValues.get(0).txnState,
+ "First: InitProducerId → ONGOING")
+ assertEquals(TransactionState.ONGOING, rotationMetadata.txnState,
+ "Second: InitProducerId with rotation → ONGOING")
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState,
+ "Third: EndTransaction → PREPARE_COMMIT")
+ }
+
+ @Test
+ def shouldHandleConcurrentTransactionsInPrepareStatesTV2(): Unit = {
+ // Test that attempting EndTransaction while already in a PREPARE state
+ // returns CONCURRENT_TRANSACTIONS error in TV2. This error is retriable -
+ // the client should wait for the current transaction to complete and then
+ // retry.
+
+ def testConcurrentInPrepareState(prepareState: TransactionState, attemptedResult: TransactionResult): Unit = {
+ // Setup ONGOING transaction with 2PC enabled
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Determine transaction result to reach desired PREPARE state
+ val txnResult = if (prepareState == TransactionState.PREPARE_COMMIT) {
+ TransactionResult.COMMIT
+ } else {
+ TransactionResult.ABORT
+ }
+
+ // Call EndTransaction to move to PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Use client producer ID
+ result.producerEpoch, // Use client epoch
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify we're in the expected PREPARE state
+ assertEquals(Errors.NONE, error, "EndTransaction should succeed")
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(prepareState, transitMetadata.txnState,
+ s"Transaction should be in $prepareState state")
+
+ val stateBefore = transitMetadata.txnState
+
+ // Attempt another EndTransaction while still in PREPARE state
+ coordinator.handleEndTransaction(
+ transactionalId,
+ result.producerId, // Same client producer ID
+ result.producerEpoch, // Same client epoch
+ attemptedResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: TV2 returns CONCURRENT_TRANSACTIONS (retriable error)
+ assertEquals(Errors.CONCURRENT_TRANSACTIONS, error,
+ s"Should return CONCURRENT_TRANSACTIONS when attempting $attemptedResult while in $prepareState")
+
+ // Verify: State unchanged after rejection
+ assertEquals(stateBefore, capturedTransitMetadata.getAllValues.get(1).txnState,
Review Comment:
Will do.
##########
core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala:
##########
@@ -2602,4 +2695,801 @@ class TransactionCoordinatorTest {
// Final verification that hasFailedEpochFence was never set to true for TV2
assertFalse(txnMetadata.hasFailedEpochFence)
}
+
+ @Test
+ def shouldIncrementClientEpochOnMultipleInitPidWithKeepPreparedTxn(): Unit = {
+ // Test that multiple calls to InitProducerId(keepPreparedTxn=true) increment the client-facing epoch
+ // while keeping the ongoing transaction epoch constant.
+
+ // Setup: Use helper to set up 2PC transaction and make first InitProducerId call
+ val capturedNewMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Verify that epoch bumps a few times (the first bump happened during the setup)
+ var iteration = 0
+ do {
+ val expectedClientEpoch = (producerEpoch + 1 + iteration).toShort
+
+ // Verify: each previous iteration must've bumped client epoch
+ assertEquals(producerId, result.producerId)
+ assertEquals(expectedClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata has proper nextProducerEpoch
+ val transitMetadata = capturedNewMetadata.getAllValues.get(iteration)
+ assertEquals(producerId, transitMetadata.nextProducerId)
+ assertEquals(expectedClientEpoch, transitMetadata.nextProducerEpoch)
+
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ iteration = iteration + 1
+ } while (iteration < 4)
+ }
+
+ @Test
+ def shouldCompletePreparedTransactionWithOngoingTxnEpoch(): Unit = {
+ // Test that a prepared transaction can be completed (committed or aborted) using the client-facing
+ // epoch after InitProducerId(keepPrepared=true), while markers are sent with the bumped ongoing epoch.
+
+ def testComplete(txnResult: TransactionResult): Unit = {
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch()
+
+ // Setup bumped from epoch 1 to epoch 2, now bump again from 2 to 3
+ // Action: Call InitProducerId again to bump client epoch once more
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: Second InitProducerId bumped client epoch from 2 to 3
+ val secondClientEpoch = (producerEpoch + 2).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(secondClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(producerEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify: captured metadata from second InitProducerId has proper nextProducerEpoch
+ val secondInitMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(producerId, secondInitMetadata.nextProducerId)
+ assertEquals(secondClientEpoch, secondInitMetadata.nextProducerEpoch)
+
+ val bumpedEpoch = (producerEpoch + 1).toShort // Epoch bumped during prepare phase
+ val tripleBumpedEpoch = (producerEpoch + 3).toShort // NextProducerEpoch bumped a third time
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Action: Call EndTransaction with second client-facing epoch
+ // After second InitProducerId(keepPrepared=true), validation checks against latest clientProducerEpoch
+ coordinator.handleEndTransaction(
+ transactionalId,
+ producerId,
+ secondClientEpoch, // Must use latest client-facing epoch for validation to pass
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify: Transaction completes successfully
+ assertEquals(Errors.NONE, error)
+
+ // Verify: Transaction transitions to expected state. Get the third call
+ // (index 2) which is the complete, index 0 was first InitProducerId,
+ // index 1 was second.
+ val transitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, transitMetadata.txnState)
+
+ // Verify: Markers are sent with bumped epoch.
+ // prepareAbortOrCommit bumps producerEpoch for TV2.
+ assertEquals(producerId, transitMetadata.producerId)
+ assertEquals(bumpedEpoch, transitMetadata.producerEpoch)
+
+ // Verify: After complete, nextProducerEpoch is bumped again
+ // (third time total).
+ assertEquals(tripleBumpedEpoch, transitMetadata.nextProducerEpoch)
+ }
+
+ // Test both commit and abort
+ testComplete(TransactionResult.COMMIT)
+ testComplete(TransactionResult.ABORT)
+ }
+
+ @Test
+ def shouldRejectAddPartitionsAfterInitProducerIdWithKeepPreparedTxn(): Unit = {
+ // Test that trying to add partitions with the client-facing epoch (after
+ // calling InitProducerId with keepPreparedTxn=true) returns
+ // INVALID_TXN_STATE. A properly implemented client should never try this,
+ // but the server validates it.
+
+ setupPrepared2PcTxnWithBumpedClientEpoch() // Return value not used - only need the setup and mocks
+
+ val clientEpoch = (producerEpoch + 1).toShort
+
+ val newPartitions = new util.HashSet[TopicPartition]()
+ newPartitions.add(new TopicPartition("topic2", 0))
+
+ // Action: Client tries AddPartitions with the client-facing epoch
+ // This is invalid because after InitProducerId(keepPreparedTxn=true), the client should
+ // not add partitions.
+ coordinator.handleAddPartitionsToTransaction(
+ transactionalId,
+ producerId,
+ clientEpoch,
+ newPartitions,
+ errorsCallback,
+ TV_2
+ )
+
+ // Verify: Returns INVALID_TXN_STATE (cannot add partitions in this state)
+ assertEquals(Errors.INVALID_TXN_STATE, error)
+ }
+
+ @Test
+ def shouldRotateProducerIdWhenClientEpochExhausted(): Unit = {
+ // Test that when the client epoch reaches Short.MaxValue, a new producer ID is allocated
+ // and the client epoch resets to 0, while the ongoing transaction identity remains unchanged.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Intermediate verify: epoch bumped to 32766, same producer ID
+ val firstClientEpoch = (Short.MaxValue - 1).toShort
+ assertEquals(producerId, result.producerId)
+ assertEquals(firstClientEpoch, result.producerEpoch)
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from first call
+ val firstMetadata = capturedTransitMetadata.getAllValues.get(0)
+ assertEquals(producerId, firstMetadata.nextProducerId)
+ assertEquals(firstClientEpoch, firstMetadata.nextProducerEpoch)
+
+ // Second call: exhaust at 32767, should rotate producer ID
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify: new producer ID allocated, epoch reset to 0
+ val rotatedProducerId = nextPid - 1 // Most recently allocated PID
+ assertTrue(rotatedProducerId != producerId, "Producer ID should have rotated")
+ assertEquals(rotatedProducerId, result.producerId)
+ assertEquals(0, result.producerEpoch)
+ // Ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ // Verify captured metadata from second call shows rotation
+ val secondMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, secondMetadata.nextProducerId)
+ assertEquals(0, secondMetadata.nextProducerEpoch)
+ // Original ongoing identity unchanged
+ assertEquals(producerId, secondMetadata.producerId)
+ assertEquals(startEpoch, secondMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldHandleMultipleProducerIdRotationsForOngoingTxn(): Unit = {
+ // Test that multiple producer ID rotations can occur during a single ongoing transaction,
+ // and the ongoing transaction identity remains constant throughout.
+
+ val startEpoch = (Short.MaxValue - 3).toShort // 32764
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction at epoch boundary and make first InitProducerId call
+ // Helper will bump: 32764 → 32765, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Track rotations starting from the state after helper's InitProducerId call
+ var rotationCount = 0
+ val targetRotations = 2
+ var currentClientEpoch: Short = result.producerEpoch // 32765 after helper
+ val allocatedProducerIds = scala.collection.mutable.Set[Long](producerId)
+
+ // Loop until we've seen 2 rotations
+ while (rotationCount < targetRotations) {
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Check if rotation occurred (epoch wrapped to 0)
+ if (result.producerEpoch == 0 && currentClientEpoch != 0) {
+ rotationCount += 1
+ // Verify new producer ID allocated
+ assertFalse(allocatedProducerIds.contains(result.producerId),
+ s"Producer ID ${result.producerId} should be new on rotation ${rotationCount}")
+ allocatedProducerIds.add(result.producerId)
+ }
+
+ // Verify ongoing transaction identity unchanged
+ assertEquals(producerId, result.ongoingTxnProducerId)
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch)
+
+ currentClientEpoch = result.producerEpoch
+ }
+
+ assertEquals(2, rotationCount, "Should have seen exactly 2 rotations")
+ assertEquals(3, allocatedProducerIds.size, "Should have 3 unique producer IDs (original + 2 rotations)")
+
+ // Verify final state
+ val finalMetadata = capturedTransitMetadata.getAllValues.get(capturedTransitMetadata.getAllValues.size - 1)
+ assertEquals(0, finalMetadata.nextProducerEpoch, "Final epoch should be 0 after rotations")
+ assertFalse(finalMetadata.nextProducerId == producerId, "Final producer ID should be rotated")
+ // Ongoing identity still original
+ assertEquals(producerId, finalMetadata.producerId)
+ assertEquals(startEpoch, finalMetadata.producerEpoch)
+ }
+
+ @Test
+ def shouldCompleteTransactionWithExhaustedEpochs(): Unit = {
+ // Test that transactions can be committed or aborted successfully even
+ // after producer ID rotation due to epoch exhaustion. This tests the case
+ // where overflow happens either during the second InitProducerId call or
+ // during endTransaction.
+ def testCompleteWithRotation(txnResult: TransactionResult, startEpoch: Short, testProducerId: Long): Unit = {
+
+ // Setup: Use helper to create ONGOING transaction and make first
+ // InitProducerId call. Helper will bump: startEpoch → startEpoch + 1,
+ // establishing dual identity.
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch, testProducerId)
+
+ // Call InitProducerId again - may rotate producer id (depends on startEpoch)
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Store client-facing producerId/epoch from InitProducerId result.
+ // Rotation may have happened during InitProducerId or later during
+ // EndTransaction.
+ val clientProducerId = result.producerId
+ val clientProducerEpoch = result.producerEpoch
+
+ // Verify ongoing transaction producerId/epoch unchanged after
+ // InitProducerId (regardless of when rotation happens).
+ assertEquals(testProducerId, result.ongoingTxnProducerId, "Ongoing ID unchanged")
+ assertEquals(startEpoch, result.ongoingTxnProducerEpoch, "Ongoing epoch unchanged")
+
+ val expectedState = if (txnResult == TransactionResult.COMMIT) {
+ TransactionState.PREPARE_COMMIT
+ } else {
+ TransactionState.PREPARE_ABORT
+ }
+
+ // Complete the transaction
+ // TV2 validates against clientProducerId/clientProducerEpoch returned from InitProducerId
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use client's producer ID from InitProducerId result
+ clientProducerEpoch, // Use client's epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Capture the TxnTransitMetadata that was passed to addTxnMarkersToSend
+ // This represents the COMPLETE_* transition metadata created by the internal prepareComplete() call
+ val capturedMarkerMetadata: ArgumentCaptor[TxnTransitMetadata] = ArgumentCaptor.forClass(classOf[TxnTransitMetadata])
+ verify(transactionMarkerChannelManager).addTxnMarkersToSend(
+ ArgumentMatchers.eq(coordinatorEpoch),
+ ArgumentMatchers.eq(txnResult),
+ any[TransactionMetadata],
+ capturedMarkerMetadata.capture()
+ )
+
+ // Verify successful completion
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to expected state.
+ // Index 0: helper's InitProducerId
+ // Index 1: test's InitProducerId (rotation)
+ // Index 2: EndTransaction
+ val prepareMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(expectedState, prepareMetadata.txnState)
+
+ // Verify that transaction markers are sent with original (not rotated)
+ // producerId and startEpoch + 1. The ongoing transaction maintains the
+ // original producerId/epoch, so markers use those credentials.
+ assertEquals(testProducerId, prepareMetadata.producerId,
+ "Markers should use original producerId (ongoing transaction)")
+ assertEquals((startEpoch + 1).toShort, prepareMetadata.producerEpoch,
+ "Markers should use startEpoch + 1 (bumped from ongoing transaction epoch)")
+
+ // Simulate markers being sent and transaction completing to
+ // COMPLETE_COMMIT/COMPLETE_ABORT. At this point, handleEndTransaction
+ // has already internally called prepareComplete() in the
+ // sendTxnMarkersCallback, so the metadata already has
+ // pendingState=COMPLETE_COMMIT/COMPLETE_ABORT. We just need to write
+ // the transition to the log to complete it.
+
+ // Use the TxnTransitMetadata created by the internal prepareComplete()
+ val txnTransitMetadata = capturedMarkerMetadata.getValue
+
+ // Append COMPLETE_* transition to log (mock calls completeTransitionTo)
+ transactionManager.appendTransactionToLog(
+ transactionalId,
+ coordinatorEpoch,
+ txnTransitMetadata,
+ _ => {}, // No-op callback since we're just completing the transition
+ _ => false, // Don't retry on error
+ null
+ )
+
+ // Verify that rotation happened (regardless of whether it occurred during
+ // InitProducerId or EndTransaction). After transaction completes, the final
+ // state should have:
+ // - A new producerId (different from the original testProducerId)
+ // - Epoch determined by when rotation occurred (see scenarios below)
+ // - prevProducerId set to the original testProducerId for retry detection
+ assertNotEquals(testProducerId, txnTransitMetadata.producerId,
+ "Producer ID should have rotated by transaction completion")
+
+ // Verify epoch overflow occurred and total epoch increments.
+ // We do 3 total epoch increments in this test:
+ // 1. In setupPrepared2PcTxnWithBumpedClientEpoch: startEpoch → startEpoch + 1
+ // 2. Second InitProducerId call: startEpoch + 1 → startEpoch + 2 (or rotation)
+ // 3. EndTransaction: bumps nextProducerEpoch again (post-rotation or triggers rotation)
+ // Since rotation happens when trying to exceed MaxValue - 1 (32766), the final
+ // epoch wraps around.
+
+ // Assert 1: Final epoch is less than startEpoch (overflow occurred)
+ assertTrue(txnTransitMetadata.producerEpoch < startEpoch,
+ s"Overflow should have occurred: finalEpoch=${txnTransitMetadata.producerEpoch} < startEpoch=$startEpoch")
+
+ // Assert 2: Total epoch increments = 3 (accounting for overflow)
+ // Formula: finalEpoch + (MaxValue - startEpoch) = total increments
+ // This accounts for increments from startEpoch to MaxValue boundary, then
+ // wrap to 0 and increments to finalEpoch.
+ val totalIncrements = txnTransitMetadata.producerEpoch + Short.MaxValue - startEpoch
+ assertEquals(3, totalIncrements,
+ s"Should have 3 total epoch increments: finalEpoch=${txnTransitMetadata.producerEpoch} + " +
+ s"(MaxValue=${Short.MaxValue} - startEpoch=$startEpoch) = $totalIncrements")
+
+ assertEquals(testProducerId, txnTransitMetadata.prevProducerId,
+ "prevProducerId should track original ID for retry detection")
+
+ // Retry EndTransaction with same client credentials - should succeed as a retry
+ coordinator.handleEndTransaction(
+ transactionalId,
+ clientProducerId, // Use same client producer ID from InitProducerId result
+ clientProducerEpoch, // Use same client epoch from InitProducerId result
+ txnResult,
+ TV_2,
+ endTxnCallback
+ )
+
+ // Verify retry succeeds with NONE
+ assertEquals(Errors.NONE, error, "Retry with client credentials should succeed")
+
+ // Reset mock to clear call history for next invocation
+ reset(transactionMarkerChannelManager)
+ }
+
+ // Test both commit and abort with rotation - use different producer IDs
+ // to avoid conflicts
+ mockPidGenerator() // Initialize PID generator
+
+ // Scenario 1: Rotation happens during InitProducerId call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=100, epoch=32765, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=100, epoch=32765,
+ // nextPid=100, nextEpoch=32766
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32766 triggers
+ // isEpochExhausted() → rotation. Ongoing, pid=100, epoch=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 4. CommitTxn/AbortTxn: bumps nextEpoch 0→1 to fence delayed requests
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=1, prevPid=100
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch, 100L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch, 200L)
+
+ // Scenario 2: Rotation happens during EndTransaction call
+ // Example flow:
+ // 1. (initial): Ongoing, pid=300, epoch=32764, nextPid=-1, nextEpoch=-1
+ // 2. InitProducerId(keepPreparedTxn): Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32765
+ // 3. InitProducerId(keepPreparedTxn): clientEpoch=32765 < 32766
+ // (not exhausted) → just epoch bump. Ongoing, pid=300, epoch=32764,
+ // nextPid=300, nextEpoch=32766
+ // 4. CommitTxn/AbortTxn: prepareAbortOrCommit bumps epoch 32764→32765,
+ // then generateTxnTransitMetadataForTxnCompletion sees
+ // isProducerEpochExhausted() → rotation. PrepareCommit, pid=32765,
+ // nextPid=<new>, nextEpoch=0
+ // 5. Complete: CompleteCommit, pid=<new>, epoch=0, prevPid=300
+ val startEpoch2 = (Short.MaxValue - 3).toShort // 32764
+ testCompleteWithRotation(TransactionResult.COMMIT, startEpoch2, 300L)
+ testCompleteWithRotation(TransactionResult.ABORT, startEpoch2, 400L)
+ }
+
+ @Test
+ def shouldHandleProducerIdRotationPersistence(): Unit = {
+ // Test that producer ID rotation metadata is correctly persisted through state transitions.
+
+ val startEpoch = (Short.MaxValue - 2).toShort // 32765
+ mockPidGenerator()
+
+ // Setup: Use helper to create ONGOING transaction and make first InitProducerId call
+ // Helper will bump: 32765 → 32766, establishing dual identity
+ val capturedTransitMetadata = setupPrepared2PcTxnWithBumpedClientEpoch(startEpoch)
+
+ // Call InitProducerId again - this will exhaust at 32767 and rotate
+ coordinator.handleInitProducerId(
+ transactionalId,
+ txnTimeoutMs,
+ enableTwoPCFlag = true,
+ keepPreparedTxn = true,
+ None,
+ initProducerIdMockCallback
+ )
+
+ // Verify rotation occurred
+ val rotatedProducerId = result.producerId
+ val rotatedEpoch = result.producerEpoch
+ assertNotEquals(producerId, rotatedProducerId, "Producer ID should have rotated")
+ assertEquals(0, rotatedEpoch, "Client epoch should reset to 0")
+
+ // Verify: Check captured metadata shows rotation with prevProducerId set
+ val rotationMetadata = capturedTransitMetadata.getAllValues.get(1)
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID should be stored in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
+ "Rotated epoch should be 0")
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId should be set to original ID for retry detection")
+
+ // Verify ongoing identity still uses original producer ID
+ assertEquals(producerId, rotationMetadata.producerId, "Ongoing transaction producerId unchanged")
+ assertEquals(startEpoch, rotationMetadata.producerEpoch, "Ongoing transaction epoch unchanged")
+
+ // Complete the transaction (commit)
+ coordinator.handleEndTransaction(
+ transactionalId,
+ rotatedProducerId, // Use client's producer ID (rotated)
+ rotatedEpoch, // Use client's epoch (0 after rotation)
+ TransactionResult.COMMIT,
+ TV_2,
+ endTxnCallback
+ )
+
+ assertEquals(Errors.NONE, error)
+
+ // Verify transition to PREPARE_COMMIT
+ val prepareCommitMetadata = capturedTransitMetadata.getAllValues.get(2)
+ assertEquals(TransactionState.PREPARE_COMMIT, prepareCommitMetadata.txnState)
+
+ // Verify: The key aspect of persistence is that prevProducerId was
+ // correctly set during rotation. This ensures that after the transaction
+ // completes to EMPTY, retry detection will work correctly by recognizing
+ // requests from the old producer ID.
+ assertEquals(producerId, rotationMetadata.prevProducerId,
+ "prevProducerId persisted for retry detection")
+
+ // Verify: The rotation information is persisted in the metadata. After
+ // this transaction completes, the next transaction will start with
+ // rotatedProducerId.
+ assertEquals(rotatedProducerId, rotationMetadata.nextProducerId,
+ "Rotated ID persisted in nextProducerId")
+ assertEquals(0, rotationMetadata.nextProducerEpoch,
Review Comment:
Will do.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org