Update to use subclassing to avoid so much change

Author: kirktrue

clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java

@@ -120,76 +120,6 @@ public class TransactionManager {
     private final Set<TopicPartition> newPartitionsInTransaction;
     private final Set<TopicPartition> pendingPartitionsInTransaction;
     private final Set<TopicPartition> partitionsInTransaction;
-
-    /**
-     * During its normal course of operations, the transaction manager transitions through different internal
-     * states (i.e. by updating {@link #currentState}) to one of those defined in {@link State}. These state transitions
-     * result from actions on one of the following classes of threads:
-     *
-     * <ul>
-     *     <li><em>Application</em> threads that invokes {@link Producer} API calls</li>
-     *     <li><em>{@link Sender}</em> thread operations</li>
-     * </ul>
-     *
-     * When an invalid state transition is detected during execution on an <em>application</em> thread, the
-     * {@link #currentState} is <em>not updated</em> and an {@link IllegalStateException} is thrown. This gives the
-     * application the opportunity to fix the issue without permanently poisoning the state of the
-     * transaction manager. The {@link Producer} API calls that perform a state transition include:
-     *
-     * <ul>
-     *     <li>{@link Producer#initTransactions()} calls {@link #initializeTransactions()}</li>
-     *     <li>{@link Producer#beginTransaction()} calls {@link #beginTransaction()}</li>
-     *     <li>{@link Producer#commitTransaction()}} calls {@link #beginCommit()}</li>
-     *     <li>{@link Producer#abortTransaction()} calls {@link #beginAbort()}
-     *     </li>
-     *     <li>{@link Producer#sendOffsetsToTransaction(Map, ConsumerGroupMetadata)} calls
-     *         {@link #sendOffsetsToTransaction(Map, ConsumerGroupMetadata)}
-     *     </li>
-     *     <li>{@link Producer#send(ProducerRecord)} (and its variants) calls
-     *         {@link #maybeAddPartition(TopicPartition)} and
-     *         {@link #maybeTransitionToErrorState(RuntimeException)}
-     *     </li>
-     * </ul>
-     *
-     * <p/>
-     *
-     * The {@link Producer} is implemented such that much of its work delegated to and performed asynchronously on the
-     * <em>{@link Sender}</em> thread. This includes record batching, network I/O, broker response handlers, etc. If an
-     * invalid state transition is detected in the <em>{@link Sender}</em> thread, in addition to throwing an
-     * {@link IllegalStateException}, the transaction manager intentionally "poisons" itself by setting its
-     * {@link #currentState} to {@link State#FATAL_ERROR}, a state from which it cannot recover.
-     *
-     * <p/>
-     *
-     * It's important to prevent possible corruption when the transaction manager has determined that it is in a
-     * fatal state. Subsequent transaction operations attempted via either the <em>application</em> or the
-     * <em>{@link Sender}</em> thread should fail. This is achieved when these operations invoke the
-     * {@link #maybeFailWithError()} method, as it causes a {@link KafkaException} to be thrown, ensuring the stated
-     * transactional guarantees are not violated.
-     *
-     * <p/>
-     *
-     * See KAFKA-14831 for more detail.
-     */
-    public interface InvalidStateTransitionAttemptHandler {
-
-        /**
-         * Callback to determine if the transaction manager's {@link #currentState} should be set to
-         * {@link State#FATAL_ERROR} to prevent possible transaction corruption.
-         *
-         * @return {@code true} to set state to {@link State#FATAL_ERROR} before throwing an exception,
-         *         {@code false} to throw an exception without first changing the state
-         */
-        boolean isFatalState();
-    }
-
-    /**
-     * The default logic is to poison the internal {@link #currentState state} if an invalid state transition is
-     * attempted on the {@link Sender.SenderThread sender's thread}.
-     */
-    public static final InvalidStateTransitionAttemptHandler DEFAULT_INVALID_STATE_TRANSITION_ATTEMPT_HANDLER = () -> Thread.currentThread() instanceof Sender.SenderThread;
-    private final InvalidStateTransitionAttemptHandler invalidStateTransitionAttemptHandler;
-
     private PendingStateTransition pendingTransition;
 
     // This is used by the TxnRequestHandlers to control how long to back off before a given request is retried.
@@ -274,15 +204,6 @@ public TransactionManager(final LogContext logContext,
                               final int transactionTimeoutMs,
                               final long retryBackoffMs,
                               final ApiVersions apiVersions) {
-        this(logContext, transactionalId, transactionTimeoutMs, retryBackoffMs, apiVersions, DEFAULT_INVALID_STATE_TRANSITION_ATTEMPT_HANDLER);
-    }
-
-    public TransactionManager(final LogContext logContext,
-                              final String transactionalId,
-                              final int transactionTimeoutMs,
-                              final long retryBackoffMs,
-                              final ApiVersions apiVersions,
-                              final InvalidStateTransitionAttemptHandler invalidStateTransitionAttemptHandler) {
         this.producerIdAndEpoch = ProducerIdAndEpoch.NONE;
         this.transactionalId = transactionalId;
         this.log = logContext.logger(TransactionManager.class);
@@ -299,7 +220,6 @@ public TransactionManager(final LogContext logContext,
         this.retryBackoffMs = retryBackoffMs;
         this.txnPartitionMap = new TxnPartitionMap(logContext);
         this.apiVersions = apiVersions;
-        this.invalidStateTransitionAttemptHandler = invalidStateTransitionAttemptHandler;
     }
 
     public synchronized TransactionalRequestResult initializeTransactions() {
@@ -1086,7 +1006,7 @@ private void transitionTo(State target, RuntimeException error) {
             String message = idString + "Invalid transition attempted from state "
                     + currentState.name() + " to state " + target.name();
 
-            if (invalidStateTransitionAttemptHandler.isFatalState()) {
+            if (shouldSetToFatalErrorState()) {
                 currentState = State.FATAL_ERROR;
                 lastError = new IllegalStateException(message);
                 throw lastError;
@@ -1109,6 +1029,63 @@ private void transitionTo(State target, RuntimeException error) {
         currentState = target;
     }
 
+    /**
+     * During its normal course of operations, the transaction manager transitions through different internal
+     * states (i.e. by updating {@link #currentState}) to one of those defined in {@link State}. These state transitions
+     * result from actions on one of the following classes of threads:
+     *
+     * <ul>
+     *     <li><em>Application</em> threads that invokes {@link Producer} API calls</li>
+     *     <li><em>{@link Sender}</em> thread operations</li>
+     * </ul>
+     *
+     * When an invalid state transition is detected during execution on an <em>application</em> thread, the
+     * {@link #currentState} is <em>not updated</em> and an {@link IllegalStateException} is thrown. This gives the
+     * application the opportunity to fix the issue without permanently poisoning the state of the
+     * transaction manager. The {@link Producer} API calls that perform a state transition include:
+     *
+     * <ul>
+     *     <li>{@link Producer#initTransactions()} calls {@link #initializeTransactions()}</li>
+     *     <li>{@link Producer#beginTransaction()} calls {@link #beginTransaction()}</li>
+     *     <li>{@link Producer#commitTransaction()}} calls {@link #beginCommit()}</li>
+     *     <li>{@link Producer#abortTransaction()} calls {@link #beginAbort()}
+     *     </li>
+     *     <li>{@link Producer#sendOffsetsToTransaction(Map, ConsumerGroupMetadata)} calls
+     *         {@link #sendOffsetsToTransaction(Map, ConsumerGroupMetadata)}
+     *     </li>
+     *     <li>{@link Producer#send(ProducerRecord)} (and its variants) calls
+     *         {@link #maybeAddPartition(TopicPartition)} and
+     *         {@link #maybeTransitionToErrorState(RuntimeException)}
+     *     </li>
+     * </ul>
+     *
+     * <p/>
+     *
+     * The {@link Producer} is implemented such that much of its work delegated to and performed asynchronously on the
+     * <em>{@link Sender}</em> thread. This includes record batching, network I/O, broker response handlers, etc. If an
+     * invalid state transition is detected in the <em>{@link Sender}</em> thread, in addition to throwing an
+     * {@link IllegalStateException}, the transaction manager intentionally "poisons" itself by setting its
+     * {@link #currentState} to {@link State#FATAL_ERROR}, a state from which it cannot recover.
+     *
+     * <p/>
+     *
+     * It's important to prevent possible corruption when the transaction manager has determined that it is in a
+     * fatal state. Subsequent transaction operations attempted via either the <em>application</em> or the
+     * <em>{@link Sender}</em> thread should fail. This is achieved when these operations invoke the
+     * {@link #maybeFailWithError()} method, as it causes a {@link KafkaException} to be thrown, ensuring the stated
+     * transactional guarantees are not violated.
+     *
+     * <p/>
+     *
+     * See KAFKA-14831 for more detail.
+     *
+     * @return {@code true} to set state to {@link State#FATAL_ERROR} before throwing an exception,
+     *         {@code false} to throw an exception without first changing the state
+     */
+    protected boolean shouldSetToFatalErrorState() {
+        return Thread.currentThread() instanceof Sender.SenderThread;
+    }
+
     private void ensureTransactional() {
         if (!isTransactional())
             throw new IllegalStateException("Transactional method invoked on a non-transactional producer.");

clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java

@@ -149,7 +149,7 @@ public class TransactionManagerTest {
 
     private RecordAccumulator accumulator = null;
     private Sender sender = null;
-    private TransactionManager transactionManager = null;
+    private TestableTransactionManager transactionManager = null;
     private Node brokerNode = null;
     private long finalizedFeaturesEpoch = 0;
 
@@ -163,12 +163,6 @@ public void setup() {
     }
 
     private void initializeTransactionManager(Optional<String> transactionalId, boolean transactionV2Enabled) {
-        initializeTransactionManager(transactionalId, transactionV2Enabled, TransactionManager.DEFAULT_INVALID_STATE_TRANSITION_ATTEMPT_HANDLER);
-    }
-
-    private void initializeTransactionManager(Optional<String> transactionalId,
-                                              boolean transactionV2Enabled,
-                                              TransactionManager.InvalidStateTransitionAttemptHandler invalidStateTransitionAttemptHandler) {
         Metrics metrics = new Metrics(time);
 
         apiVersions.update("0", new NodeApiVersions(Arrays.asList(
@@ -194,8 +188,8 @@ private void initializeTransactionManager(Optional<String> transactionalId,
                 .setMinVersionLevel(transactionV2Enabled ? (short) 2 : (short) 1)),
             finalizedFeaturesEpoch));
         finalizedFeaturesEpoch += 1;
-        this.transactionManager = new TransactionManager(logContext, transactionalId.orElse(null),
-                transactionTimeoutMs, DEFAULT_RETRY_BACKOFF_MS, apiVersions, invalidStateTransitionAttemptHandler);
+        this.transactionManager = new TestableTransactionManager(logContext, transactionalId.orElse(null),
+                transactionTimeoutMs, DEFAULT_RETRY_BACKOFF_MS, apiVersions);
 
         int batchSize = 16 * 1024;
         int deliveryTimeoutMs = 3000;
@@ -1044,7 +1038,7 @@ public void testTransactionManagerDisablesV2() {
                 .setMaxVersionLevel((short) 1)
                 .setMinVersionLevel((short) 1)),
             0));
-        this.transactionManager = new TransactionManager(logContext, transactionalId,
+        this.transactionManager = new TestableTransactionManager(logContext, transactionalId,
             transactionTimeoutMs, DEFAULT_RETRY_BACKOFF_MS, apiVersions);
 
         int batchSize = 16 * 1024;
@@ -3808,8 +3802,8 @@ public void testBackgroundInvalidStateTransitionIsFatal() {
         // The default logic is to poison the transaction manager on an invalid state transition attempt if that
         // attempt happens on the Sender thread. Here the logic is altered to always poison on invalid attempts to
         // mimic as though it's being invoked by the Sender.
-        TransactionManager.InvalidStateTransitionAttemptHandler testHandler = () -> true;
-        initializeTransactionManager(Optional.of(transactionalId), true, testHandler);
+        initializeTransactionManager(Optional.of(transactionalId), true);
+        transactionManager.setShouldSetToFatalErrorStateOverride(true);
         doInitTransactions();
         assertTrue(transactionManager.isTransactional());
 
@@ -4382,4 +4376,26 @@ private void runUntil(Supplier<Boolean> condition) {
         ProducerTestUtils.runUntil(sender, condition);
     }
 
+    private static class TestableTransactionManager extends TransactionManager {
+
+        private Optional<Boolean> shouldSetToFatalErrorStateOverride;
+
+        public TestableTransactionManager(LogContext logContext,
+                                          String transactionalId,
+                                          int transactionTimeoutMs,
+                                          long retryBackoffMs,
+                                          ApiVersions apiVersions) {
+            super(logContext, transactionalId, transactionTimeoutMs, retryBackoffMs, apiVersions);
+            this.shouldSetToFatalErrorStateOverride = Optional.empty();
+        }
+
+        private void setShouldSetToFatalErrorStateOverride(boolean override) {
+            shouldSetToFatalErrorStateOverride = Optional.of(override);
+        }
+
+        @Override
+        protected boolean shouldSetToFatalErrorState() {
+            return shouldSetToFatalErrorStateOverride.orElseGet(super::shouldSetToFatalErrorState);
+        }
+    }
 }