From f54975c33135140351c50370282e86c49c81bbdd Mon Sep 17 00:00:00 2001
From: Josep Prat <josep.prat@aiven.io>
Date: Mon, 5 Feb 2024 16:43:53 +0100
Subject: [PATCH 001/258] MINOR: Fix code listings in quickstart.html (#10767)

* MINOR: Fix code listings in quickstart

Uses the right syntax highlighting

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Luke Chen <showuon@gmail.com>
---
 docs/quickstart.html | 35 ++++++++++++++---------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/docs/quickstart.html b/docs/quickstart.html
index e0c842793473b..455fd8b267c24 100644
--- a/docs/quickstart.html
+++ b/docs/quickstart.html
@@ -152,7 +152,7 @@ <h4 class="anchor-heading">
 
         <pre class="line-numbers"><code class="language-bash">$ bin/kafka-topics.sh --describe --topic quickstart-events --bootstrap-server localhost:9092
 Topic: quickstart-events        TopicId: NPmZHyhbR9y00wMglMH2sg PartitionCount: 1       ReplicationFactor: 1	Configs:
-    Topic: quickstart-events Partition: 0    Leader: 0   Replicas: 0 Isr: 0</code></pre>
+Topic: quickstart-events Partition: 0    Leader: 0   Replicas: 0 Isr: 0</code></pre>
     </div>
 
     <div class="quickstart-step">
@@ -173,8 +173,8 @@ <h4 class="anchor-heading">
         </p>
 
         <pre class="line-numbers"><code class="language-bash">$ bin/kafka-console-producer.sh --topic quickstart-events --bootstrap-server localhost:9092
-This is my first event
-This is my second event</code></pre>
+&gt;This is my first event
+&gt;This is my second event</code></pre>
 
         <p>
             You can stop the producer client with <code>Ctrl-C</code> at any time.
@@ -233,19 +233,16 @@ <h4 class="anchor-heading">
             Edit the <code class="language-bash">config/connect-standalone.properties</code> file, add or change the <code>plugin.path</code> configuration property match the following, and save the file:
         </p>
 
-        <pre class="brush: bash;">
-&gt; echo "plugin.path=libs/connect-file-{{fullDotVersion}}.jar"</pre>
+        <pre class="line-numbers"><code class="language-bash">$ echo "plugin.path=libs/connect-file-{{fullDotVersion}}.jar &gt;&gt; config/connect-standalone.properties"</code></pre>
 
         <p>
             Then, start by creating some seed data to test with:
         </p>
 
-        <pre class="brush: bash;">
-&gt; echo -e "foo\nbar" > test.txt</pre>
+        <pre class="line-numbers"><code class="language-bash">$ echo -e "foo\nbar" &gt; test.txt</code></pre>
         Or on Windows:
-        <pre class="brush: bash;">
-&gt; echo foo> test.txt
-&gt; echo bar>> test.txt</pre>
+        <pre class="line-numbers"><code class="language-bash">$ echo foo&gt; test.txt
+$ echo bar&gt;&gt; test.txt</code></pre>
 
         <p>
             Next, we'll start two connectors running in <i>standalone</i> mode, which means they run in a single, local, dedicated
@@ -255,8 +252,7 @@ <h4 class="anchor-heading">
             class to instantiate, and any other configuration required by the connector.
         </p>
 
-        <pre class="brush: bash;">
-&gt; bin/connect-standalone.sh config/connect-standalone.properties config/connect-file-source.properties config/connect-file-sink.properties</pre>
+        <pre class="line-numbers"><code class="language-bash">$ bin/connect-standalone.sh config/connect-standalone.properties config/connect-file-source.properties config/connect-file-sink.properties</code></pre>
 
         <p>
             These sample configuration files, included with Kafka, use the default local cluster configuration you started earlier
@@ -273,10 +269,9 @@ <h4 class="anchor-heading">
         </p>
 
 
-        <pre class="brush: bash;">
-&gt; more test.sink.txt
+        <pre class="line-numbers"><code class="language-bash">$ more test.sink.txt
 foo
-bar</pre>
+bar</code></pre>
 
         <p>
             Note that the data is being stored in the Kafka topic <code>connect-test</code>, so we can also run a console consumer to see the
@@ -284,16 +279,14 @@ <h4 class="anchor-heading">
         </p>
 
 
-        <pre class="brush: bash;">
-&gt; bin/kafka-console-consumer.sh --bootstrap-server localhost:9092 --topic connect-test --from-beginning
+        <pre class="line-numbers"><code class="language-bash">$ bin/kafka-console-consumer.sh --bootstrap-server localhost:9092 --topic connect-test --from-beginning
 {"schema":{"type":"string","optional":false},"payload":"foo"}
 {"schema":{"type":"string","optional":false},"payload":"bar"}
-...</pre>
+&hellip;</code></pre>
 
         <p>The connectors continue to process data, so we can add data to the file and see it move through the pipeline:</p>
 
-        <pre class="brush: bash;">
-&gt; echo Another line>> test.txt</pre>
+        <pre class="line-numbers"><code class="language-bash">$ echo Another line&gt;&gt; test.txt</code></pre>
 
         <p>You should see the line appear in the console consumer output and in the sink file.</p>
 
@@ -318,7 +311,7 @@ <h4 class="anchor-heading">
 
         <p>To give you a first taste, here's how one would implement the popular <code>WordCount</code> algorithm:</p>
 
-        <pre class="line-numbers"><code class="language-bash">KStream&lt;String, String&gt; textLines = builder.stream("quickstart-events");
+        <pre class="line-numbers"><code class="language-java">KStream&lt;String, String&gt; textLines = builder.stream("quickstart-events");
 
 KTable&lt;String, Long&gt; wordCounts = textLines
             .flatMapValues(line -&gt; Arrays.asList(line.toLowerCase().split(" ")))

From a63131aab8b2425a39ce67c8304df10628fd920f Mon Sep 17 00:00:00 2001
From: "Gyeongwon, Do" <dct012@gmail.com>
Date: Tue, 6 Feb 2024 00:57:10 +0900
Subject: [PATCH 002/258] KAFKA-15717: Added KRaft support in
 LeaderEpochIntegrationTest (#15225)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 .../epoch/LeaderEpochIntegrationTest.scala    | 76 ++++++++++++-------
 .../scala/unit/kafka/utils/TestUtils.scala    |  5 ++
 2 files changed, 53 insertions(+), 28 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
index 8771413269c5e..ba0becd4eecff 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
@@ -18,33 +18,34 @@ package kafka.server.epoch
 
 import kafka.cluster.BrokerEndPoint
 import kafka.server.KafkaConfig._
-import kafka.server.{BlockingSend, KafkaServer, BrokerBlockingSender}
+import kafka.server.{BlockingSend, BrokerBlockingSender, KafkaBroker, QuorumTestHarness}
 import kafka.utils.Implicits._
 import kafka.utils.TestUtils._
-import kafka.utils.{Logging, TestUtils}
-import kafka.server.QuorumTestHarness
+import kafka.utils.{Logging, TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.Errors._
 import org.apache.kafka.common.serialization.StringSerializer
 import org.apache.kafka.common.utils.{LogContext, SystemTime}
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition
-import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic
-import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.{OffsetForLeaderPartition, OffsetForLeaderTopic, OffsetForLeaderTopicCollection}
 import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.requests.{OffsetsForLeaderEpochRequest, OffsetsForLeaderEpochResponse}
 import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.{AfterEach, Test}
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.ValueSource
 
 import scala.jdk.CollectionConverters._
-import scala.collection.Map
+import scala.collection.{Map, Seq}
 import scala.collection.mutable.ListBuffer
 
 class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
-  var brokers: ListBuffer[KafkaServer] = ListBuffer()
+  var brokers: ListBuffer[KafkaBroker] = ListBuffer()
   val topic1 = "foo"
   val topic2 = "bar"
   val t1p0 = new TopicPartition(topic1, 0)
@@ -63,13 +64,14 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     super.tearDown()
   }
 
-  @Test
-  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader(): Unit = {
-    brokers ++= (0 to 1).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader(quorum: String): Unit = {
+    brokers ++= (0 to 1).map { id => createBroker(fromProps(createBrokerConfig(id, zkConnectOrNull))) }
 
     // Given two topics with replication of a single partition
     for (topic <- List(topic1, topic2)) {
-      createTopic(zkClient, topic, Map(0 -> Seq(0, 1)), servers = brokers)
+      createTopic(topic, Map(0 -> Seq(0, 1)))
     }
 
     // When we send four messages
@@ -95,17 +97,18 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     waitUntilTrue(() => messagesHaveLeaderEpoch(brokers(0), expectedLeaderEpoch, 4), "Leader epoch should be 1")
   }
 
-  @Test
-  def shouldSendLeaderEpochRequestAndGetAResponse(): Unit = {
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def shouldSendLeaderEpochRequestAndGetAResponse(quorum: String): Unit = {
 
     //3 brokers, put partition on 100/101 and then pretend to be 102
-    brokers ++= (100 to 102).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+    brokers ++= (100 to 102).map { id => createBroker(fromProps(createBrokerConfig(id, zkConnectOrNull))) }
 
     val assignment1 = Map(0 -> Seq(100), 1 -> Seq(101))
-    TestUtils.createTopic(zkClient, topic1, assignment1, brokers)
+    createTopic(topic1, assignment1)
 
     val assignment2 = Map(0 -> Seq(100))
-    TestUtils.createTopic(zkClient, topic2, assignment2, brokers)
+    createTopic(topic2, assignment2)
 
     //Send messages equally to the two partitions, then half as many to a third
     producer = createProducer(plaintextBootstrapServers(brokers), acks = -1)
@@ -142,17 +145,22 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     fetcher1.close()
   }
 
-  @Test
-  def shouldIncreaseLeaderEpochBetweenLeaderRestarts(): Unit = {
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def shouldIncreaseLeaderEpochBetweenLeaderRestarts(quorum: String): Unit = {
     //Setup: we are only interested in the single partition on broker 101
-    brokers += createServer(fromProps(createBrokerConfig(100, zkConnect)))
-    assertEquals(100, TestUtils.waitUntilControllerElected(zkClient))
+    brokers += createBroker(fromProps(createBrokerConfig(100, zkConnectOrNull)))
+    if (isKRaftTest()) {
+      assertEquals(controllerServer.config.nodeId, TestUtils.waitUntilQuorumLeaderElected(controllerServer))
+    } else {
+      assertEquals(100, TestUtils.waitUntilControllerElected(zkClient))
+    }
 
-    brokers += createServer(fromProps(createBrokerConfig(101, zkConnect)))
+    brokers += createBroker(fromProps(createBrokerConfig(101, zkConnectOrNull)))
 
     def leo() = brokers(1).replicaManager.localLog(tp).get.logEndOffset
 
-    TestUtils.createTopic(zkClient, tp.topic, Map(tp.partition -> Seq(101)), brokers)
+    createTopic(tp.topic, Map(tp.partition -> Seq(101)))
     producer = createProducer(plaintextBootstrapServers(brokers), acks = -1)
 
     //1. Given a single message
@@ -232,7 +240,7 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     assertEquals(-1, fetcher.leaderOffsetsFor(epoch5)(t1p0).endOffset())
   }
 
-  private def sender(from: KafkaServer, to: KafkaServer): BlockingSend = {
+  private def sender(from: KafkaBroker, to: KafkaBroker): BlockingSend = {
     val node = from.metadataCache.getAliveBrokerNode(to.config.brokerId,
       from.config.interBrokerListenerName).get
     val endPoint = new BrokerEndPoint(node.id(), node.host(), node.port())
@@ -245,13 +253,13 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     }, "Epoch didn't change")
   }
 
-  private  def messagesHaveLeaderEpoch(broker: KafkaServer, expectedLeaderEpoch: Int, minOffset: Int): Boolean = {
+  private  def messagesHaveLeaderEpoch(broker: KafkaBroker, expectedLeaderEpoch: Int, minOffset: Int): Boolean = {
     var result = true
     for (topic <- List(topic1, topic2)) {
       val tp = new TopicPartition(topic, 0)
-      val leo = broker.getLogManager.getLog(tp).get.logEndOffset
+      val leo = broker.logManager.getLog(tp).get.logEndOffset
       result = result && leo > 0 && brokers.forall { broker =>
-        broker.getLogManager.getLog(tp).get.logSegments.stream.allMatch { segment =>
+        broker.logManager.getLog(tp).get.logSegments.stream.allMatch { segment =>
           if (segment.read(minOffset, Integer.MAX_VALUE) == null) {
             false
           } else {
@@ -277,6 +285,18 @@ class LeaderEpochIntegrationTest extends QuorumTestHarness with Logging {
     producer.close()
   }
 
+  private def createTopic(topic: String, partitionReplicaAssignment: collection.Map[Int, Seq[Int]]): Unit = {
+    resource(createAdminClient(brokers, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))) { admin =>
+      TestUtils.createTopicWithAdmin(
+        admin = admin,
+        topic = topic,
+        replicaAssignment = partitionReplicaAssignment,
+        brokers = brokers,
+        controllers = controllerServers
+      )
+    }
+  }
+
   /**
     * Simulates how the Replica Fetcher Thread requests leader offsets for epochs
     */
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 16a7e95c0bfbd..19eac0bcfdb8d 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -1341,6 +1341,11 @@ object TestUtils extends Logging {
     controllerId.getOrElse(throw new AssertionError(s"Controller not elected after $timeout ms"))
   }
 
+  def waitUntilQuorumLeaderElected(controllerServer: ControllerServer, timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
+    val (leaderAndEpoch, _) = computeUntilTrue(controllerServer.raftManager.leaderAndEpoch, waitTime = timeout)(_.leaderId().isPresent)
+    leaderAndEpoch.leaderId().orElseThrow(() => new AssertionError(s"Quorum Controller leader not elected after $timeout ms"))
+  }
+
   def awaitLeaderChange[B <: KafkaBroker](
       brokers: Seq[B],
       tp: TopicPartition,

From 24a79c2613e32b1d1615f6f34d63225259740d0e Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Mon, 5 Feb 2024 13:55:57 -0500
Subject: [PATCH 003/258] KAFKA-14683: Migrate WorkerSinkTaskTest to Mockito
 (1/3)  (#14663)

Reviewers: Greg Harris <greg.harris@aiven.io>
---
 .../runtime/WorkerSinkTaskMockitoTest.java    | 672 ++++++++++++++++++
 .../connect/runtime/WorkerSinkTaskTest.java   | 488 -------------
 2 files changed, 672 insertions(+), 488 deletions(-)
 create mode 100644 connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java

diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
new file mode 100644
index 0000000000000..b41f571d6357e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
@@ -0,0 +1,672 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.runtime;
+
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.WorkerSinkTask.SinkTaskMetricsGroup;
+import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
+import org.apache.kafka.connect.runtime.errors.ErrorReporter;
+import org.apache.kafka.connect.runtime.errors.ProcessingContext;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
+import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.junit.MockitoRule;
+import org.mockito.quality.Strictness;
+import org.mockito.stubbing.Answer;
+import org.mockito.stubbing.OngoingStubbing;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Supplier;
+
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.StrictStubs.class)
+public class WorkerSinkTaskMockitoTest {
+    // These are fixed to keep this code simpler. In this example we assume byte[] raw values
+    // with mix of integer/string in Connect
+    private static final String TOPIC = "test";
+    private static final int PARTITION = 12;
+    private static final int PARTITION2 = 13;
+    private static final int PARTITION3 = 14;
+    private static final long FIRST_OFFSET = 45;
+    private static final Schema KEY_SCHEMA = Schema.INT32_SCHEMA;
+    private static final int KEY = 12;
+    private static final Schema VALUE_SCHEMA = Schema.STRING_SCHEMA;
+    private static final String VALUE = "VALUE";
+    private static final byte[] RAW_KEY = "key".getBytes();
+    private static final byte[] RAW_VALUE = "value".getBytes();
+
+    private static final TopicPartition TOPIC_PARTITION = new TopicPartition(TOPIC, PARTITION);
+    private static final TopicPartition TOPIC_PARTITION2 = new TopicPartition(TOPIC, PARTITION2);
+    private static final TopicPartition TOPIC_PARTITION3 = new TopicPartition(TOPIC, PARTITION3);
+
+    private static final Set<TopicPartition> INITIAL_ASSIGNMENT =
+            new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2));
+
+    private static final Map<String, String> TASK_PROPS = new HashMap<>();
+
+    static {
+        TASK_PROPS.put(SinkConnector.TOPICS_CONFIG, TOPIC);
+        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, SinkTask.class.getName());
+    }
+
+    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
+
+    private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private ConnectorTaskId taskId1 = new ConnectorTaskId("job", 1);
+    private TargetState initialState = TargetState.STARTED;
+    private MockTime time;
+    private WorkerSinkTask workerTask;
+    @Mock
+    private SinkTask sinkTask;
+    private ArgumentCaptor<WorkerSinkTaskContext> sinkTaskContext = ArgumentCaptor.forClass(WorkerSinkTaskContext.class);
+    private WorkerConfig workerConfig;
+    private MockConnectMetrics metrics;
+    @Mock
+    private PluginClassLoader pluginLoader;
+    @Mock
+    private Converter keyConverter;
+    @Mock
+    private Converter valueConverter;
+    @Mock
+    private HeaderConverter headerConverter;
+    @Mock
+    private TransformationChain<ConsumerRecord<byte[], byte[]>, SinkRecord> transformationChain;
+    @Mock
+    private TaskStatus.Listener statusListener;
+    @Mock
+    private StatusBackingStore statusBackingStore;
+    @Mock
+    private KafkaConsumer<byte[], byte[]> consumer;
+    @Mock
+    private ErrorHandlingMetrics errorHandlingMetrics;
+    private ArgumentCaptor<ConsumerRebalanceListener> rebalanceListener = ArgumentCaptor.forClass(ConsumerRebalanceListener.class);
+    @Rule
+    public final MockitoRule rule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS);
+
+    private long recordsReturnedTp1;
+    private long recordsReturnedTp3;
+
+    @Before
+    public void setUp() {
+        time = new MockTime();
+        Map<String, String> workerProps = new HashMap<>();
+        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerConfig = new StandaloneConfig(workerProps);
+        metrics = new MockConnectMetrics(time);
+        recordsReturnedTp1 = 0;
+        recordsReturnedTp3 = 0;
+    }
+
+    private void createTask(TargetState initialState) {
+        createTask(initialState, keyConverter, valueConverter, headerConverter);
+    }
+
+    private void createTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter) {
+        createTask(initialState, keyConverter, valueConverter, headerConverter, RetryWithToleranceOperatorTest.noopOperator(), Collections::emptyList);
+    }
+
+    private void createTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter,
+                            RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator,
+                            Supplier<List<ErrorReporter<ConsumerRecord<byte[], byte[]>>>> errorReportersSupplier) {
+        workerTask = new WorkerSinkTask(
+                taskId, sinkTask, statusListener, initialState, workerConfig, ClusterConfigState.EMPTY, metrics,
+                keyConverter, valueConverter, errorHandlingMetrics, headerConverter,
+                transformationChain, consumer, pluginLoader, time,
+                retryWithToleranceOperator, null, statusBackingStore, errorReportersSupplier);
+    }
+
+    @After
+    public void tearDown() {
+        if (metrics != null) metrics.stop();
+    }
+
+    @Test
+    public void testStartPaused() {
+        createTask(TargetState.PAUSED);
+
+        expectPollInitialAssignment();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+
+        time.sleep(10000L);
+        verify(consumer).pause(INITIAL_ASSIGNMENT);
+
+        assertSinkMetricValue("partition-count", 2);
+        assertTaskMetricValue("status", "paused");
+        assertTaskMetricValue("running-ratio", 0.0);
+        assertTaskMetricValue("pause-ratio", 1.0);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
+    }
+
+    @Test
+    public void testPause() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1))
+                // Pause
+                .thenThrow(new WakeupException())
+                // Offset commit as requested when pausing; No records returned by consumer.poll()
+                .thenAnswer(expectConsumerPoll(0))
+                // And unpause
+                .thenThrow(new WakeupException())
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // initial assignment
+        verifyPollInitialAssignment();
+
+        workerTask.iteration(); // fetch some data
+        // put should've been called twice now (initial assignment & poll)
+        verify(sinkTask, times(2)).put(anyList());
+
+        workerTask.transitionTo(TargetState.PAUSED);
+        time.sleep(10_000L);
+
+        assertSinkMetricValue("partition-count", 2);
+        assertSinkMetricValue("sink-record-read-total", 1.0);
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        assertSinkMetricValue("sink-record-active-count", 1.0);
+        assertSinkMetricValue("sink-record-active-count-max", 1.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.333333);
+        assertSinkMetricValue("offset-commit-seq-no", 0.0);
+        assertSinkMetricValue("offset-commit-completion-rate", 0.0);
+        assertSinkMetricValue("offset-commit-completion-total", 0.0);
+        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 1.0);
+        assertTaskMetricValue("batch-size-avg", 0.5);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
+
+        workerTask.iteration(); // wakeup
+        // Pause
+        verify(statusListener).onPause(taskId);
+        verify(consumer).pause(INITIAL_ASSIGNMENT);
+        verify(consumer).wakeup();
+
+        // Offset commit as requested when pausing; No records returned by consumer.poll()
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+
+        workerTask.iteration(); // now paused
+        time.sleep(30000L);
+
+        assertSinkMetricValue("offset-commit-seq-no", 1.0);
+        assertSinkMetricValue("offset-commit-completion-rate", 0.0333);
+        assertSinkMetricValue("offset-commit-completion-total", 1.0);
+        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "paused");
+        assertTaskMetricValue("running-ratio", 0.25);
+        assertTaskMetricValue("pause-ratio", 0.75);
+        verify(sinkTask, times(3)).put(anyList());
+
+        workerTask.transitionTo(TargetState.STARTED);
+        workerTask.iteration(); // wakeup
+        workerTask.iteration(); // now unpaused
+
+        // And unpause
+        verify(statusListener).onResume(taskId);
+        verify(consumer, times(2)).wakeup();
+        INITIAL_ASSIGNMENT.forEach(tp -> {
+            verify(consumer).resume(Collections.singleton(tp));
+        });
+        verify(sinkTask, times(4)).put(anyList());
+    }
+
+    @Test
+    public void testShutdown() throws Exception {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+        sinkTaskContext.getValue().requestCommit(); // Force an offset commit
+
+        // second iteration
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+
+        workerTask.iteration();
+        verify(sinkTask, times(2)).put(anyList());
+
+        doAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+            rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+            return null;
+        }).when(consumer).close();
+
+        workerTask.stop();
+        verify(consumer).wakeup();
+
+        workerTask.close();
+        verify(sinkTask).stop();
+        verify(consumer).close();
+        verify(headerConverter).close();
+    }
+
+    @Test
+    public void testErrorInRebalancePartitionLoss() {
+        RuntimeException exception = new RuntimeException("Revocation error");
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                });
+
+        doThrow(exception).when(sinkTask).close(INITIAL_ASSIGNMENT);
+
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+
+        try {
+            workerTask.iteration();
+            fail("Poll should have raised the rebalance exception");
+        } catch (RuntimeException e) {
+            assertEquals(exception, e);
+        }
+    }
+
+    @Test
+    public void testErrorInRebalancePartitionRevocation() {
+        RuntimeException exception = new RuntimeException("Revocation error");
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                });
+
+        expectRebalanceRevocationError(exception);
+
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+        try {
+            workerTask.iteration();
+            fail("Poll should have raised the rebalance exception");
+        } catch (RuntimeException e) {
+            assertEquals(exception, e);
+        }
+    }
+
+    @Test
+    public void testErrorInRebalancePartitionAssignment() {
+        RuntimeException exception = new RuntimeException("Assignment error");
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                });
+
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+
+        expectRebalanceAssignmentError(exception);
+        try {
+            workerTask.iteration();
+            fail("Poll should have raised the rebalance exception");
+        } catch (RuntimeException e) {
+            assertEquals(exception, e);
+        } finally {
+            verify(sinkTask).close(INITIAL_ASSIGNMENT);
+        }
+    }
+
+    @Test
+    public void testPartialRevocationAndAssignment() {
+        createTask(initialState);
+
+        when(consumer.assignment())
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(Collections.singleton(TOPIC_PARTITION2))
+                .thenReturn(Collections.singleton(TOPIC_PARTITION2))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3)))
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT);
+
+        INITIAL_ASSIGNMENT.forEach(tp -> when(consumer.position(tp)).thenReturn(FIRST_OFFSET));
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        when(consumer.poll(any(Duration.class)))
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                })
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
+                    return ConsumerRecords.empty();
+                })
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
+                    return ConsumerRecords.empty();
+                })
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsLost(Collections.singleton(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION));
+                    return ConsumerRecords.empty();
+                });
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        when(consumer.position(TOPIC_PARTITION3)).thenReturn(FIRST_OFFSET);
+
+        // First iteration--first call to poll, first consumer assignment
+        workerTask.iteration();
+        verifyPollInitialAssignment();
+
+        // Second iteration--second call to poll, partial consumer revocation
+        workerTask.iteration();
+        verify(sinkTask).close(Collections.singleton(TOPIC_PARTITION));
+        verify(sinkTask, times(2)).put(Collections.emptyList());
+
+        // Third iteration--third call to poll, partial consumer assignment
+        workerTask.iteration();
+        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION3));
+        verify(sinkTask, times(3)).put(Collections.emptyList());
+
+        // Fourth iteration--fourth call to poll, one partition lost; can't commit offsets for it, one new partition assigned
+        workerTask.iteration();
+        verify(sinkTask).close(Collections.singleton(TOPIC_PARTITION3));
+        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION));
+        verify(sinkTask, times(4)).put(Collections.emptyList());
+    }
+
+    @Test
+    public void testMetricsGroup() {
+        SinkTaskMetricsGroup group = new SinkTaskMetricsGroup(taskId, metrics);
+        SinkTaskMetricsGroup group1 = new SinkTaskMetricsGroup(taskId1, metrics);
+        for (int i = 0; i != 10; ++i) {
+            group.recordRead(1);
+            group.recordSend(2);
+            group.recordPut(3);
+            group.recordPartitionCount(4);
+            group.recordOffsetSequenceNumber(5);
+        }
+        Map<TopicPartition, OffsetAndMetadata> committedOffsets = new HashMap<>();
+        committedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        group.recordCommittedOffsets(committedOffsets);
+        Map<TopicPartition, OffsetAndMetadata> consumedOffsets = new HashMap<>();
+        consumedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 10));
+        group.recordConsumedOffsets(consumedOffsets);
+
+        for (int i = 0; i != 20; ++i) {
+            group1.recordRead(1);
+            group1.recordSend(2);
+            group1.recordPut(30);
+            group1.recordPartitionCount(40);
+            group1.recordOffsetSequenceNumber(50);
+        }
+        committedOffsets = new HashMap<>();
+        committedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        committedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 3));
+        group1.recordCommittedOffsets(committedOffsets);
+        consumedOffsets = new HashMap<>();
+        consumedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 20));
+        consumedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 30));
+        group1.recordConsumedOffsets(consumedOffsets);
+
+        assertEquals(0.333, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-read-rate"), 0.001d);
+        assertEquals(0.667, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-send-rate"), 0.001d);
+        assertEquals(9, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-active-count"), 0.001d);
+        assertEquals(4, metrics.currentMetricValueAsDouble(group.metricGroup(), "partition-count"), 0.001d);
+        assertEquals(5, metrics.currentMetricValueAsDouble(group.metricGroup(), "offset-commit-seq-no"), 0.001d);
+        assertEquals(3, metrics.currentMetricValueAsDouble(group.metricGroup(), "put-batch-max-time-ms"), 0.001d);
+
+        // Close the group
+        group.close();
+
+        for (MetricName metricName : group.metricGroup().metrics().metrics().keySet()) {
+            // Metrics for this group should no longer exist
+            assertFalse(group.metricGroup().groupId().includes(metricName));
+        }
+        // Sensors for this group should no longer exist
+        assertNull(group.metricGroup().metrics().getSensor("source-record-poll"));
+        assertNull(group.metricGroup().metrics().getSensor("source-record-write"));
+        assertNull(group.metricGroup().metrics().getSensor("poll-batch-time"));
+
+        assertEquals(0.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-read-rate"), 0.001d);
+        assertEquals(1.333, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-send-rate"), 0.001d);
+        assertEquals(45, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-active-count"), 0.001d);
+        assertEquals(40, metrics.currentMetricValueAsDouble(group1.metricGroup(), "partition-count"), 0.001d);
+        assertEquals(50, metrics.currentMetricValueAsDouble(group1.metricGroup(), "offset-commit-seq-no"), 0.001d);
+        assertEquals(30, metrics.currentMetricValueAsDouble(group1.metricGroup(), "put-batch-max-time-ms"), 0.001d);
+    }
+
+    private void expectRebalanceRevocationError(RuntimeException e) {
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+        doThrow(e).when(sinkTask).close(INITIAL_ASSIGNMENT);
+    }
+
+    private void expectRebalanceAssignmentError(RuntimeException e) {
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+        when(consumer.position(TOPIC_PARTITION)).thenReturn(FIRST_OFFSET);
+        when(consumer.position(TOPIC_PARTITION2)).thenReturn(FIRST_OFFSET);
+
+        doThrow(e).when(sinkTask).open(INITIAL_ASSIGNMENT);
+    }
+
+    private void verifyInitializeTask() {
+        verify(consumer).subscribe(eq(asList(TOPIC)), rebalanceListener.capture());
+        verify(sinkTask).initialize(sinkTaskContext.capture());
+        verify(sinkTask).start(TASK_PROPS);
+    }
+
+    private OngoingStubbing<ConsumerRecords<byte[], byte[]>> expectPollInitialAssignment() {
+        when(consumer.assignment()).thenReturn(INITIAL_ASSIGNMENT);
+        INITIAL_ASSIGNMENT.forEach(tp -> when(consumer.position(tp)).thenReturn(FIRST_OFFSET));
+
+        return when(consumer.poll(any(Duration.class))).thenAnswer(
+                invocation -> {
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                }
+        );
+    }
+
+    private void verifyPollInitialAssignment() {
+        verify(sinkTask).open(INITIAL_ASSIGNMENT);
+        verify(consumer, atLeastOnce()).assignment();
+        verify(sinkTask).put(Collections.emptyList());
+    }
+
+    private Answer<ConsumerRecords<byte[], byte[]>> expectConsumerPoll(final int numMessages) {
+        return expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, new RecordHeaders());
+    }
+
+    private Answer<ConsumerRecords<byte[], byte[]>> expectConsumerPoll(final int numMessages, final long timestamp, final TimestampType timestampType, Headers headers) {
+        return invocation -> {
+            List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
+            for (int i = 0; i < numMessages; i++)
+                records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + i, timestamp, timestampType,
+                        0, 0, RAW_KEY, RAW_VALUE, headers, Optional.empty()));
+            recordsReturnedTp1 += numMessages;
+            return new ConsumerRecords<>(
+                    numMessages > 0 ?
+                            Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records)
+                            : Collections.emptyMap()
+            );
+        };
+    }
+
+    private void expectConversionAndTransformation(final String topicPrefix, final Headers headers) {
+        when(keyConverter.toConnectData(TOPIC, headers, RAW_KEY)).thenReturn(new SchemaAndValue(KEY_SCHEMA, KEY));
+        when(valueConverter.toConnectData(TOPIC, headers, RAW_VALUE)).thenReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE));
+
+        for (Header header : headers) {
+            when(headerConverter.toConnectHeader(TOPIC, header.key(), header.value())).thenReturn(new SchemaAndValue(VALUE_SCHEMA, new String(header.value())));
+        }
+
+        expectTransformation(topicPrefix);
+    }
+
+    @SuppressWarnings("unchecked")
+    private void expectTransformation(final String topicPrefix) {
+        when(transformationChain.apply(any(ProcessingContext.class), any(SinkRecord.class))).thenAnswer((Answer<SinkRecord>)
+                invocation -> {
+                    SinkRecord origRecord = invocation.getArgument(1);
+                    return topicPrefix != null && !topicPrefix.isEmpty()
+                            ? origRecord.newRecord(
+                            topicPrefix + origRecord.topic(),
+                            origRecord.kafkaPartition(),
+                            origRecord.keySchema(),
+                            origRecord.key(),
+                            origRecord.valueSchema(),
+                            origRecord.value(),
+                            origRecord.timestamp(),
+                            origRecord.headers()
+                    ) : origRecord;
+                });
+    }
+
+    private void expectTaskGetTopic() {
+        when(statusBackingStore.getTopic(anyString(), anyString())).thenAnswer((Answer<TopicStatus>) invocation -> {
+            String connector = invocation.getArgument(0, String.class);
+            String topic = invocation.getArgument(1, String.class);
+            return new TopicStatus(topic, new ConnectorTaskId(connector, 0), Time.SYSTEM.milliseconds());
+        });
+    }
+
+    private void assertSinkMetricValue(String name, double expected) {
+        MetricGroup sinkTaskGroup = workerTask.sinkTaskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertTaskMetricValue(String name, double expected) {
+        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(taskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertTaskMetricValue(String name, String expected) {
+        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
+        String measured = metrics.currentMetricValueAsString(taskGroup, name);
+        assertEquals(expected, measured);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
index d2582501934a9..66edf6a07170c 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
@@ -22,7 +22,6 @@
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
-import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.header.Header;
@@ -38,7 +37,6 @@
 import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
-import org.apache.kafka.connect.runtime.WorkerSinkTask.SinkTaskMetricsGroup;
 import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
 import org.apache.kafka.connect.runtime.errors.ErrorReporter;
 import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
@@ -205,175 +203,6 @@ public void tearDown() {
         if (metrics != null) metrics.stop();
     }
 
-    @Test
-    public void testStartPaused() throws Exception {
-        createTask(TargetState.PAUSED);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        consumer.pause(INITIAL_ASSIGNMENT);
-        PowerMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        time.sleep(10000L);
-
-        assertSinkMetricValue("partition-count", 2);
-        assertTaskMetricValue("status", "paused");
-        assertTaskMetricValue("running-ratio", 0.0);
-        assertTaskMetricValue("pause-ratio", 1.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPause() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        // Pause
-        statusListener.onPause(taskId);
-        EasyMock.expectLastCall();
-        expectConsumerWakeup();
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        consumer.pause(INITIAL_ASSIGNMENT);
-        PowerMock.expectLastCall();
-
-        // Offset commit as requested when pausing; No records returned by consumer.poll()
-        sinkTask.preCommit(EasyMock.anyObject());
-        EasyMock.expectLastCall().andStubReturn(Collections.emptyMap());
-        expectConsumerPoll(0);
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        // And unpause
-        statusListener.onResume(taskId);
-        EasyMock.expectLastCall();
-        expectConsumerWakeup();
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            consumer.resume(Collections.singleton(tp));
-            PowerMock.expectLastCall();
-        });
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // initial assignment
-        workerTask.iteration(); // fetch some data
-        workerTask.transitionTo(TargetState.PAUSED);
-        time.sleep(10000L);
-
-        assertSinkMetricValue("partition-count", 2);
-        assertSinkMetricValue("sink-record-read-total", 1.0);
-        assertSinkMetricValue("sink-record-send-total", 1.0);
-        assertSinkMetricValue("sink-record-active-count", 1.0);
-        assertSinkMetricValue("sink-record-active-count-max", 1.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.333333);
-        assertSinkMetricValue("offset-commit-seq-no", 0.0);
-        assertSinkMetricValue("offset-commit-completion-rate", 0.0);
-        assertSinkMetricValue("offset-commit-completion-total", 0.0);
-        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 1.0);
-        assertTaskMetricValue("batch-size-avg", 0.5);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
-
-        workerTask.iteration(); // wakeup
-        workerTask.iteration(); // now paused
-        time.sleep(30000L);
-
-        assertSinkMetricValue("offset-commit-seq-no", 1.0);
-        assertSinkMetricValue("offset-commit-completion-rate", 0.0333);
-        assertSinkMetricValue("offset-commit-completion-total", 1.0);
-        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "paused");
-        assertTaskMetricValue("running-ratio", 0.25);
-        assertTaskMetricValue("pause-ratio", 0.75);
-
-        workerTask.transitionTo(TargetState.STARTED);
-        workerTask.iteration(); // wakeup
-        workerTask.iteration(); // now unpaused
-        //printMetrics();
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testShutdown() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        // first iteration
-        expectPollInitialAssignment();
-
-        // second iteration
-        EasyMock.expect(sinkTask.preCommit(EasyMock.anyObject())).andReturn(Collections.emptyMap());
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        // WorkerSinkTask::stop
-        consumer.wakeup();
-        PowerMock.expectLastCall();
-        sinkTask.stop();
-        PowerMock.expectLastCall();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        // WorkerSinkTask::close
-        consumer.close();
-        PowerMock.expectLastCall().andAnswer(() -> {
-            rebalanceListener.getValue().onPartitionsRevoked(
-                INITIAL_ASSIGNMENT
-            );
-            return null;
-        });
-        headerConverter.close();
-        PowerMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        sinkTaskContext.getValue().requestCommit(); // Force an offset commit
-        workerTask.iteration();
-        workerTask.stop();
-        workerTask.close();
-
-        PowerMock.verifyAll();
-    }
-
     @Test
     public void testPollRedelivery() throws Exception {
         createTask(initialState);
@@ -550,152 +379,6 @@ public void testPollRedeliveryWithConsumerRebalance() throws Exception {
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testErrorInRebalancePartitionLoss() throws Exception {
-        RuntimeException exception = new RuntimeException("Revocation error");
-
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-        expectRebalanceLossError(exception);
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
-        }
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testErrorInRebalancePartitionRevocation() throws Exception {
-        RuntimeException exception = new RuntimeException("Revocation error");
-
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-        expectRebalanceRevocationError(exception);
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
-        }
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testErrorInRebalancePartitionAssignment() throws Exception {
-        RuntimeException exception = new RuntimeException("Assignment error");
-
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-        expectRebalanceAssignmentError(exception);
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
-        }
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPartialRevocationAndAssignment() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
-                rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
-                return ConsumerRecords.empty();
-            });
-        EasyMock.expect(consumer.assignment()).andReturn(Collections.singleton(TOPIC_PARTITION)).times(2);
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-        sinkTask.close(Collections.singleton(TOPIC_PARTITION));
-        EasyMock.expectLastCall();
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
-                rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
-                return ConsumerRecords.empty();
-            });
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET);
-        sinkTask.open(Collections.singleton(TOPIC_PARTITION3));
-        EasyMock.expectLastCall();
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsLost(Collections.singleton(TOPIC_PARTITION3));
-                rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION));
-                return ConsumerRecords.empty();
-            });
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(3);
-        sinkTask.close(Collections.singleton(TOPIC_PARTITION3));
-        EasyMock.expectLastCall();
-        EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET);
-        sinkTask.open(Collections.singleton(TOPIC_PARTITION));
-        EasyMock.expectLastCall();
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        // First iteration--first call to poll, first consumer assignment
-        workerTask.iteration();
-        // Second iteration--second call to poll, partial consumer revocation
-        workerTask.iteration();
-        // Third iteration--third call to poll, partial consumer assignment
-        workerTask.iteration();
-        // Fourth iteration--fourth call to poll, one partition lost; can't commit offsets for it, one new partition assigned
-        workerTask.iteration();
-
-        PowerMock.verifyAll();
-    }
-
     @Test
     public void testPreCommitFailureAfterPartialRevocationAndAssignment() throws Exception {
         createTask(initialState);
@@ -1783,67 +1466,6 @@ public void testTopicsRegex() {
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testMetricsGroup() {
-        SinkTaskMetricsGroup group = new SinkTaskMetricsGroup(taskId, metrics);
-        SinkTaskMetricsGroup group1 = new SinkTaskMetricsGroup(taskId1, metrics);
-        for (int i = 0; i != 10; ++i) {
-            group.recordRead(1);
-            group.recordSend(2);
-            group.recordPut(3);
-            group.recordPartitionCount(4);
-            group.recordOffsetSequenceNumber(5);
-        }
-        Map<TopicPartition, OffsetAndMetadata> committedOffsets = new HashMap<>();
-        committedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        group.recordCommittedOffsets(committedOffsets);
-        Map<TopicPartition, OffsetAndMetadata> consumedOffsets = new HashMap<>();
-        consumedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 10));
-        group.recordConsumedOffsets(consumedOffsets);
-
-        for (int i = 0; i != 20; ++i) {
-            group1.recordRead(1);
-            group1.recordSend(2);
-            group1.recordPut(30);
-            group1.recordPartitionCount(40);
-            group1.recordOffsetSequenceNumber(50);
-        }
-        committedOffsets = new HashMap<>();
-        committedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 2));
-        committedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 3));
-        group1.recordCommittedOffsets(committedOffsets);
-        consumedOffsets = new HashMap<>();
-        consumedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 20));
-        consumedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 30));
-        group1.recordConsumedOffsets(consumedOffsets);
-
-        assertEquals(0.333, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-read-rate"), 0.001d);
-        assertEquals(0.667, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-send-rate"), 0.001d);
-        assertEquals(9, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-active-count"), 0.001d);
-        assertEquals(4, metrics.currentMetricValueAsDouble(group.metricGroup(), "partition-count"), 0.001d);
-        assertEquals(5, metrics.currentMetricValueAsDouble(group.metricGroup(), "offset-commit-seq-no"), 0.001d);
-        assertEquals(3, metrics.currentMetricValueAsDouble(group.metricGroup(), "put-batch-max-time-ms"), 0.001d);
-
-        // Close the group
-        group.close();
-
-        for (MetricName metricName : group.metricGroup().metrics().metrics().keySet()) {
-            // Metrics for this group should no longer exist
-            assertFalse(group.metricGroup().groupId().includes(metricName));
-        }
-        // Sensors for this group should no longer exist
-        assertNull(group.metricGroup().metrics().getSensor("source-record-poll"));
-        assertNull(group.metricGroup().metrics().getSensor("source-record-write"));
-        assertNull(group.metricGroup().metrics().getSensor("poll-batch-time"));
-
-        assertEquals(0.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-read-rate"), 0.001d);
-        assertEquals(1.333, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-send-rate"), 0.001d);
-        assertEquals(45, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-active-count"), 0.001d);
-        assertEquals(40, metrics.currentMetricValueAsDouble(group1.metricGroup(), "partition-count"), 0.001d);
-        assertEquals(50, metrics.currentMetricValueAsDouble(group1.metricGroup(), "offset-commit-seq-no"), 0.001d);
-        assertEquals(30, metrics.currentMetricValueAsDouble(group1.metricGroup(), "put-batch-max-time-ms"), 0.001d);
-    }
-
     @Test
     public void testHeaders() throws Exception {
         Headers headers = new RecordHeaders();
@@ -1994,53 +1616,6 @@ private void expectInitializeTask() {
         PowerMock.expectLastCall();
     }
 
-    private void expectRebalanceLossError(RuntimeException e) {
-        sinkTask.close(new HashSet<>(INITIAL_ASSIGNMENT));
-        EasyMock.expectLastCall().andThrow(e);
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT);
-                return ConsumerRecords.empty();
-            });
-    }
-
-    private void expectRebalanceRevocationError(RuntimeException e) {
-        sinkTask.close(new HashSet<>(INITIAL_ASSIGNMENT));
-        EasyMock.expectLastCall().andThrow(e);
-
-        sinkTask.preCommit(EasyMock.anyObject());
-        EasyMock.expectLastCall().andReturn(Collections.emptyMap());
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                return ConsumerRecords.empty();
-            });
-    }
-
-    private void expectRebalanceAssignmentError(RuntimeException e) {
-        sinkTask.close(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall();
-
-        sinkTask.preCommit(EasyMock.anyObject());
-        EasyMock.expectLastCall().andReturn(Collections.emptyMap());
-
-        EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(FIRST_OFFSET);
-
-        sinkTask.open(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall().andThrow(e);
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
-                return ConsumerRecords.empty();
-            });
-    }
-
     private void expectPollInitialAssignment() {
         sinkTask.open(INITIAL_ASSIGNMENT);
         EasyMock.expectLastCall();
@@ -2057,12 +1632,6 @@ private void expectPollInitialAssignment() {
         EasyMock.expectLastCall();
     }
 
-    private void expectConsumerWakeup() {
-        consumer.wakeup();
-        EasyMock.expectLastCall();
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andThrow(new WakeupException());
-    }
-
     private void expectConsumerPoll(final int numMessages) {
         expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, emptyHeaders());
     }
@@ -2180,63 +1749,6 @@ private void assertTaskMetricValue(String name, String expected) {
         assertEquals(expected, measured);
     }
 
-    private void printMetrics() {
-        System.out.println();
-        sinkMetricValue("sink-record-read-rate");
-        sinkMetricValue("sink-record-read-total");
-        sinkMetricValue("sink-record-send-rate");
-        sinkMetricValue("sink-record-send-total");
-        sinkMetricValue("sink-record-active-count");
-        sinkMetricValue("sink-record-active-count-max");
-        sinkMetricValue("sink-record-active-count-avg");
-        sinkMetricValue("partition-count");
-        sinkMetricValue("offset-commit-seq-no");
-        sinkMetricValue("offset-commit-completion-rate");
-        sinkMetricValue("offset-commit-completion-total");
-        sinkMetricValue("offset-commit-skip-rate");
-        sinkMetricValue("offset-commit-skip-total");
-        sinkMetricValue("put-batch-max-time-ms");
-        sinkMetricValue("put-batch-avg-time-ms");
-
-        taskMetricValue("status-unassigned");
-        taskMetricValue("status-running");
-        taskMetricValue("status-paused");
-        taskMetricValue("status-failed");
-        taskMetricValue("status-destroyed");
-        taskMetricValue("running-ratio");
-        taskMetricValue("pause-ratio");
-        taskMetricValue("offset-commit-max-time-ms");
-        taskMetricValue("offset-commit-avg-time-ms");
-        taskMetricValue("batch-size-max");
-        taskMetricValue("batch-size-avg");
-        taskMetricValue("offset-commit-failure-percentage");
-        taskMetricValue("offset-commit-success-percentage");
-    }
-
-    private double sinkMetricValue(String metricName) {
-        MetricGroup sinkTaskGroup = workerTask.sinkTaskMetricsGroup().metricGroup();
-        double value = metrics.currentMetricValueAsDouble(sinkTaskGroup, metricName);
-        System.out.println("** " + metricName + "=" + value);
-        return value;
-    }
-
-    private double taskMetricValue(String metricName) {
-        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
-        double value = metrics.currentMetricValueAsDouble(taskGroup, metricName);
-        System.out.println("** " + metricName + "=" + value);
-        return value;
-    }
-
-
-    private void assertMetrics(int minimumPollCountExpected) {
-        MetricGroup sinkTaskGroup = workerTask.sinkTaskMetricsGroup().metricGroup();
-        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
-        double readRate = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-read-rate");
-        double readTotal = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-read-total");
-        double sendRate = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-send-rate");
-        double sendTotal = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-send-total");
-    }
-
     private RecordHeaders emptyHeaders() {
         return new RecordHeaders();
     }

From 5f35b41e92c502c731775a2b417d30eae9a48e0e Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Mon, 5 Feb 2024 11:39:47 -0800
Subject: [PATCH 004/258] KAFKA-15834: Remove more leaky
 NamedTopologyIntegrationTest tests (#15243)

Signed-off-by: Greg Harris <greg.harris@aiven.io>
Reviewers: Anna Sophie Blee-Goldman <sophie@responsive.dev>
---
 .../NamedTopologyIntegrationTest.java         | 157 ------------------
 1 file changed, 157 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/NamedTopologyIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/NamedTopologyIntegrationTest.java
index e1d123526897e..9d9c77850e2d9 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/NamedTopologyIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/NamedTopologyIntegrationTest.java
@@ -18,8 +18,6 @@
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.serialization.IntegerDeserializer;
-import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.LongDeserializer;
 import org.apache.kafka.common.serialization.LongSerializer;
 import org.apache.kafka.common.serialization.Serdes;
@@ -36,7 +34,6 @@
 import org.apache.kafka.streams.errors.MissingSourceTopicException;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.errors.StreamsUncaughtExceptionHandler;
-import org.apache.kafka.streams.errors.StreamsUncaughtExceptionHandler.StreamThreadExceptionResponse;
 import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
 import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
 import org.apache.kafka.streams.kstream.Consumed;
@@ -58,7 +55,6 @@
 import org.apache.kafka.streams.utils.UniqueTopicSerdeScope;
 import org.apache.kafka.test.TestUtils;
 
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedList;
@@ -81,7 +77,6 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -92,17 +87,13 @@
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.safeUniqueTestName;
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitForApplicationState;
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.waitUntilMinKeyValueRecordsReceived;
-import static org.apache.kafka.streams.processor.internals.ClientUtils.extractThreadId;
-import static org.apache.kafka.test.TestUtils.retryOnExceptionWithTimeout;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
-import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singleton;
-import static java.util.Collections.singletonList;
 
 @Timeout(600)
 @Tag("integration")
@@ -550,31 +541,6 @@ public void shouldAllowRemovingAndAddingNamedTopologyToRunningApplicationWithMul
         assertThat(waitUntilMinKeyValueRecordsReceived(consumerConfig, OUTPUT_STREAM_2, 5), equalTo(COUNT_OUTPUT_DATA));
     }
 
-    @Test
-    public void shouldRemoveOneNamedTopologyWhileAnotherContinuesProcessing() throws Exception {
-        CLUSTER.createTopic(DELAYED_INPUT_STREAM_1, 2, 1);
-        CLUSTER.createTopic(DELAYED_INPUT_STREAM_2, 2, 1);
-
-        try {
-            topology1Builder.stream(DELAYED_INPUT_STREAM_1).groupBy((k, v) -> k).count(IN_MEMORY_STORE).toStream().to(OUTPUT_STREAM_1);
-            topology2Builder.stream(DELAYED_INPUT_STREAM_2).map((k, v) -> {
-                throw new IllegalStateException("Should not process any records for removed topology-2");
-            });
-            streams.addNamedTopology(topology1Builder.build());
-            streams.addNamedTopology(topology2Builder.build());
-            IntegrationTestUtils.startApplicationAndWaitUntilRunning(streams);
-
-            streams.removeNamedTopology("topology-2").all().get();
-
-            produceToInputTopics(DELAYED_INPUT_STREAM_1, STANDARD_INPUT_DATA);
-            produceToInputTopics(DELAYED_INPUT_STREAM_2, STANDARD_INPUT_DATA);
-
-            assertThat(waitUntilMinKeyValueRecordsReceived(consumerConfig, OUTPUT_STREAM_1, 5), equalTo(COUNT_OUTPUT_DATA));
-        } finally {
-            CLUSTER.deleteTopics(DELAYED_INPUT_STREAM_1, DELAYED_INPUT_STREAM_2);
-        }
-    }
-
     @Test
     public void shouldRemoveAndReplaceTopologicallyIncompatibleNamedTopology() throws Exception {
         CLUSTER.createTopics(SUM_OUTPUT, COUNT_OUTPUT);
@@ -716,129 +682,6 @@ public void shouldAddToEmptyInitialTopologyRemoveResetOffsetsThenAddSameNamedTop
         CLUSTER.deleteTopicsAndWait(SUM_OUTPUT, COUNT_OUTPUT);
     }
 
-    @Test
-    public void shouldWaitForMissingInputTopicsToBeCreated() throws Exception {
-        setupSecondKafkaStreams();
-        topology1Builder.stream(NEW_STREAM).groupBy((k, v) -> k).count(IN_MEMORY_STORE).toStream().to(OUTPUT_STREAM_1);
-        topology1Builder2.stream(NEW_STREAM).groupBy((k, v) -> k).count(IN_MEMORY_STORE).toStream().to(OUTPUT_STREAM_1);
-
-        final TrackingExceptionHandler handler = new  TrackingExceptionHandler();
-        streams.setUncaughtExceptionHandler(handler);
-        streams2.setUncaughtExceptionHandler(handler);
-
-        streams.addNamedTopology(topology1Builder.build());
-        streams2.addNamedTopology(topology1Builder2.build());
-        IntegrationTestUtils.startApplicationAndWaitUntilRunning(asList(streams, streams2));
-
-        retryOnExceptionWithTimeout(() -> {
-            final Throwable error = handler.nextError(TOPOLOGY_1);
-            assertThat(error, notNullValue());
-            assertThat(error.getCause().getClass(), is(MissingSourceTopicException.class));
-        });
-
-        try {
-            CLUSTER.createTopic(NEW_STREAM, 2, 1);
-            produceToInputTopics(NEW_STREAM, STANDARD_INPUT_DATA);
-
-            final List<KeyValue<String, Integer>> output =
-                waitUntilMinKeyValueRecordsReceived(consumerConfig, OUTPUT_STREAM_1, 5);
-            output.retainAll(COUNT_OUTPUT_DATA);
-
-            assertThat(output, equalTo(COUNT_OUTPUT_DATA));
-
-            // Make sure the threads were not actually killed and replaced
-            assertThat(streams.metadataForLocalThreads().size(), equalTo(2));
-            assertThat(streams2.metadataForLocalThreads().size(), equalTo(2));
-
-            final Set<String> localThreadsNames = streams.metadataForLocalThreads().stream()
-                .map(t -> extractThreadId(t.threadName()))
-                .collect(Collectors.toSet());
-            final Set<String> localThreadsNames2 = streams2.metadataForLocalThreads().stream()
-                .map(t -> extractThreadId(t.threadName()))
-                .collect(Collectors.toSet());
-
-            assertThat(localThreadsNames.contains("StreamThread-1"), is(true));
-            assertThat(localThreadsNames.contains("StreamThread-2"), is(true));
-            assertThat(localThreadsNames2.contains("StreamThread-1"), is(true));
-            assertThat(localThreadsNames2.contains("StreamThread-2"), is(true));
-        } finally {
-            CLUSTER.deleteTopicsAndWait(NEW_STREAM);
-        }
-    }
-
-    @Test
-    public void shouldBackOffTaskAndEmitDataWithinSameTopology() throws Exception {
-        CLUSTER.createTopic(DELAYED_INPUT_STREAM_1, 2, 1);
-        CLUSTER.createTopic(DELAYED_INPUT_STREAM_2, 2, 1);
-
-        try {
-            final AtomicInteger noOutputExpected = new AtomicInteger(0);
-            final AtomicInteger outputExpected = new AtomicInteger(0);
-            props.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
-            props.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 15000L);
-            props.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory(appId).getPath());
-            props.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
-            props.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
-
-            // Discard the pre-created streams and replace with test-specific topology
-            streams.close();
-            streams = new KafkaStreamsNamedTopologyWrapper(props);
-            streams.setUncaughtExceptionHandler(exception -> StreamThreadExceptionResponse.REPLACE_THREAD);
-
-            final NamedTopologyBuilder builder = streams.newNamedTopologyBuilder("topology_A");
-            builder.stream(DELAYED_INPUT_STREAM_1).peek((k, v) -> outputExpected.incrementAndGet()).to(OUTPUT_STREAM_1);
-            builder.stream(DELAYED_INPUT_STREAM_2)
-                .peek((k, v) -> {
-                    throw new RuntimeException("Kaboom");
-                })
-                .peek((k, v) -> noOutputExpected.incrementAndGet())
-                .to(OUTPUT_STREAM_2);
-
-            streams.addNamedTopology(builder.build());
-
-            IntegrationTestUtils.startApplicationAndWaitUntilRunning(streams);
-            IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp(
-                DELAYED_INPUT_STREAM_2,
-                singletonList(
-                    new KeyValue<>(1, "A")
-                ),
-                TestUtils.producerConfig(
-                    CLUSTER.bootstrapServers(),
-                    IntegerSerializer.class,
-                    StringSerializer.class,
-                    new Properties()),
-                0L);
-            IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp(
-                DELAYED_INPUT_STREAM_1,
-                Arrays.asList(
-                    new KeyValue<>(1, "A"),
-                    new KeyValue<>(1, "B")
-                ),
-                TestUtils.producerConfig(
-                    CLUSTER.bootstrapServers(),
-                    IntegerSerializer.class,
-                    StringSerializer.class,
-                    new Properties()),
-                0L);
-            IntegrationTestUtils.waitUntilFinalKeyValueRecordsReceived(
-                TestUtils.consumerConfig(
-                    CLUSTER.bootstrapServers(),
-                    IntegerDeserializer.class,
-                    StringDeserializer.class
-                ),
-                OUTPUT_STREAM_1,
-                Arrays.asList(
-                    new KeyValue<>(1, "A"),
-                    new KeyValue<>(1, "B")
-                )
-            );
-            assertThat(noOutputExpected.get(), equalTo(0));
-            assertThat(outputExpected.get(), equalTo(2));
-        } finally {
-            CLUSTER.deleteTopics(DELAYED_INPUT_STREAM_1, DELAYED_INPUT_STREAM_2);
-        }
-    }
-
     /**
      * Validates that each metadata object has only partitions & state stores for its specific topology name and
      * asserts that {@code left} and {@code right} differ only by {@link StreamsMetadata#hostInfo()}

From 54ef99f42f8aba2614287295329a68be320cd581 Mon Sep 17 00:00:00 2001
From: Jim Galasyn <jim.galasyn@confluent.io>
Date: Mon, 5 Feb 2024 12:03:15 -0800
Subject: [PATCH 005/258] Update Streams API broker compat table for AK 3.7
 (#15051)

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 docs/streams/upgrade-guide.html | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/streams/upgrade-guide.html b/docs/streams/upgrade-guide.html
index 80f53b8572014..14544077403cd 100644
--- a/docs/streams/upgrade-guide.html
+++ b/docs/streams/upgrade-guide.html
@@ -1425,7 +1425,7 @@ <h3 class="anchor-heading"><a id="streams_api_broker_compat" class="anchor-link"
             <td>Kafka Streams API (rows)</td>
             <td>0.10.0.x</td>
             <td>0.10.1.x and 0.10.2.x</td>
-            <td>0.11.0.x and<br>1.0.x and<br>1.1.x and<br>2.0.x and<br>2.1.x and<br>2.2.x and<br>2.3.x and<br>2.4.x and<br>2.5.x and<br>2.6.x and<br>2.7.x and<br>2.8.x and<br>3.0.x and<br>3.1.x and<br>3.2.x and<br>3.3.x and<br>3.4.x and<br>3.5.x and<br>3.6.x</td>
+            <td>0.11.0.x and<br>1.0.x and<br>1.1.x and<br>2.0.x and<br>2.1.x and<br>2.2.x and<br>2.3.x and<br>2.4.x and<br>2.5.x and<br>2.6.x and<br>2.7.x and<br>2.8.x and<br>3.0.x and<br>3.1.x and<br>3.2.x and<br>3.3.x and<br>3.4.x and<br>3.5.x and<br>3.6.x and<br>3.7.x</td>
           </tr>
           <tr>
             <td>0.10.0.x</td>
@@ -1452,7 +1452,7 @@ <h3 class="anchor-heading"><a id="streams_api_broker_compat" class="anchor-link"
             <td>compatible; requires message format 0.10 or higher;<br>if message headers are used, message format 0.11<br>or higher required</td>
           </tr>
           <tr>
-            <td>2.2.1 and<br>2.3.x and<br>2.4.x and<br>2.5.x and<br>2.6.x and<br>2.7.x and<br>2.8.x and<br>3.0.x and<br>3.1.x and<br>3.2.x and<br>3.3.x and<br>3.4.x and<br>3.5.x and<br>3.6.x</td>
+            <td>2.2.1 and<br>2.3.x and<br>2.4.x and<br>2.5.x and<br>2.6.x and<br>2.7.x and<br>2.8.x and<br>3.0.x and<br>3.1.x and<br>3.2.x and<br>3.3.x and<br>3.4.x and<br>3.5.x and<br>3.6.x and<br>3.7.x</td>
             <td></td>
             <td></td>
             <td>compatible; requires message format 0.11 or higher;<br>enabling exactly-once v2 requires 2.4.x or higher</td>

From 3c20d4e54a4efcd83859dcf6bda4eaa015f6c8ee Mon Sep 17 00:00:00 2001
From: Said Boudjelda <bmscomp@gmail.com>
Date: Tue, 6 Feb 2024 12:18:13 +0100
Subject: [PATCH 006/258] MINOR: Upgrade maven artifact version to 3.9.6
 (#15309)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 LICENSE-binary             | 4 ++--
 NOTICE-binary              | 2 +-
 gradle/dependencies.gradle | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/LICENSE-binary b/LICENSE-binary
index c116ca6b8e56d..de9367464889f 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -240,7 +240,7 @@ jetty-util-9.4.53.v20231009
 jetty-util-ajax-9.4.53.v20231009
 jose4j-0.9.4
 lz4-java-1.8.0
-maven-artifact-3.8.8
+maven-artifact-3.9.6
 metrics-core-4.1.12.1
 metrics-core-2.2.0
 netty-buffer-4.1.100.Final
@@ -253,7 +253,7 @@ netty-transport-classes-epoll-4.1.100.Final
 netty-transport-native-epoll-4.1.100.Final
 netty-transport-native-unix-common-4.1.100.Final
 opentelemetry-proto-1.0.0-alpha
-plexus-utils-3.3.1
+plexus-utils-3.5.1
 reflections-0.10.2
 reload4j-1.2.25
 rocksdbjni-7.9.2
diff --git a/NOTICE-binary b/NOTICE-binary
index 82920af2fc951..d3207a131e25b 100644
--- a/NOTICE-binary
+++ b/NOTICE-binary
@@ -560,7 +560,7 @@ The Apache Software Foundation (http://www.apache.org/).
 
 
 Maven Artifact
-Copyright 2001-2019 The Apache Software Foundation
+Copyright 2001-2024 The Apache Software Foundation
 
 This product includes software developed at
 The Apache Software Foundation (http://www.apache.org/).
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 49a57d1a38e7a..6067d51f3674c 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -143,7 +143,7 @@ versions += [
   kafka_35: "3.5.2",
   kafka_36: "3.6.1",
   lz4: "1.8.0",
-  mavenArtifact: "3.8.8",
+  mavenArtifact: "3.9.6",
   metrics: "2.2.0",
   netty: "4.1.100.Final",
   opentelemetryProto: "1.0.0-alpha",

From d6068189a5c41055f8cf7bb63986a63df19cef9d Mon Sep 17 00:00:00 2001
From: Mickael Maison <mimaison@users.noreply.github.com>
Date: Tue, 6 Feb 2024 19:09:54 +0100
Subject: [PATCH 007/258] MINOR: Update LICENSE-binary file (#15322)

Reviewers: Josep Prat <josep.prat@aiven.io>
---
 LICENSE-binary | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/LICENSE-binary b/LICENSE-binary
index de9367464889f..c3b370da06ea8 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -212,7 +212,7 @@ commons-cli-1.4
 commons-collections-3.2.2
 commons-digester-2.1
 commons-io-2.11.0
-commons-lang3-3.8.1
+commons-lang3-3.12.0
 commons-logging-1.2
 commons-validator-1.7
 error_prone_annotations-2.10.0
@@ -223,6 +223,7 @@ jackson-dataformat-csv-2.16.0
 jackson-datatype-jdk8-2.16.0
 jackson-jaxrs-base-2.16.0
 jackson-jaxrs-json-provider-2.16.0
+jackson-module-afterburner-2.16.0
 jackson-module-jaxb-annotations-2.16.0
 jackson-module-scala_2.13-2.16.0
 jackson-module-scala_2.12-2.16.0

From af0aceb4d718f1b647114857b374806c7c2b67a0 Mon Sep 17 00:00:00 2001
From: yicheny <siran0611@gmail.com>
Date: Wed, 7 Feb 2024 15:30:01 +0800
Subject: [PATCH 008/258] MINOR: fix word spelling mistakes (#15331)

fix word spelling mistakes

Reviewers: Luke Chen <showuon@gmail.com>
---
 raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
index e9fca1d6b2b5b..fd98774355cd6 100644
--- a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
+++ b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
@@ -975,7 +975,7 @@ private CompletableFuture<FetchResponseData> handleFetchRequest(
             || isPartitionDiverged(partitionResponse)
             || isPartitionSnapshotted(partitionResponse)) {
             // Reply immediately if any of the following is true
-            // 1. The response contains an errror
+            // 1. The response contains an error
             // 2. There are records in the response
             // 3. The fetching replica doesn't want to wait for the partition to contain new data
             // 4. The fetching replica needs to truncate because the log diverged

From 08b68583fa1042c9b9473e4dc9352a3b45c772fb Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Wed, 7 Feb 2024 01:06:13 -0800
Subject: [PATCH 009/258] =?UTF-8?q?KAFKA-16037:=20Update=20VerifiableConsu?=
 =?UTF-8?q?mer=20to=20support=20KIP-848=E2=80=99s=20group=20protocol=20con?=
 =?UTF-8?q?fig=20(#15325)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add the optional --group-protocol command line option that can be set in the system tests

Reviewers: Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../org/apache/kafka/tools/VerifiableConsumer.java    | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
index 55acea53eea49..5fe66a5998834 100644
--- a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
+++ b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
@@ -34,6 +34,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.GroupProtocol;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
@@ -528,6 +529,15 @@ private static ArgumentParser argParser() {
                 .metavar("TOPIC")
                 .help("Consumes messages from this topic.");
 
+        parser.addArgument("--group-protocol")
+                .action(store())
+                .required(false)
+                .type(String.class)
+                .setDefault(GroupProtocol.CLASSIC.name)
+                .metavar("GROUP_PROTOCOL")
+                .dest("groupProtocol")
+                .help(String.format("Group protocol (must be one of %s)", Utils.join(GroupProtocol.values(), ", ")));
+
         parser.addArgument("--group-id")
                 .action(store())
                 .required(true)
@@ -617,6 +627,7 @@ public static VerifiableConsumer createFromArgs(ArgumentParser parser, String[]
             }
         }
 
+        consumerProps.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, res.getString("groupProtocol"));
         consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, res.getString("groupId"));
 
         String groupInstanceId = res.getString("groupInstanceId");

From c000b1fae2bd7d4b76713a53508f128a13431ab6 Mon Sep 17 00:00:00 2001
From: David Arthur <mumrah@gmail.com>
Date: Wed, 7 Feb 2024 15:54:59 -0500
Subject: [PATCH 010/258] MINOR: Fix some MetadataDelta handling issues during
 ZK migration (#15327)

Reviewers: Colin P. McCabe <cmccabe@apache.org>
---
 .../migration/ZkAclMigrationClientTest.scala  | 162 +++++++++++++++++-
 .../ZkConfigMigrationClientTest.scala         |   4 +-
 .../zk/migration/ZkMigrationClientTest.scala  |   2 +-
 .../org/apache/kafka/image/AclsDelta.java     |  13 --
 .../migration/KRaftMigrationDriver.java       |   5 +-
 .../migration/KRaftMigrationZkWriter.java     |  70 ++++----
 .../migration/KRaftMigrationZkWriterTest.java |   6 +-
 7 files changed, 206 insertions(+), 56 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
index 0107b9fa95400..bb4d3e646c7b7 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
@@ -21,14 +21,14 @@ import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString
 import kafka.utils.TestUtils
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.acl._
-import org.apache.kafka.common.metadata.AccessControlEntryRecord
+import org.apache.kafka.common.metadata.{AccessControlEntryRecord, RemoveAccessControlEntryRecord}
 import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourcePatternFilter, ResourceType}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.SecurityUtils
 import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
 import org.apache.kafka.metadata.migration.KRaftMigrationZkWriter
 import org.apache.kafka.server.common.ApiMessageAndVersion
-import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue}
+import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue, fail}
 import org.junit.jupiter.api.Test
 
 import scala.collection.mutable
@@ -169,7 +169,7 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
     val image = delta.apply(MetadataProvenance.EMPTY)
 
     // load snapshot to Zookeeper.
-    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient)
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, fail(_))
     kraftMigrationZkWriter.handleSnapshot(image, (_, _, operation) => { migrationState = operation.apply(migrationState) })
 
     // Verify the new ACLs in Zookeeper.
@@ -189,4 +189,160 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
         AclPermissionType.fromCode(acl1Resource3.permissionType())),
       resource3AclsInZk.head.ace)
   }
+
+  def user(user: String): String = {
+    new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user).toString
+  }
+
+  def acl(resourceName: String,
+          resourceType: ResourceType,
+          resourcePattern: PatternType,
+          principal: String,
+          host: String = "*",
+          operation: AclOperation = AclOperation.READ,
+          permissionType: AclPermissionType = AclPermissionType.ALLOW
+  ): AccessControlEntryRecord = {
+    new AccessControlEntryRecord()
+      .setId(Uuid.randomUuid())
+      .setHost(host)
+      .setOperation(operation.code())
+      .setPrincipal(principal)
+      .setPermissionType(permissionType.code())
+      .setPatternType(resourcePattern.code())
+      .setResourceName(resourceName)
+      .setResourceType(resourceType.code())
+  }
+
+  @Test
+  def testDeleteOneAclOfMany(): Unit = {
+    zkClient.createAclPaths()
+    val topicName = "topic-" + Uuid.randomUuid()
+    val resource = new ResourcePattern(ResourceType.TOPIC, topicName, PatternType.LITERAL)
+
+    // Create a delta with some ACLs
+    val delta = new MetadataDelta(MetadataImage.EMPTY)
+    val acl1 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("alice"))
+    val acl2 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("bob"))
+    val acl3 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("carol"))
+    delta.replay(acl1)
+    delta.replay(acl2)
+    delta.replay(acl3)
+    val image = delta.apply(MetadataProvenance.EMPTY)
+
+    // Sync image to ZK
+    val errorLogs = mutable.Buffer[String]()
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, errorLogs.append)
+    kraftMigrationZkWriter.handleSnapshot(image, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+
+    // verify 3 ACLs in ZK
+    val aclsInZk = zkClient.getVersionedAclsForResource(resource).acls
+    assertEquals(3, aclsInZk.size)
+
+    // Delete one of the ACLs
+    val delta2 = new MetadataDelta.Builder()
+      .setImage(image)
+      .build()
+    delta2.replay(new RemoveAccessControlEntryRecord().setId(acl3.id()))
+    val image2 = delta2.apply(MetadataProvenance.EMPTY)
+    kraftMigrationZkWriter.handleDelta(image, image2, delta2, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+
+    // verify the other 2 ACLs are still in ZK
+    val aclsInZk2 = zkClient.getVersionedAclsForResource(resource).acls
+    assertEquals(2, aclsInZk2.size)
+    assertEquals(0, errorLogs.size)
+
+    // Add another ACL
+    val acl4 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("carol"))
+    delta2.replay(acl4)
+    val image3 = delta2.apply(MetadataProvenance.EMPTY)
+
+    // This is a contrived error case. In practice, we will never pass the same image as prev and current.
+    // The point of this is to exercise the case of a deleted ACL missing from the prev image.
+    kraftMigrationZkWriter.handleDelta(image3, image3, delta2, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+
+    val aclsInZk3 = zkClient.getVersionedAclsForResource(resource).acls
+    assertEquals(3, aclsInZk3.size)
+    assertEquals(1, errorLogs.size)
+    assertEquals(s"Cannot delete ACL ${acl3.id()} from ZK since it is missing from previous AclImage", errorLogs.head)
+  }
+
+  @Test
+  def testAclUpdateAndDelete(): Unit = {
+    zkClient.createAclPaths()
+    val errorLogs = mutable.Buffer[String]()
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, errorLogs.append)
+
+    val topicName = "topic-" + Uuid.randomUuid()
+    val otherName = "other-" + Uuid.randomUuid()
+    val literalResource = new ResourcePattern(ResourceType.TOPIC, topicName, PatternType.LITERAL)
+    val prefixedResource = new ResourcePattern(ResourceType.TOPIC, topicName, PatternType.PREFIXED)
+    val otherResource = new ResourcePattern(ResourceType.TOPIC, otherName, PatternType.LITERAL)
+
+    // Create a delta with some ACLs
+    val acl1 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("alice"))
+    val acl2 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("bob"))
+    val acl3 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("carol"))
+    val acl4 = acl(topicName, ResourceType.TOPIC, PatternType.LITERAL, user("dave"))
+
+    val delta1 = new MetadataDelta(MetadataImage.EMPTY)
+    delta1.replay(acl1)
+    delta1.replay(acl2)
+    delta1.replay(acl3)
+    delta1.replay(acl4)
+
+    val image1 = delta1.apply(MetadataProvenance.EMPTY)
+    kraftMigrationZkWriter.handleDelta(MetadataImage.EMPTY, image1, delta1, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+    assertEquals(4, zkClient.getVersionedAclsForResource(literalResource).acls.size)
+    assertEquals(0, zkClient.getVersionedAclsForResource(prefixedResource).acls.size)
+    assertEquals(0, zkClient.getVersionedAclsForResource(otherResource).acls.size)
+    assertEquals(0, errorLogs.size)
+
+    val acl5 = acl(topicName, ResourceType.TOPIC, PatternType.PREFIXED, user("alice"))
+    val acl6 = acl(topicName, ResourceType.TOPIC, PatternType.PREFIXED, user("bob"))
+    val acl7 = acl(otherName, ResourceType.TOPIC, PatternType.LITERAL, user("carol"))
+    val acl8 = acl(otherName, ResourceType.TOPIC, PatternType.LITERAL, user("dave"))
+
+    // Add two prefixed and two "other" ACLs, delete one of the literal ACLs
+    val delta2 = new MetadataDelta.Builder().setImage(image1).build()
+    delta2.replay(acl5)
+    delta2.replay(acl6)
+    delta2.replay(acl7)
+    delta2.replay(acl8)
+    delta2.replay(new RemoveAccessControlEntryRecord().setId(acl1.id()))
+
+    val image2 = delta2.apply(MetadataProvenance.EMPTY)
+    kraftMigrationZkWriter.handleDelta(image1, image2, delta2, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+    assertEquals(3, zkClient.getVersionedAclsForResource(literalResource).acls.size)
+    assertEquals(2, zkClient.getVersionedAclsForResource(prefixedResource).acls.size)
+    assertEquals(2, zkClient.getVersionedAclsForResource(otherResource).acls.size)
+    assertEquals(0, errorLogs.size)
+
+    // Delete and add ACL for literal resource, remove both prefixed ACLs, add another "other"
+    val acl9 = acl(otherName, ResourceType.TOPIC, PatternType.LITERAL, user("eve"))
+    val delta3 = new MetadataDelta.Builder().setImage(image2).build()
+    delta3.replay(acl1)
+    delta3.replay(new RemoveAccessControlEntryRecord().setId(acl2.id()))
+    delta3.replay(new RemoveAccessControlEntryRecord().setId(acl5.id()))
+    delta3.replay(new RemoveAccessControlEntryRecord().setId(acl6.id()))
+    delta3.replay(acl9)
+
+    val image3 = delta3.apply(MetadataProvenance.EMPTY)
+    kraftMigrationZkWriter.handleDelta(image2, image3, delta3, (_, _, operation) => {
+      migrationState = operation.apply(migrationState)
+    })
+    assertEquals(3, zkClient.getVersionedAclsForResource(literalResource).acls.size)
+    assertEquals(0, zkClient.getVersionedAclsForResource(prefixedResource).acls.size)
+    assertEquals(3, zkClient.getVersionedAclsForResource(otherResource).acls.size)
+    assertEquals(0, errorLogs.size)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
index 16b15386bacdd..2aff7edf7135f 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
@@ -40,7 +40,7 @@ import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.server.common.ApiMessageAndVersion
 import org.apache.kafka.server.config.ConfigType
 import org.apache.kafka.server.util.MockRandom
-import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue}
+import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue, fail}
 import org.junit.jupiter.api.Test
 
 import java.util
@@ -326,7 +326,7 @@ class ZkConfigMigrationClientTest extends ZkMigrationTestHarness {
     val image = delta.apply(MetadataProvenance.EMPTY)
 
     // load snapshot to Zookeeper.
-    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient)
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, fail(_))
     kraftMigrationZkWriter.handleSnapshot(image, (_, _, operation) => {
       migrationState = operation.apply(migrationState)
     })
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationClientTest.scala
index 9b71bd6b0e96a..345e2b05d8678 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationClientTest.scala
@@ -317,7 +317,7 @@ class ZkMigrationClientTest extends ZkMigrationTestHarness {
 
   @Test
   def testTopicAndBrokerConfigsMigrationWithSnapshots(): Unit = {
-    val kraftWriter = new KRaftMigrationZkWriter(migrationClient)
+    val kraftWriter = new KRaftMigrationZkWriter(migrationClient, fail(_))
 
     // Add add some topics and broker configs and create new image.
     val topicName = "testTopic"
diff --git a/metadata/src/main/java/org/apache/kafka/image/AclsDelta.java b/metadata/src/main/java/org/apache/kafka/image/AclsDelta.java
index 15e9a69c1939b..3a38d52aca3ab 100644
--- a/metadata/src/main/java/org/apache/kafka/image/AclsDelta.java
+++ b/metadata/src/main/java/org/apache/kafka/image/AclsDelta.java
@@ -25,12 +25,10 @@
 import org.apache.kafka.server.common.MetadataVersion;
 
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 
@@ -40,7 +38,6 @@
 public final class AclsDelta {
     private final AclsImage image;
     private final Map<Uuid, Optional<StandardAcl>> changes = new LinkedHashMap<>();
-    private final Set<StandardAcl> deleted = new HashSet<>();
 
     public AclsDelta(AclsImage image) {
         this.image = image;
@@ -56,15 +53,6 @@ public Map<Uuid, Optional<StandardAcl>> changes() {
         return changes;
     }
 
-    /**
-     * Return a Set of the ACLs which were deleted in this delta. This is used by the ZK migration components.
-     *
-     * @return Set of deleted ACLs
-     */
-    public Set<StandardAcl> deleted() {
-        return deleted;
-    }
-
     void finishSnapshot() {
         for (Entry<Uuid, StandardAcl> entry : image.acls().entrySet()) {
             if (!changes.containsKey(entry.getKey())) {
@@ -93,7 +81,6 @@ public void replay(AccessControlEntryRecord record) {
     public void replay(RemoveAccessControlEntryRecord record) {
         if (image.acls().containsKey(record.id())) {
             changes.put(record.id(), Optional.empty());
-            deleted.add(image.acls().get(record.id()));
         } else if (changes.containsKey(record.id())) {
             changes.remove(record.id());
             // No need to track a ACL that was added and deleted within the same delta
diff --git a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationDriver.java b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationDriver.java
index f33f0790e0a7b..4c796f9eade73 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationDriver.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationDriver.java
@@ -134,7 +134,8 @@ public long nextPollTimeMs() {
         this.time = time;
         LogContext logContext = new LogContext("[KRaftMigrationDriver id=" + nodeId + "] ");
         this.controllerMetrics = controllerMetrics;
-        this.log = logContext.logger(KRaftMigrationDriver.class);
+        Logger log = logContext.logger(KRaftMigrationDriver.class);
+        this.log = log;
         this.migrationState = MigrationDriverState.UNINITIALIZED;
         this.migrationLeadershipState = ZkMigrationLeadershipState.EMPTY;
         this.eventQueue = new KafkaEventQueue(Time.SYSTEM, logContext, "controller-" + nodeId + "-migration-driver-");
@@ -144,7 +145,7 @@ public long nextPollTimeMs() {
         this.initialZkLoadHandler = initialZkLoadHandler;
         this.faultHandler = faultHandler;
         this.quorumFeatures = quorumFeatures;
-        this.zkMetadataWriter = new KRaftMigrationZkWriter(zkMigrationClient);
+        this.zkMetadataWriter = new KRaftMigrationZkWriter(zkMigrationClient, log::error);
         this.recordRedactor = new RecordRedactor(configSchema);
         this.minBatchSize = minBatchSize;
     }
diff --git a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
index 6c82c9cb9ccf9..8a7e148d579e0 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
@@ -60,8 +60,8 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 
 public class KRaftMigrationZkWriter {
 
@@ -82,11 +82,14 @@ public class KRaftMigrationZkWriter {
 
 
     private final MigrationClient migrationClient;
+    private final Consumer<String> errorLogger;
 
     public KRaftMigrationZkWriter(
-        MigrationClient migrationClient
+        MigrationClient migrationClient,
+        Consumer<String> errorLogger
     ) {
         this.migrationClient = migrationClient;
+        this.errorLogger = errorLogger;
     }
 
     public void handleSnapshot(MetadataImage image, KRaftMigrationOperationConsumer operationConsumer) {
@@ -122,7 +125,7 @@ public boolean handleDelta(
             updated = true;
         }
         if (delta.aclsDelta() != null) {
-            handleAclsDelta(image.acls(), delta.aclsDelta(), operationConsumer);
+            handleAclsDelta(previousImage.acls(), image.acls(), delta.aclsDelta(), operationConsumer);
             updated = true;
         }
         if (delta.delegationTokenDelta() != null) {
@@ -593,6 +596,7 @@ void handleAclsSnapshot(AclsImage image, KRaftMigrationOperationConsumer operati
         });
 
         newResources.forEach(resourcePattern -> {
+            // newResources is generated from allAclsInSnapshot, and we don't remove from that map, so this unguarded .get() is safe
             Set<AccessControlEntry> accessControlEntries = allAclsInSnapshot.get(resourcePattern);
             String name = "Writing " + accessControlEntries.size() + " for resource " + resourcePattern;
             operationConsumer.accept(UPDATE_ACL, name, migrationState ->
@@ -612,43 +616,45 @@ void handleAclsSnapshot(AclsImage image, KRaftMigrationOperationConsumer operati
         });
     }
 
-    void handleAclsDelta(AclsImage image, AclsDelta delta, KRaftMigrationOperationConsumer operationConsumer) {
-        // Compute the resource patterns that were changed
-        Set<ResourcePattern> resourcesWithChangedAcls = delta.changes().values()
-            .stream()
-            .filter(Optional::isPresent)
-            .map(Optional::get)
-            .map(this::resourcePatternFromAcl)
-            .collect(Collectors.toSet());
-
-        Set<ResourcePattern> resourcesWithDeletedAcls = delta.deleted()
-            .stream()
-            .map(this::resourcePatternFromAcl)
-            .collect(Collectors.toSet());
-
+    void handleAclsDelta(AclsImage prevImage, AclsImage image, AclsDelta delta, KRaftMigrationOperationConsumer operationConsumer) {
         // Need to collect all ACLs for any changed resource pattern
         Map<ResourcePattern, List<AccessControlEntry>> aclsToWrite = new HashMap<>();
-        image.acls().forEach((uuid, standardAcl) -> {
-            ResourcePattern resourcePattern = resourcePatternFromAcl(standardAcl);
-            boolean removed = resourcesWithDeletedAcls.remove(resourcePattern);
-            // If a resource pattern is present in the delta as a changed or deleted acl, need to include it
-            if (resourcesWithChangedAcls.contains(resourcePattern) || removed) {
-                aclsToWrite.computeIfAbsent(resourcePattern, __ -> new ArrayList<>()).add(
-                    new AccessControlEntry(standardAcl.principal(), standardAcl.host(), standardAcl.operation(), standardAcl.permissionType())
-                );
+        delta.changes().forEach((aclId, aclChange) -> {
+            if (aclChange.isPresent()) {
+                ResourcePattern resourcePattern = resourcePatternFromAcl(aclChange.get());
+                aclsToWrite.put(resourcePattern, new ArrayList<>());
+            } else {
+                // We need to look in the previous image to get deleted ACLs resource pattern
+                StandardAcl deletedAcl = prevImage.acls().get(aclId);
+                if (deletedAcl == null) {
+                    errorLogger.accept("Cannot delete ACL " + aclId + " from ZK since it is missing from previous AclImage");
+                } else {
+                    ResourcePattern resourcePattern = resourcePatternFromAcl(deletedAcl);
+                    aclsToWrite.put(resourcePattern, new ArrayList<>());
+                }
             }
         });
 
-        resourcesWithDeletedAcls.forEach(deletedResource -> {
-            String name = "Deleting resource " + deletedResource + " which has no more ACLs";
-            operationConsumer.accept(DELETE_ACL, name, migrationState ->
-                migrationClient.aclClient().deleteResource(deletedResource, migrationState));
+        // Iterate through the new image to collect any ACLs for these changed resources
+        image.acls().forEach((uuid, standardAcl) -> {
+            ResourcePattern resourcePattern = resourcePatternFromAcl(standardAcl);
+            List<AccessControlEntry> entries = aclsToWrite.get(resourcePattern);
+            if (entries != null) {
+                entries.add(new AccessControlEntry(standardAcl.principal(), standardAcl.host(), standardAcl.operation(), standardAcl.permissionType()));
+            }
         });
 
+        // If there are no more ACLs for a resource, delete it. Otherwise, update it with the new set of ACLs
         aclsToWrite.forEach((resourcePattern, accessControlEntries) -> {
-            String name = "Writing " + accessControlEntries.size() + " for resource " + resourcePattern;
-            operationConsumer.accept(UPDATE_ACL, name, migrationState ->
-                migrationClient.aclClient().writeResourceAcls(resourcePattern, accessControlEntries, migrationState));
+            if (accessControlEntries.isEmpty()) {
+                String name = "Deleting resource " + resourcePattern + " which has no more ACLs";
+                operationConsumer.accept(DELETE_ACL, name, migrationState ->
+                    migrationClient.aclClient().deleteResource(resourcePattern, migrationState));
+            } else {
+                String name = "Writing " + accessControlEntries.size() + " for resource " + resourcePattern;
+                operationConsumer.accept(UPDATE_ACL, name, migrationState ->
+                    migrationClient.aclClient().writeResourceAcls(resourcePattern, accessControlEntries, migrationState));
+            }
         });
     }
 
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriterTest.java b/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriterTest.java
index e823d6bdaca60..0e109f6b9c74a 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriterTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriterTest.java
@@ -80,7 +80,7 @@ public void iterateTopics(EnumSet<TopicVisitorInterest> interests, TopicVisitor
             .setConfigMigrationClient(configClient)
             .build();
 
-        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient);
+        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient, __ -> { });
 
         MetadataImage image = new MetadataImage(
             MetadataProvenance.EMPTY,
@@ -120,7 +120,7 @@ public void testReconcileSnapshotEmptyZk() {
             .setAclMigrationClient(aclClient)
             .build();
 
-        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient);
+        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient, __ -> { });
 
         MetadataImage image = new MetadataImage(
             MetadataProvenance.EMPTY,
@@ -179,7 +179,7 @@ public void iterateTopics(EnumSet<TopicVisitorInterest> interests, TopicVisitor
             .setAclMigrationClient(aclClient)
             .build();
 
-        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient);
+        KRaftMigrationZkWriter writer = new KRaftMigrationZkWriter(migrationClient, __ -> { });
 
         MetadataImage image = new MetadataImage(
             MetadataProvenance.EMPTY,

From 87c7fc0cd73be49c7d942ad23caf266d5b71263c Mon Sep 17 00:00:00 2001
From: Kohei Nozaki <kohei@apache.org>
Date: Thu, 8 Feb 2024 06:02:42 +0900
Subject: [PATCH 011/258] KAFKA-16055: Thread unsafe use of HashMap stored in
 QueryableStoreProvider#storeProviders (#15121)

This PR replaces a HashMap by a ConcurrentHashMap so that the local state store queries can be made from multiple threads. See this for additional context: https://lists.apache.org/thread/gpct1275bfqovlckptn3lvf683qpoxol

Reviewers: Anna Sophie Blee-Goldman <ableegoldman@apache.org>, Sagar Rao <sagarmeansocean@gmail.com>
---
 .../kafka/streams/state/internals/QueryableStoreProvider.java | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/QueryableStoreProvider.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/QueryableStoreProvider.java
index 07cf0ee86079a..7ed11917b9cb8 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/QueryableStoreProvider.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/QueryableStoreProvider.java
@@ -20,7 +20,7 @@
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.state.QueryableStoreType;
 
-import java.util.HashMap;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -36,7 +36,7 @@ public class QueryableStoreProvider {
     private final GlobalStateStoreProvider globalStoreProvider;
 
     public QueryableStoreProvider(final GlobalStateStoreProvider globalStateStoreProvider) {
-        this.storeProviders = new HashMap<>();
+        this.storeProviders = new ConcurrentHashMap<>();
         this.globalStoreProvider = globalStateStoreProvider;
     }
 

From ab44d5bb63267caad26733801c1cef3622f3a4d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hy=20=28=ED=95=98=EC=9D=B4=29?= <yello1109@daum.net>
Date: Thu, 8 Feb 2024 13:56:14 +0530
Subject: [PATCH 012/258] MINOR: Remove unused logger in
 PlaintextChannelBuilder (#15301)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
Co-authored-by: highluck <yello1109>
---
 .../apache/kafka/common/network/PlaintextChannelBuilder.java   | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java b/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
index 0b4738afa2c07..4ff41f81b594c 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
@@ -23,8 +23,6 @@
 import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.PlaintextAuthenticationContext;
 import org.apache.kafka.common.utils.Utils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -35,7 +33,6 @@
 import java.util.function.Supplier;
 
 public class PlaintextChannelBuilder implements ChannelBuilder {
-    private static final Logger log = LoggerFactory.getLogger(PlaintextChannelBuilder.class);
     private final ListenerName listenerName;
     private Map<String, ?> configs;
 

From cd328d84acc956deb61964d238ca803249aa6314 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Thu, 8 Feb 2024 19:09:54 +0800
Subject: [PATCH 013/258] KAFKA-15201: Allow git push to fail gracefully
 (#14645)

Reviewers: Divij Vaidya <diviv@amazon.com>, Qichao Chu <qichao@uber.com>
---
 release.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/release.py b/release.py
index 43c5809861554..6b5b538656173 100755
--- a/release.py
+++ b/release.py
@@ -730,7 +730,31 @@ def select_gpg_key():
     fail("Ok, giving up")
 if not user_ok("Ok to push RC tag %s (y/n)?: " % rc_tag):
     fail("Ok, giving up")
-cmd("Pushing RC tag", "git push %s %s" % (PUSH_REMOTE_NAME, rc_tag))
+
+print(f"Pushing RC tag {rc_tag} to {PUSH_REMOTE_NAME}")
+try:
+    push_command = f"git push {PUSH_REMOTE_NAME} {rc_tag}".split()
+    output = subprocess.check_output(push_command, stderr=subprocess.STDOUT)
+    print_output(output.decode('utf-8'))
+    if "error" in output.decode('utf-8'):
+        print("*********************************************")
+        print("*** ERROR when trying to perform git push ***")
+        print("*********************************************")
+        print(output)
+        print("")
+        print("Due the failure of git push, the program will exit here. Please note that: ")
+        print(f"1) You are still at branch {release_version}, not {starting_branch}")
+        print(f"2) Tag {rc_tag} is still present locally")
+        print("")
+        print(f"In order to restart the workflow, you will have to manually switch back to the original branch and delete the branch {release_version} and tag {rc_tag}")
+        sys.exit(1)
+except Exception as e:
+    print(f"Failed when trying to git push {rc_tag}. Error: {e}"
+    print("You may need to clean up branches/tags yourself before retrying.")
+    print("Due the failure of git push, the program will exit here. Please note that: ")
+    print(f"1) You are still at branch {release_version}, not {starting_branch}")
+    print(f"2) Tag {rc_tag} is still present locally")
+    sys.exit(1)
 
 # Move back to starting branch and clean out the temporary release branch (e.g. 1.0.0) we used to generate everything
 cmd("Resetting repository working state", "git reset --hard HEAD && git checkout %s" % starting_branch, shell=True)

From 116bc000c8c6533552321fe1d395629c2aa00bd9 Mon Sep 17 00:00:00 2001
From: David Arthur <mumrah@gmail.com>
Date: Thu, 8 Feb 2024 18:44:42 -0500
Subject: [PATCH 014/258] MINOR: fix scala compile issue (#15343)

Reviewers: David Jacot <djacot@confluent.io>
---
 .../unit/kafka/zk/migration/ZkAclMigrationClientTest.scala    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
index bb4d3e646c7b7..a4045f5794871 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
@@ -231,7 +231,7 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
 
     // Sync image to ZK
     val errorLogs = mutable.Buffer[String]()
-    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, errorLogs.append)
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, msg => errorLogs.append(msg))
     kraftMigrationZkWriter.handleSnapshot(image, (_, _, operation) => {
       migrationState = operation.apply(migrationState)
     })
@@ -276,7 +276,7 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
   def testAclUpdateAndDelete(): Unit = {
     zkClient.createAclPaths()
     val errorLogs = mutable.Buffer[String]()
-    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, errorLogs.append)
+    val kraftMigrationZkWriter = new KRaftMigrationZkWriter(migrationClient, msg => errorLogs.append(msg))
 
     val topicName = "topic-" + Uuid.randomUuid()
     val otherName = "other-" + Uuid.randomUuid()

From adbffcacb43814b8bbd11015845c844f2d5e186a Mon Sep 17 00:00:00 2001
From: Matthias Berndt <matthias_berndt@gmx.de>
Date: Fri, 9 Feb 2024 11:57:32 +0100
Subject: [PATCH 015/258] MINOR: Remove collections-compat dependency for
 streams module when on Scala 2.13 (#15239)

Reviewers: Divij Vaidya <diviv@amazon.com>, Josep Prat <jlprat@apache.org>

---------

Co-authored-by: Matthias Berndt <matthias.berndt@ttmzero.com>
---
 build.gradle | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/build.gradle b/build.gradle
index c43b9c1e0b42a..d365233465108 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2314,8 +2314,14 @@ project(':streams:streams-scala') {
     api project(':streams')
 
     api libs.scalaLibrary
-    api libs.scalaCollectionCompat
-
+    if ( versions.baseScala == '2.12' ) {
+      // Scala-Collection-Compat isn't required when compiling with Scala 2.13 or later,
+      // and having it in the dependencies could lead to classpath conflicts in Scala 3
+      // projects that use kafka-streams-kafka_2.13 (because we don't have a Scala 3 version yet)
+      // but also pull in scala-collection-compat_3 via another dependency.
+      // So we make sure to not include it in the dependencies.
+      api libs.scalaCollectionCompat
+    }
     testImplementation project(':core')
     testImplementation project(':core').sourceSets.test.output
     testImplementation project(':server-common').sourceSets.test.output

From 489a7dd71eebbf4f5effedeea1f3341854ddc75e Mon Sep 17 00:00:00 2001
From: "Gyeongwon, Do" <dct012@gmail.com>
Date: Fri, 9 Feb 2024 20:07:20 +0900
Subject: [PATCH 016/258] MINOR: Improve Code Style (#15319)

- Removing ! and Unused Imports
- Put a space after the control structure's defining keyword.
- remove unnecessary whitespace a space after the method name in higher-order function invocations.

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 .../scala/kafka/admin/ConfigCommand.scala     |  4 +--
 .../controller/ControllerEventManager.scala   |  2 +-
 core/src/main/scala/kafka/log/LocalLog.scala  |  2 +-
 .../src/main/scala/kafka/log/LogCleaner.scala | 12 +++----
 .../scala/kafka/log/LogCleanerManager.scala   |  2 +-
 .../src/main/scala/kafka/log/LogManager.scala |  4 +--
 .../main/scala/kafka/serializer/Decoder.scala |  2 +-
 .../kafka/server/ConfigAdminManager.scala     |  8 ++---
 .../main/scala/kafka/server/KafkaServer.scala |  2 +-
 .../scala/kafka/server/ZkAdminManager.scala   |  6 ++--
 .../server/metadata/ZkMetadataCache.scala     |  4 +--
 .../main/scala/kafka/tools/MirrorMaker.scala  |  2 +-
 .../kafka/tools/StateChangeLogMerger.scala    |  2 +-
 .../src/main/scala/kafka/utils/FileLock.scala |  2 +-
 .../main/scala/kafka/utils/Throttler.scala    |  4 +--
 .../main/scala/kafka/utils/ToolsUtils.scala   |  4 +--
 .../kafka/utils/VerifiableProperties.scala    | 18 +++++------
 ...thLegacyMessageFormatIntegrationTest.scala |  4 +--
 .../api/PlaintextAdminIntegrationTest.scala   |  2 +-
 .../ProducerIdsIntegrationTest.scala          |  2 +-
 .../kafka/server/KRaftClusterTest.scala       |  2 +-
 .../kafka/raft/KafkaMetadataLogTest.scala     |  8 ++---
 .../scala/other/kafka/StressTestLog.scala     |  4 +--
 .../other/kafka/TestLinearWriteSpeed.scala    | 16 +++++-----
 .../unit/kafka/cluster/PartitionTest.scala    |  2 +-
 .../ControllerIntegrationTest.scala           |  2 +-
 .../transaction/ProducerIdManagerTest.scala   |  2 +-
 .../integration/KafkaServerTestHarness.scala  |  2 +-
 .../AbstractLogCleanerIntegrationTest.scala   |  2 +-
 .../kafka/log/LogCleanerManagerTest.scala     |  8 ++---
 .../scala/unit/kafka/log/LogCleanerTest.scala | 32 +++++++++----------
 .../scala/unit/kafka/log/LogLoaderTest.scala  | 18 +++++------
 .../scala/unit/kafka/log/LogManagerTest.scala |  4 +--
 .../scala/unit/kafka/log/LogSegmentTest.scala |  8 ++---
 .../unit/kafka/log/OffsetIndexTest.scala      | 14 ++++----
 .../scala/unit/kafka/log/OffsetMapTest.scala  | 12 +++----
 .../scala/unit/kafka/log/TimeIndexTest.scala  |  2 +-
 .../scala/unit/kafka/log/UnifiedLogTest.scala | 20 ++++++------
 .../security/authorizer/AuthorizerTest.scala  |  4 +--
 .../AbstractApiVersionsRequestTest.scala      |  2 +-
 .../server/AbstractFetcherThreadTest.scala    |  2 +-
 .../server/DeleteTopicsRequestTest.scala      |  2 +-
 .../ThrottledChannelExpirationTest.scala      |  4 +--
 ...opicIdWithOldInterBrokerProtocolTest.scala |  2 +-
 .../kafka/tools/DumpLogSegmentsTest.scala     |  2 +-
 .../unit/kafka/utils/CoreUtilsTest.scala      |  2 +-
 .../unit/kafka/utils/SchedulerTest.scala      |  2 +-
 .../scala/unit/kafka/utils/TestUtils.scala    |  8 ++---
 .../unit/kafka/zk/AdminZkClientTest.scala     |  2 +-
 .../kafka/server/AssignmentsManager.java      |  2 --
 50 files changed, 139 insertions(+), 141 deletions(-)

diff --git a/core/src/main/scala/kafka/admin/ConfigCommand.scala b/core/src/main/scala/kafka/admin/ConfigCommand.scala
index 0d82d23d7ed7f..bcbeedc468e4e 100644
--- a/core/src/main/scala/kafka/admin/ConfigCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConfigCommand.scala
@@ -979,10 +979,10 @@ object ConfigCommand extends Logging {
         val isAddConfigFilePresent = options.has(addConfigFile)
         val isDeleteConfigPresent = options.has(deleteConfig)
 
-        if(isAddConfigPresent && isAddConfigFilePresent)
+        if (isAddConfigPresent && isAddConfigFilePresent)
           throw new IllegalArgumentException("Only one of --add-config or --add-config-file must be specified")
 
-        if(!isAddConfigPresent && !isAddConfigFilePresent && !isDeleteConfigPresent)
+        if (!isAddConfigPresent && !isAddConfigFilePresent && !isDeleteConfigPresent)
           throw new IllegalArgumentException("At least one of --add-config, --add-config-file, or --delete-config must be specified with --alter")
       }
     }
diff --git a/core/src/main/scala/kafka/controller/ControllerEventManager.scala b/core/src/main/scala/kafka/controller/ControllerEventManager.scala
index 9a5a6fe1abcf2..431f06daa1dba 100644
--- a/core/src/main/scala/kafka/controller/ControllerEventManager.scala
+++ b/core/src/main/scala/kafka/controller/ControllerEventManager.scala
@@ -109,7 +109,7 @@ class ControllerEventManager(controllerId: Int,
     queuedEvent
   }
 
-  def clearAndPut(event: ControllerEvent): QueuedEvent = inLock(putLock){
+  def clearAndPut(event: ControllerEvent): QueuedEvent = inLock(putLock) {
     val preemptedEvents = new ArrayList[QueuedEvent]()
     queue.drainTo(preemptedEvents)
     preemptedEvents.forEach(_.preempt(processor))
diff --git a/core/src/main/scala/kafka/log/LocalLog.scala b/core/src/main/scala/kafka/log/LocalLog.scala
index c649ad536bda2..b2121f5312d7b 100644
--- a/core/src/main/scala/kafka/log/LocalLog.scala
+++ b/core/src/main/scala/kafka/log/LocalLog.scala
@@ -485,7 +485,7 @@ class LocalLog(@volatile private var _dir: File,
             s" =max(provided offset = $expectedNextOffset, LEO = $logEndOffset) while it already exists. Existing " +
             s"segment is ${segments.get(newOffset)}.")
         }
-      } else if (!segments.isEmpty && newOffset < activeSegment.baseOffset) {
+      } else if (segments.nonEmpty && newOffset < activeSegment.baseOffset) {
         throw new KafkaException(
           s"Trying to roll a new log segment for topic partition $topicPartition with " +
             s"start offset $newOffset =max(provided offset = $expectedNextOffset, LEO = $logEndOffset) lower than start offset of the active segment $activeSegment")
diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala
index 0422eee767bb8..8098ea237e06d 100644
--- a/core/src/main/scala/kafka/log/LogCleaner.scala
+++ b/core/src/main/scala/kafka/log/LogCleaner.scala
@@ -931,7 +931,7 @@ private[log] class Cleaner(val id: Int,
    */
   private def growBuffers(maxLogMessageSize: Int): Unit = {
     val maxBufferSize = math.max(maxLogMessageSize, maxIoBufferSize)
-    if(readBuffer.capacity >= maxBufferSize || writeBuffer.capacity >= maxBufferSize)
+    if (readBuffer.capacity >= maxBufferSize || writeBuffer.capacity >= maxBufferSize)
       throw new IllegalStateException("This log contains a message larger than maximum allowable size of %s.".format(maxBufferSize))
     val newSize = math.min(this.readBuffer.capacity * 2, maxBufferSize)
     info(s"Growing cleaner I/O buffers from ${readBuffer.capacity} bytes to $newSize bytes.")
@@ -943,9 +943,9 @@ private[log] class Cleaner(val id: Int,
    * Restore the I/O buffer capacity to its original size
    */
   private def restoreBuffers(): Unit = {
-    if(this.readBuffer.capacity > this.ioBufferSize)
+    if (this.readBuffer.capacity > this.ioBufferSize)
       this.readBuffer = ByteBuffer.allocate(this.ioBufferSize)
-    if(this.writeBuffer.capacity > this.ioBufferSize)
+    if (this.writeBuffer.capacity > this.ioBufferSize)
       this.writeBuffer = ByteBuffer.allocate(this.ioBufferSize)
   }
 
@@ -964,13 +964,13 @@ private[log] class Cleaner(val id: Int,
   private[log] def groupSegmentsBySize(segments: Iterable[LogSegment], maxSize: Int, maxIndexSize: Int, firstUncleanableOffset: Long): List[Seq[LogSegment]] = {
     var grouped = List[List[LogSegment]]()
     var segs = segments.toList
-    while(segs.nonEmpty) {
+    while (segs.nonEmpty) {
       var group = List(segs.head)
       var logSize = segs.head.size.toLong
       var indexSize = segs.head.offsetIndex.sizeInBytes.toLong
       var timeIndexSize = segs.head.timeIndex.sizeInBytes.toLong
       segs = segs.tail
-      while(segs.nonEmpty &&
+      while (segs.nonEmpty &&
             logSize + segs.head.size <= maxSize &&
             indexSize + segs.head.offsetIndex.sizeInBytes <= maxIndexSize &&
             timeIndexSize + segs.head.timeIndex.sizeInBytes <= maxIndexSize &&
@@ -1124,7 +1124,7 @@ private[log] class Cleaner(val id: Int,
       stats.indexBytesRead(bytesRead)
 
       // if we didn't read even one complete message, our read buffer may be too small
-      if(position == startPosition)
+      if (position == startPosition)
         growBuffersOrFail(segment.log, position, maxLogMessageSize, records)
     }
 
diff --git a/core/src/main/scala/kafka/log/LogCleanerManager.scala b/core/src/main/scala/kafka/log/LogCleanerManager.scala
index fa67f9c321ff6..f69ec42c716c6 100755
--- a/core/src/main/scala/kafka/log/LogCleanerManager.scala
+++ b/core/src/main/scala/kafka/log/LogCleanerManager.scala
@@ -298,7 +298,7 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
         case Some(s) =>
           throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be aborted and paused since it is in $s state.")
       }
-      while(!isCleaningInStatePaused(topicPartition))
+      while (!isCleaningInStatePaused(topicPartition))
         pausedCleaningCond.await(100, TimeUnit.MILLISECONDS)
     }
   }
diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala
index 41ebdd1a4c1df..db96c497c5537 100755
--- a/core/src/main/scala/kafka/log/LogManager.scala
+++ b/core/src/main/scala/kafka/log/LogManager.scala
@@ -1323,7 +1323,7 @@ class LogManager(logDirs: Seq[File],
    * data directories by fewest partitions.
    */
   private def nextLogDirs(): List[File] = {
-    if(_liveLogDirs.size == 1) {
+    if (_liveLogDirs.size == 1) {
       List(_liveLogDirs.peek())
     } else {
       // count the number of logs in each parent directory (including 0 for empty directories
@@ -1438,7 +1438,7 @@ class LogManager(logDirs: Seq[File],
         val timeSinceLastFlush = time.milliseconds - log.lastFlushTime
         debug(s"Checking if flush is needed on ${topicPartition.topic} flush interval ${log.config.flushMs}" +
               s" last flushed ${log.lastFlushTime} time since last flush: $timeSinceLastFlush")
-        if(timeSinceLastFlush >= log.config.flushMs)
+        if (timeSinceLastFlush >= log.config.flushMs)
           log.flush(false)
       } catch {
         case e: Throwable =>
diff --git a/core/src/main/scala/kafka/serializer/Decoder.scala b/core/src/main/scala/kafka/serializer/Decoder.scala
index 4b8c545f52beb..1ad554100812a 100644
--- a/core/src/main/scala/kafka/serializer/Decoder.scala
+++ b/core/src/main/scala/kafka/serializer/Decoder.scala
@@ -44,7 +44,7 @@ class DefaultDecoder(props: VerifiableProperties = null) extends Decoder[Array[B
  */
 class StringDecoder(props: VerifiableProperties = null) extends Decoder[String] {
   val encoding =
-    if(props == null)
+    if (props == null)
       StandardCharsets.UTF_8.name()
     else
       props.getString("serializer.encoding", StandardCharsets.UTF_8.name())
diff --git a/core/src/main/scala/kafka/server/ConfigAdminManager.scala b/core/src/main/scala/kafka/server/ConfigAdminManager.scala
index 9d56197022d87..4e1ee24b59292 100644
--- a/core/src/main/scala/kafka/server/ConfigAdminManager.scala
+++ b/core/src/main/scala/kafka/server/ConfigAdminManager.scala
@@ -144,7 +144,7 @@ class ConfigAdminManager(nodeId: Int,
             case BROKER =>
               // The resource name must be either blank (if setting a cluster config) or
               // the ID of this specific broker.
-              if (!configResource.name().isEmpty) {
+              if (configResource.name().nonEmpty) {
                 validateResourceNameIsCurrentNodeId(resource.resourceName())
               }
               validateBrokerConfigChange(resource, configResource)
@@ -168,7 +168,7 @@ class ConfigAdminManager(nodeId: Int,
     resource: IAlterConfigsResource,
     configResource: ConfigResource
   ): Unit = {
-    val perBrokerConfig = !configResource.name().isEmpty
+    val perBrokerConfig = configResource.name().nonEmpty
     val persistentProps = configRepository.config(configResource)
     val configProps = conf.dynamicConfig.fromPersistentProps(persistentProps, perBrokerConfig)
     val alterConfigOps = resource.configs().asScala.map {
@@ -193,7 +193,7 @@ class ConfigAdminManager(nodeId: Int,
     configResource: ConfigResource
   ): Unit = {
     try {
-      conf.dynamicConfig.validate(props, !configResource.name().isEmpty)
+      conf.dynamicConfig.validate(props, configResource.name().nonEmpty)
     } catch {
       case e: ApiException => throw e
       //KAFKA-13609: InvalidRequestException is not really the right exception here if the
@@ -243,7 +243,7 @@ class ConfigAdminManager(nodeId: Int,
           }
           resourceType match {
             case BROKER =>
-              if (!configResource.name().isEmpty) {
+              if (configResource.name().nonEmpty) {
                 validateResourceNameIsCurrentNodeId(resource.resourceName())
               }
               validateBrokerConfigChange(resource, configResource)
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index 6f159111c4d1b..bea0a252fb77d 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -661,7 +661,7 @@ class KafkaServer(
 
   protected def createRemoteLogManager(): Option[RemoteLogManager] = {
     if (config.remoteLogManagerConfig.enableRemoteStorageSystem()) {
-      if(config.logDirs.size > 1) {
+      if (config.logDirs.size > 1) {
         throw new KafkaException("Tiered storage is not supported with multiple log dirs.")
       }
 
diff --git a/core/src/main/scala/kafka/server/ZkAdminManager.scala b/core/src/main/scala/kafka/server/ZkAdminManager.scala
index 694546cf81228..dce32ae3fcc13 100644
--- a/core/src/main/scala/kafka/server/ZkAdminManager.scala
+++ b/core/src/main/scala/kafka/server/ZkAdminManager.scala
@@ -661,7 +661,7 @@ class ZkAdminManager(val config: KafkaConfig,
     val exactUser = wantExact(userComponent)
     val exactClientId = wantExact(clientIdComponent)
 
-    def wantExcluded(component: Option[ClientQuotaFilterComponent]): Boolean = strict && !component.isDefined
+    def wantExcluded(component: Option[ClientQuotaFilterComponent]): Boolean = strict && component.isEmpty
     val excludeUser = wantExcluded(userComponent)
     val excludeClientId = wantExcluded(clientIdComponent)
 
@@ -803,7 +803,7 @@ class ZkAdminManager(val config: KafkaConfig,
   private val attemptToDescribeUserThatDoesNotExist = "Attempt to describe a user credential that does not exist"
 
   def describeUserScramCredentials(users: Option[Seq[String]]): DescribeUserScramCredentialsResponseData = {
-    val describingAllUsers = !users.isDefined || users.get.isEmpty
+    val describingAllUsers = users.isEmpty || users.get.isEmpty
     val retval = new DescribeUserScramCredentialsResponseData()
     val userResults = mutable.Map[String, DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult]()
 
@@ -970,7 +970,7 @@ class ZkAdminManager(val config: KafkaConfig,
     val initiallyValidUserMechanismPairs = (upsertions.filter(upsertion => !invalidUsers.contains(upsertion.name)).map(upsertion => (upsertion.name, upsertion.mechanism)) ++
       deletions.filter(deletion => !invalidUsers.contains(deletion.name)).map(deletion => (deletion.name, deletion.mechanism)))
 
-    val usersWithDuplicateUserMechanismPairs = initiallyValidUserMechanismPairs.groupBy(identity).filter (
+    val usersWithDuplicateUserMechanismPairs = initiallyValidUserMechanismPairs.groupBy(identity).filter(
       userMechanismPairAndOccurrencesTuple => userMechanismPairAndOccurrencesTuple._2.length > 1).keys.map(userMechanismPair => userMechanismPair._1).toSet
     usersWithDuplicateUserMechanismPairs.foreach { user =>
       retval.results.add(new AlterUserScramCredentialsResult()
diff --git a/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala b/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala
index c980d9320d5e6..2a6a880d54ac0 100755
--- a/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala
+++ b/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala
@@ -695,12 +695,12 @@ class ZkMetadataCache(
    *                           minExpectedEpoch within timeoutMs.
    */
   def waitUntilFeatureEpochOrThrow(minExpectedEpoch: Long, timeoutMs: Long): Unit = {
-    if(minExpectedEpoch < 0L) {
+    if (minExpectedEpoch < 0L) {
       throw new IllegalArgumentException(
         s"Expected minExpectedEpoch >= 0, but $minExpectedEpoch was provided.")
     }
 
-    if(timeoutMs < 0L) {
+    if (timeoutMs < 0L) {
       throw new IllegalArgumentException(s"Expected timeoutMs >= 0, but $timeoutMs was provided.")
     }
     val waitEndTimeNanos = System.nanoTime() + (timeoutMs * 1000000)
diff --git a/core/src/main/scala/kafka/tools/MirrorMaker.scala b/core/src/main/scala/kafka/tools/MirrorMaker.scala
index d7cf6dbb275ce..d7332c58270e1 100755
--- a/core/src/main/scala/kafka/tools/MirrorMaker.scala
+++ b/core/src/main/scala/kafka/tools/MirrorMaker.scala
@@ -239,7 +239,7 @@ object MirrorMaker extends Logging {
           exitingOnSendFailure = true
           fatal("Mirror maker thread failure due to ", t)
       } finally {
-        CoreUtils.swallow ({
+        CoreUtils.swallow({
           info("Flushing producer.")
           producer.flush()
 
diff --git a/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala b/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
index f2c796bbfdd33..4ac918a63c157 100755
--- a/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
+++ b/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
@@ -88,7 +88,7 @@ object StateChangeLogMerger extends Logging {
                               .ofType(classOf[String])
                               .defaultsTo("9999-12-31 23:59:59,999")
                               
-    if(args.isEmpty)
+    if (args.isEmpty)
       CommandLineUtils.printUsageAndExit(parser, "A tool for merging the log files from several brokers to reconnstruct a unified history of what happened.")
 
 
diff --git a/core/src/main/scala/kafka/utils/FileLock.scala b/core/src/main/scala/kafka/utils/FileLock.scala
index a7ae1aa735b0b..7a5afc9fb958b 100644
--- a/core/src/main/scala/kafka/utils/FileLock.scala
+++ b/core/src/main/scala/kafka/utils/FileLock.scala
@@ -65,7 +65,7 @@ class FileLock(val file: File) extends Logging {
   def unlock(): Unit = {
     this synchronized {
       trace(s"Releasing lock on ${file.getAbsolutePath}")
-      if(flock != null)
+      if (flock != null)
         flock.release()
     }
   }
diff --git a/core/src/main/scala/kafka/utils/Throttler.scala b/core/src/main/scala/kafka/utils/Throttler.scala
index 1e77624a9fbe5..286343cd4449f 100644
--- a/core/src/main/scala/kafka/utils/Throttler.scala
+++ b/core/src/main/scala/kafka/utils/Throttler.scala
@@ -95,13 +95,13 @@ object Throttler {
     val interval = 30000
     var start = System.currentTimeMillis
     var total = 0
-    while(true) {
+    while (true) {
       val value = rand.nextInt(1000)
       Thread.sleep(1)
       throttler.maybeThrottle(value)
       total += value
       val now = System.currentTimeMillis
-      if(now - start >= interval) {
+      if (now - start >= interval) {
         println(total / (interval/1000.0))
         start = now
         total = 0
diff --git a/core/src/main/scala/kafka/utils/ToolsUtils.scala b/core/src/main/scala/kafka/utils/ToolsUtils.scala
index 8f3ae49b7aa51..7a2aa3111c557 100644
--- a/core/src/main/scala/kafka/utils/ToolsUtils.scala
+++ b/core/src/main/scala/kafka/utils/ToolsUtils.scala
@@ -25,14 +25,14 @@ import scala.collection.mutable
 object ToolsUtils {
 
   def validatePortOrDie(parser: OptionParser, hostPort: String): Unit = {
-    val hostPorts: Array[String] = if(hostPort.contains(','))
+    val hostPorts: Array[String] = if (hostPort.contains(','))
       hostPort.split(",")
     else
       Array(hostPort)
     val validHostPort = hostPorts.filter { hostPortData =>
       org.apache.kafka.common.utils.Utils.getPort(hostPortData) != null
     }
-    val isValid = !validHostPort.isEmpty && validHostPort.length == hostPorts.length
+    val isValid = validHostPort.nonEmpty && validHostPort.length == hostPorts.length
     if (!isValid)
       CommandLineUtils.printUsageAndExit(parser, "Please provide valid host:port like host1:9091,host2:9092\n ")
   }
diff --git a/core/src/main/scala/kafka/utils/VerifiableProperties.scala b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
index 33ba418fc8c65..15f7978d2d847 100755
--- a/core/src/main/scala/kafka/utils/VerifiableProperties.scala
+++ b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
@@ -44,7 +44,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
   def getProperty(name: String): String = {
     val value = props.getProperty(name)
     referenceSet.add(name)
-    if(value == null) value else value.trim()
+    if (value == null) value else value.trim()
   }
 
   /**
@@ -75,7 +75,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    */
   private def getIntInRange(name: String, default: Int, range: (Int, Int)): Int = {
     val v =
-      if(containsKey(name))
+      if (containsKey(name))
         getProperty(name).toInt
       else
         default
@@ -85,7 +85,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
 
  private def getShortInRange(name: String, default: Short, range: (Short, Short)): Short = {
     val v =
-      if(containsKey(name))
+      if (containsKey(name))
         getProperty(name).toShort
       else
         default
@@ -118,7 +118,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    */
   private def getLongInRange(name: String, default: Long, range: (Long, Long)): Long = {
     val v =
-      if(containsKey(name))
+      if (containsKey(name))
         getProperty(name).toLong
       else
         default
@@ -140,7 +140,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    * @param default The default value for the property if not present
    */
   def getDouble(name: String, default: Double): Double = {
-    if(containsKey(name))
+    if (containsKey(name))
       getDouble(name)
     else
       default
@@ -153,7 +153,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    * @return the boolean value
    */
   def getBoolean(name: String, default: Boolean): Boolean = {
-    if(!containsKey(name))
+    if (!containsKey(name))
       default
     else {
       val v = getProperty(name)
@@ -168,7 +168,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    * Get a string property, or, if no such property is defined, return the given default value
    */
   def getString(name: String, default: String): String = {
-    if(containsKey(name))
+    if (containsKey(name))
       getProperty(name)
     else
       default
@@ -190,7 +190,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
       val m = Csv.parseCsvMap(getString(name, "")).asScala
       m.foreach {
         case(key, value) => 
-          if(!valid(value))
+          if (!valid(value))
             throw new IllegalArgumentException("Invalid entry '%s' = '%s' for property '%s'".format(key, value, name))
       }
       m
@@ -202,7 +202,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
   def verify(): Unit = {
     info("Verifying properties")
     val propNames = Collections.list(props.propertyNames).asScala.map(_.toString).sorted
-    for(key <- propNames) {
+    for (key <- propNames) {
       if (!referenceSet.contains(key) && !key.startsWith("external"))
         warn("Property %s is not valid".format(key))
       else
diff --git a/core/src/test/scala/integration/kafka/api/ConsumerWithLegacyMessageFormatIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/ConsumerWithLegacyMessageFormatIntegrationTest.scala
index 634d2b52c1cfe..c338865c7d73d 100644
--- a/core/src/test/scala/integration/kafka/api/ConsumerWithLegacyMessageFormatIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ConsumerWithLegacyMessageFormatIntegrationTest.scala
@@ -35,7 +35,7 @@ class ConsumerWithLegacyMessageFormatIntegrationTest extends AbstractConsumerTes
   override protected def brokerPropertyOverrides(properties: Properties): Unit = {
     // legacy message formats are only supported with IBP < 3.0
     // KRaft mode is not supported for inter.broker.protocol.version = 2.8, The minimum version required is 3.0-IV1"
-    if(!isKRaftTest())
+    if (!isKRaftTest())
       properties.put(KafkaConfig.InterBrokerProtocolVersionProp, "2.8")
   }
 
@@ -91,7 +91,7 @@ class ConsumerWithLegacyMessageFormatIntegrationTest extends AbstractConsumerTes
     assertEquals(20, timestampTopic1P1.timestamp)
     assertEquals(Optional.of(0), timestampTopic1P1.leaderEpoch)
 
-    if(!isKRaftTest()) {
+    if (!isKRaftTest()) {
       assertNull(timestampOffsets.get(new TopicPartition(topic2, 0)), "null should be returned when message format is 0.9.0")
       assertNull(timestampOffsets.get(new TopicPartition(topic2, 1)), "null should be returned when message format is 0.9.0")
     }
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
index c291859d559ab..54be3337625d4 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -105,7 +105,7 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
     var nodeStrs: List[String] = null
     do {
       val nodes = client.describeCluster().nodes().get().asScala
-      nodeStrs = nodes.map ( node => s"${node.host}:${node.port}" ).toList.sorted
+      nodeStrs = nodes.map(node => s"${node.host}:${node.port}").toList.sorted
     } while (nodeStrs.size < brokerStrs.size)
     assertEquals(brokerStrs.mkString(","), nodeStrs.mkString(","))
   }
diff --git a/core/src/test/scala/integration/kafka/coordinator/transaction/ProducerIdsIntegrationTest.scala b/core/src/test/scala/integration/kafka/coordinator/transaction/ProducerIdsIntegrationTest.scala
index 519c2bcf08804..2a225d7c72196 100644
--- a/core/src/test/scala/integration/kafka/coordinator/transaction/ProducerIdsIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/coordinator/transaction/ProducerIdsIntegrationTest.scala
@@ -100,7 +100,7 @@ class ProducerIdsIntegrationTest {
     val deadline = 5.seconds.fromNow
     var shouldRetry = true
     var response: InitProducerIdResponse = null
-    while(shouldRetry && deadline.hasTimeLeft()) {
+    while (shouldRetry && deadline.hasTimeLeft()) {
       val data = new InitProducerIdRequestData()
         .setProducerEpoch(RecordBatch.NO_PRODUCER_EPOCH)
         .setProducerId(RecordBatch.NO_PRODUCER_ID)
diff --git a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
index 048caa859f54e..130d0e5642ea4 100644
--- a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
@@ -310,7 +310,7 @@ class KRaftClusterTest {
         filter = ClientQuotaFilter.contains(
           List(ClientQuotaFilterComponent.ofEntity("user", "testkit")).asJava)
 
-        TestUtils.tryUntilNoAssertionError(){
+        TestUtils.tryUntilNoAssertionError() {
           val results = admin.describeClientQuotas(filter).entities().get()
           assertEquals(2, results.size(), "Broker did not see two client quotas")
           assertEquals(9999.0, results.get(entity).get("producer_byte_rate"), 1e-6)
diff --git a/core/src/test/scala/kafka/raft/KafkaMetadataLogTest.scala b/core/src/test/scala/kafka/raft/KafkaMetadataLogTest.scala
index 51d31b11eaccb..508a82a511c4b 100644
--- a/core/src/test/scala/kafka/raft/KafkaMetadataLogTest.scala
+++ b/core/src/test/scala/kafka/raft/KafkaMetadataLogTest.scala
@@ -849,7 +849,7 @@ final class KafkaMetadataLogTest {
     val log = buildMetadataLog(tempDir, mockTime, config)
 
     // Generate some segments
-    for(_ <- 0 to 100) {
+    for (_ <- 0 to 100) {
       append(log, 47, 1) // An odd number of records to avoid offset alignment
     }
     assertFalse(log.maybeClean(), "Should not clean since HW was still 0")
@@ -889,12 +889,12 @@ final class KafkaMetadataLogTest {
     )
     val log = buildMetadataLog(tempDir, mockTime, config)
 
-    for(_ <- 0 to 1000) {
+    for (_ <- 0 to 1000) {
       append(log, 1, 1)
     }
     log.updateHighWatermark(new LogOffsetMetadata(1001))
 
-    for(offset <- Seq(100, 200, 300, 400, 500, 600)) {
+    for (offset <- Seq(100, 200, 300, 400, 500, 600)) {
       val snapshotId = new OffsetAndEpoch(offset, 1)
       TestUtils.resource(log.storeSnapshot(snapshotId).get()) { snapshot =>
         append(snapshot, 10)
@@ -924,7 +924,7 @@ final class KafkaMetadataLogTest {
     )
     val log = buildMetadataLog(tempDir, mockTime, config)
 
-    for(_ <- 0 to 2000) {
+    for (_ <- 0 to 2000) {
       append(log, 1, 1)
     }
     log.updateHighWatermark(new LogOffsetMetadata(2000))
diff --git a/core/src/test/scala/other/kafka/StressTestLog.scala b/core/src/test/scala/other/kafka/StressTestLog.scala
index ae86fb9215300..4c73155fda49d 100755
--- a/core/src/test/scala/other/kafka/StressTestLog.scala
+++ b/core/src/test/scala/other/kafka/StressTestLog.scala
@@ -70,7 +70,7 @@ object StressTestLog {
         Utils.delete(dir)
     })
 
-    while(running.get) {
+    while (running.get) {
       Thread.sleep(1000)
       println("Reader offset = %d, writer offset = %d".format(reader.currentOffset, writer.currentOffset))
       writer.checkProgress()
@@ -83,7 +83,7 @@ object StressTestLog {
 
     override def run(): Unit = {
       try {
-        while(running.get)
+        while (running.get)
           work()
       } catch {
         case e: Exception => {
diff --git a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
index 061c907d04c79..33a7a8ffc0f29 100755
--- a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
+++ b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
@@ -116,12 +116,12 @@ object TestLinearWriteSpeed {
     val writables = new Array[Writable](numFiles)
     val scheduler = new KafkaScheduler(1)
     scheduler.startup()
-    for(i <- 0 until numFiles) {
-      if(options.has(mmapOpt)) {
+    for (i <- 0 until numFiles) {
+      if (options.has(mmapOpt)) {
         writables(i) = new MmapWritable(new File(dir, "kafka-test-" + i + ".dat"), bytesToWrite / numFiles, buffer)
-      } else if(options.has(channelOpt)) {
+      } else if (options.has(channelOpt)) {
         writables(i) = new ChannelWritable(new File(dir, "kafka-test-" + i + ".dat"), buffer)
-      } else if(options.has(logOpt)) {
+      } else if (options.has(logOpt)) {
         val segmentSize = rand.nextInt(512)*1024*1024 + 64*1024*1024 // vary size to avoid herd effect
         val logProperties = new Properties()
         logProperties.put(TopicConfig.SEGMENT_BYTES_CONFIG, segmentSize: java.lang.Integer)
@@ -143,7 +143,7 @@ object TestLinearWriteSpeed {
     var written = 0L
     var totalWritten = 0L
     var lastReport = beginTest
-    while(totalWritten + bufferSize < bytesToWrite) {
+    while (totalWritten + bufferSize < bytesToWrite) {
       val start = System.nanoTime
       val writeSize = writables((count % numFiles).toInt.abs).write()
       val elapsed = System.nanoTime - start
@@ -152,7 +152,7 @@ object TestLinearWriteSpeed {
       written += writeSize
       count += 1
       totalWritten += writeSize
-      if((start - lastReport)/(1000.0*1000.0) > reportingInterval.doubleValue) {
+      if ((start - lastReport)/(1000.0*1000.0) > reportingInterval.doubleValue) {
         val elapsedSecs = (start - lastReport) / (1000.0*1000.0*1000.0)
         val mb = written / (1024.0*1024.0)
         println("%10.3f\t%10.3f\t%10.3f".format(mb / elapsedSecs, totalLatency / count.toDouble / (1000.0*1000.0), maxLatency / (1000.0 * 1000.0)))
@@ -160,12 +160,12 @@ object TestLinearWriteSpeed {
         written = 0
         maxLatency = 0L
         totalLatency = 0L
-      } else if(written > maxThroughputBytes * (reportingInterval / 1000.0)) {
+      } else if (written > maxThroughputBytes * (reportingInterval / 1000.0)) {
         // if we have written enough, just sit out this reporting interval
         val lastReportMs = lastReport / (1000*1000)
         val now = System.nanoTime / (1000*1000)
         val sleepMs = lastReportMs + reportingInterval - now
-        if(sleepMs > 0)
+        if (sleepMs > 0)
           Thread.sleep(sleepMs)
       }
     }
diff --git a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
index 90bc0ac566169..602bad3d5a8ad 100644
--- a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
@@ -3778,7 +3778,7 @@ class PartitionTest extends AbstractPartitionTest {
     val fetchParams = followerFetchParams(
       replicaId,
       maxBytes = maxBytes,
-      replicaEpoch = if (!replicaEpoch.isDefined) defaultBrokerEpoch(replicaId) else replicaEpoch.get
+      replicaEpoch = if (replicaEpoch.isEmpty) defaultBrokerEpoch(replicaId) else replicaEpoch.get
     )
 
     val fetchPartitionData = new FetchRequest.PartitionData(
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
index 8a7bba87c6bda..32f8753044507 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
@@ -90,7 +90,7 @@ class ControllerIntegrationTest extends QuorumTestHarness {
     waitUntilControllerEpoch(firstControllerEpoch, "broker failed to set controller epoch")
     servers.head.shutdown()
     servers.head.awaitShutdown()
-    TestUtils.waitUntilTrue(() => !zkClient.getControllerId.isDefined, "failed to kill controller")
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isEmpty, "failed to kill controller")
     waitUntilControllerEpoch(firstControllerEpoch, "controller epoch was not persisted after broker failure")
   }
 
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
index 8b5d44a2a364f..17e35ffde65ac 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
@@ -156,7 +156,7 @@ class ProducerIdManagerTest {
 
     for ( _ <- 0 until numThreads) {
       requestHandlerThreadPool.submit(() => {
-        while(latch.getCount > 0) {
+        while (latch.getCount > 0) {
           val result = manager.generateProducerId()
           result match {
             case Success(pid) =>
diff --git a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
index 88cfaae969477..1cf5450afc25d 100755
--- a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
@@ -283,7 +283,7 @@ abstract class KafkaServerTestHarness extends QuorumTestHarness {
     }
     if (configs.isEmpty)
       throw new KafkaException("Must supply at least one server config.")
-    for(i <- _brokers.indices if !alive(i)) {
+    for (i <- _brokers.indices if !alive(i)) {
       if (reconfigure) {
         _brokers(i) = createBrokerFromConfig(configs(i))
       }
diff --git a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
index 95c7237d79f6d..b2115d4279521 100644
--- a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
@@ -144,7 +144,7 @@ abstract class AbstractLogCleanerIntegrationTest {
 
   def writeDups(numKeys: Int, numDups: Int, log: UnifiedLog, codec: CompressionType,
                 startKey: Int = 0, magicValue: Byte = RecordBatch.CURRENT_MAGIC_VALUE): Seq[(Int, String, Long)] = {
-    for(_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
+    for (_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
       val value = counter.toString
       val appendInfo = log.appendAsLeader(TestUtils.singletonRecords(value = value.getBytes, codec = codec,
         key = key.toString.getBytes, magicValue = magicValue), leaderEpoch = 0)
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
index 0f74694245efa..049bc95b37ce3 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
@@ -544,7 +544,7 @@ class LogCleanerManagerTest extends Logging {
 
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
-    while(log.numberOfSegments < 8)
+    while (log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, time.milliseconds()), leaderEpoch = 0)
 
     log.updateHighWatermark(50)
@@ -566,7 +566,7 @@ class LogCleanerManagerTest extends Logging {
 
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
-    while(log.numberOfSegments < 8)
+    while (log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, time.milliseconds()), leaderEpoch = 0)
 
     log.updateHighWatermark(log.logEndOffset)
@@ -590,7 +590,7 @@ class LogCleanerManagerTest extends Logging {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     val t0 = time.milliseconds
-    while(log.numberOfSegments < 4)
+    while (log.numberOfSegments < 4)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, t0), leaderEpoch = 0)
 
     val activeSegAtT0 = log.activeSegment
@@ -705,7 +705,7 @@ class LogCleanerManagerTest extends Logging {
     val logProps = new Properties()
     logProps.put(TopicConfig.SEGMENT_BYTES_CONFIG, 1024: java.lang.Integer)
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
-    while(log.numberOfSegments < 8)
+    while (log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, time.milliseconds()), leaderEpoch = 0)
 
     val cleanerManager: LogCleanerManager = createCleanerManager(log)
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
index bcf9fc02242e6..161493457e1f3 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
@@ -130,7 +130,7 @@ class LogCleanerTest extends Logging {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     // append messages to the log until we have four segments
-    while(log.numberOfSegments < 4)
+    while (log.numberOfSegments < 4)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
     val keysFound = LogTestUtils.keysInLog(log)
     assertEquals(0L until log.logEndOffset, keysFound)
@@ -889,7 +889,7 @@ class LogCleanerTest extends Logging {
 
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
-    while(log.numberOfSegments < 2)
+    while (log.numberOfSegments < 2)
       log.appendAsLeader(record(log.logEndOffset.toInt, Array.fill(largeMessageSize)(0: Byte)), leaderEpoch = 0)
     val keysFound = LogTestUtils.keysInLog(log)
     assertEquals(0L until log.logEndOffset, keysFound)
@@ -961,7 +961,7 @@ class LogCleanerTest extends Logging {
 
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
-    while(log.numberOfSegments < 2)
+    while (log.numberOfSegments < 2)
       log.appendAsLeader(record(log.logEndOffset.toInt, Array.fill(largeMessageSize)(0: Byte)), leaderEpoch = 0)
     val keysFound = LogTestUtils.keysInLog(log)
     assertEquals(0L until log.logEndOffset, keysFound)
@@ -987,16 +987,16 @@ class LogCleanerTest extends Logging {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     // append messages with the keys 0 through N
-    while(log.numberOfSegments < 2)
+    while (log.numberOfSegments < 2)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
     // delete all even keys between 0 and N
     val leo = log.logEndOffset
-    for(key <- 0 until leo.toInt by 2)
+    for (key <- 0 until leo.toInt by 2)
       log.appendAsLeader(tombstoneRecord(key), leaderEpoch = 0)
 
     // append some new unique keys to pad out to a new active segment
-    while(log.numberOfSegments < 4)
+    while (log.numberOfSegments < 4)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0, log.activeSegment.baseOffset))
@@ -1179,14 +1179,14 @@ class LogCleanerTest extends Logging {
     val numTotalSegments = 7
 
     // append messages with the keys 0 through N-1, values equal offset
-    while(log.numberOfSegments <= numCleanableSegments)
+    while (log.numberOfSegments <= numCleanableSegments)
       log.appendAsLeader(record(log.logEndOffset.toInt % N, log.logEndOffset.toInt), leaderEpoch = 0)
 
     // at this point one message past the cleanable segments has been added
     // the entire segment containing the first uncleanable offset should not be cleaned.
     val firstUncleanableOffset = log.logEndOffset + 1  // +1  so it is past the baseOffset
 
-    while(log.numberOfSegments < numTotalSegments - 1)
+    while (log.numberOfSegments < numTotalSegments - 1)
       log.appendAsLeader(record(log.logEndOffset.toInt % N, log.logEndOffset.toInt), leaderEpoch = 0)
 
     // the last (active) segment has just one message
@@ -1269,14 +1269,14 @@ class LogCleanerTest extends Logging {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     // append unkeyed messages
-    while(log.numberOfSegments < 2)
+    while (log.numberOfSegments < 2)
       log.appendAsLeader(unkeyedRecord(log.logEndOffset.toInt), leaderEpoch = 0)
     val numInvalidMessages = unkeyedMessageCountInLog(log)
 
     val sizeWithUnkeyedMessages = log.size
 
     // append keyed messages
-    while(log.numberOfSegments < 3)
+    while (log.numberOfSegments < 3)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
     val expectedSizeAfterCleaning = log.size - sizeWithUnkeyedMessages
@@ -1326,7 +1326,7 @@ class LogCleanerTest extends Logging {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     // append messages to the log until we have four segments
-    while(log.numberOfSegments < 4)
+    while (log.numberOfSegments < 4)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
     val keys = LogTestUtils.keysInLog(log)
@@ -1352,7 +1352,7 @@ class LogCleanerTest extends Logging {
 
     // append some messages to the log
     var i = 0
-    while(log.numberOfSegments < 10) {
+    while (log.numberOfSegments < 10) {
       log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0)
       i += 1
     }
@@ -1399,7 +1399,7 @@ class LogCleanerTest extends Logging {
     val v="val".getBytes()
 
     //create 3 segments
-    for(i <- 0 until 3){
+    for (i <- 0 until 3) {
       log.appendAsLeader(TestUtils.singletonRecords(value = v, key = k), leaderEpoch = 0)
       //0 to Int.MaxValue is Int.MaxValue+1 message, -1 will be the last message of i-th segment
       val records = messageWithOffset(k, v, (i + 1L) * (Int.MaxValue + 1L) -1 )
@@ -1552,7 +1552,7 @@ class LogCleanerTest extends Logging {
       val endOffset = map.latestOffset + 1
       assertEquals(end, endOffset, "Last offset should be the end offset.")
       assertEquals(end-start, map.size, "Should have the expected number of messages in the map.")
-      for(i <- start until end)
+      for (i <- start until end)
         assertEquals(i.toLong, map.get(key(i)), "Should find all the keys")
       assertEquals(-1L, map.get(key(start - 1)), "Should not find a value too small")
       assertEquals(-1L, map.get(key(end)), "Should not find a value too large")
@@ -1987,7 +1987,7 @@ class LogCleanerTest extends Logging {
   }
 
   private def writeToLog(log: UnifiedLog, keysAndValues: Iterable[(Int, Int)], offsetSeq: Iterable[Long]): Iterable[Long] = {
-    for(((key, value), offset) <- keysAndValues.zip(offsetSeq))
+    for (((key, value), offset) <- keysAndValues.zip(offsetSeq))
       yield log.appendAsFollower(messageWithOffset(key, value, offset)).lastOffset
   }
 
@@ -2158,7 +2158,7 @@ class FakeOffsetMap(val slots: Int) extends OffsetMap {
 
   override def get(key: ByteBuffer): Long = {
     val k = keyFor(key)
-    if(map.containsKey(k))
+    if (map.containsKey(k))
       map.get(k)
     else
       -1L
diff --git a/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala b/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
index f44285bd908df..5dcca54961bb8 100644
--- a/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
@@ -792,7 +792,7 @@ class LogLoaderTest {
     val indexInterval = 3 * messageSize
     val logConfig = LogTestUtils.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = indexInterval, segmentIndexBytes = 4096)
     var log = createLog(logDir, logConfig)
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(messageSize),
         timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     assertEquals(numMessages, log.logEndOffset,
@@ -842,7 +842,7 @@ class LogLoaderTest {
     val numMessages = 200
     val logConfig = LogTestUtils.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
     var log = createLog(logDir, logConfig)
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(10), timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     val indexFiles = log.logSegments.asScala.map(_.offsetIndexFile)
     val timeIndexFiles = log.logSegments.asScala.map(_.timeIndexFile)
@@ -857,7 +857,7 @@ class LogLoaderTest {
     assertEquals(numMessages, log.logEndOffset, "Should have %d messages when log is reopened".format(numMessages))
     assertTrue(log.logSegments.asScala.head.offsetIndex.entries > 0, "The index should have been rebuilt")
     assertTrue(log.logSegments.asScala.head.timeIndex.entries > 0, "The time index should have been rebuilt")
-    for(i <- 0 until numMessages) {
+    for (i <- 0 until numMessages) {
       assertEquals(i, LogTestUtils.readLog(log, i, 100).records.batches.iterator.next().lastOffset)
       if (i == 0)
         assertEquals(log.logSegments.asScala.head.baseOffset, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
@@ -908,21 +908,21 @@ class LogLoaderTest {
     val numMessages = 200
     val logConfig = LogTestUtils.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
     var log = createLog(logDir, logConfig)
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(10), timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     val indexFiles = log.logSegments.asScala.map(_.offsetIndexFile())
     val timeIndexFiles = log.logSegments.asScala.map(_.timeIndexFile())
     log.close()
 
     // corrupt all the index files
-    for( file <- indexFiles) {
+    for ( file <- indexFiles) {
       val bw = new BufferedWriter(new FileWriter(file))
       bw.write("  ")
       bw.close()
     }
 
     // corrupt all the index files
-    for( file <- timeIndexFiles) {
+    for ( file <- timeIndexFiles) {
       val bw = new BufferedWriter(new FileWriter(file))
       bw.write("  ")
       bw.close()
@@ -931,7 +931,7 @@ class LogLoaderTest {
     // reopen the log with recovery point=0 so that the segment recovery can be triggered
     log = createLog(logDir, logConfig, lastShutdownClean = false)
     assertEquals(numMessages, log.logEndOffset, "Should have %d messages when log is reopened".format(numMessages))
-    for(i <- 0 until numMessages) {
+    for (i <- 0 until numMessages) {
       assertEquals(i, LogTestUtils.readLog(log, i, 100).records.batches.iterator.next().lastOffset)
       if (i == 0)
         assertEquals(log.logSegments.asScala.head.baseOffset, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
@@ -1728,7 +1728,7 @@ class LogLoaderTest {
     val indexInterval = 3 * messageSize
     val logConfig = LogTestUtils.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = indexInterval, segmentIndexBytes = 4096)
     var log = createLog(logDir, logConfig)
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(messageSize),
         timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     assertEquals(numMessages, log.logEndOffset,
@@ -1745,7 +1745,7 @@ class LogLoaderTest {
     assertEquals(0, log.activeSegment.timeIndex.entries, "Should have same number of time index entries as before.")
     log.activeSegment.sanityCheck(true) // this should not throw because the LogLoader created the empty active log index file during recovery
 
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(messageSize),
         timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     log.roll()
diff --git a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
index 3aca8536bdb68..2dcc3e5a35cc0 100755
--- a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
@@ -353,7 +353,7 @@ class LogManagerTest {
   def testCleanupExpiredSegments(): Unit = {
     val log = logManager.getOrCreateLog(new TopicPartition(name, 0), topicId = None)
     var offset = 0L
-    for(_ <- 0 until 200) {
+    for (_ <- 0 until 200) {
       val set = TestUtils.singletonRecords("test".getBytes())
       val info = log.appendAsLeader(set, leaderEpoch = 0)
       offset = info.lastOffset
@@ -499,7 +499,7 @@ class LogManagerTest {
     logManager = createLogManager(dirs)
 
     // verify that logs are always assigned to the least loaded partition
-    for(partition <- 0 until 20) {
+    for (partition <- 0 until 20) {
       logManager.getOrCreateLog(new TopicPartition("test", partition), topicId = None)
       assertEquals(partition + 1, logManager.allLogs.size, "We should have created the right number of logs")
       val counts = logManager.allLogs.groupBy(_.dir.getParent).values.map(_.size)
diff --git a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
index ebdcab0ccf98b..b4e278c380230 100644
--- a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
@@ -324,12 +324,12 @@ class LogSegmentTest {
   @Test
   def testRecoveryFixesCorruptIndex(): Unit = {
     val seg = createSegment(0)
-    for(i <- 0 until 100)
+    for (i <- 0 until 100)
       seg.append(i, RecordBatch.NO_TIMESTAMP, -1L, records(i, i.toString))
     val indexFile = seg.offsetIndexFile
     TestUtils.writeNonsenseToFile(indexFile, 5, indexFile.length.toInt)
     seg.recover(newProducerStateManager(), Optional.empty())
-    for(i <- 0 until 100) {
+    for (i <- 0 until 100) {
       val records = seg.read(i, 1, seg.size(), true).records.records
       assertEquals(i, records.iterator.next().offset)
     }
@@ -450,12 +450,12 @@ class LogSegmentTest {
   @Test
   def testRecoveryFixesCorruptTimeIndex(): Unit = {
     val seg = createSegment(0)
-    for(i <- 0 until 100)
+    for (i <- 0 until 100)
       seg.append(i, i * 10, i, records(i, i.toString))
     val timeIndexFile = seg.timeIndexFile
     TestUtils.writeNonsenseToFile(timeIndexFile, 5, timeIndexFile.length.toInt)
     seg.recover(newProducerStateManager(), Optional.empty())
-    for(i <- 0 until 100) {
+    for (i <- 0 until 100) {
       assertEquals(i, seg.findOffsetByTimestamp(i * 10, 0L).get.offset)
       if (i < 99)
         assertEquals(i + 1, seg.findOffsetByTimestamp(i * 10 + 1, 0L).get.offset)
diff --git a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
index 40b6195cc4641..0c2afa02162a5 100644
--- a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
@@ -45,7 +45,7 @@ class OffsetIndexTest {
   
   @AfterEach
   def teardown(): Unit = {
-    if(this.idx != null)
+    if (this.idx != null)
       this.idx.file.delete()
   }
 
@@ -62,7 +62,7 @@ class OffsetIndexTest {
     vals.foreach{x => idx.append(x._1, x._2)}
     
     // should be able to find all those values
-    for((logical, physical) <- vals)
+    for ((logical, physical) <- vals)
       assertEquals(new OffsetPosition(logical, physical), idx.lookup(logical),
         "Should be able to find values that are present.")
       
@@ -70,9 +70,9 @@ class OffsetIndexTest {
     val valMap = new immutable.TreeMap[Long, (Long, Int)]() ++ vals.map(p => (p._1, p))
     val offsets = (idx.baseOffset until vals.last._1.toInt).toArray
     Collections.shuffle(Arrays.asList(offsets))
-    for(offset <- offsets.take(30)) {
+    for (offset <- offsets.take(30)) {
       val rightAnswer = 
-        if(offset < valMap.firstKey)
+        if (offset < valMap.firstKey)
           new OffsetPosition(idx.baseOffset, 0)
         else
           new OffsetPosition(valMap.to(offset).last._1, valMap.to(offset).last._2._2)
@@ -85,7 +85,7 @@ class OffsetIndexTest {
   def lookupExtremeCases(): Unit = {
     assertEquals(new OffsetPosition(idx.baseOffset, 0), idx.lookup(idx.baseOffset),
       "Lookup on empty file")
-    for(i <- 0 until idx.maxEntries)
+    for (i <- 0 until idx.maxEntries)
       idx.append(idx.baseOffset + i + 1, i)
     // check first and last entry
     assertEquals(new OffsetPosition(idx.baseOffset, 0), idx.lookup(idx.baseOffset))
@@ -107,7 +107,7 @@ class OffsetIndexTest {
   
   @Test
   def appendTooMany(): Unit = {
-    for(i <- 0 until idx.maxEntries) {
+    for (i <- 0 until idx.maxEntries) {
       val offset = idx.baseOffset + i + 1
       idx.append(offset, i)
     }
@@ -161,7 +161,7 @@ class OffsetIndexTest {
   def truncate(): Unit = {
     val idx = new OffsetIndex(nonExistentTempFile(), 0L, 10 * 8)
     idx.truncate()
-    for(i <- 1 until 10)
+    for (i <- 1 until 10)
       idx.append(i, i)
       
     // now check the last offset after various truncate points and validate that we can still append to the index.      
diff --git a/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala b/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
index f3d584f5d1287..ccdf69490e4ac 100644
--- a/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
+++ b/core/src/test/scala/unit/kafka/log/OffsetMapTest.scala
@@ -36,12 +36,12 @@ class OffsetMapTest {
   @Test
   def testClear(): Unit = {
     val map = new SkimpyOffsetMap(4000)
-    for(i <- 0 until 10)
+    for (i <- 0 until 10)
       map.put(key(i), i)
-    for(i <- 0 until 10)
+    for (i <- 0 until 10)
       assertEquals(i.toLong, map.get(key(i)))
     map.clear()
-    for(i <- 0 until 10)
+    for (i <- 0 until 10)
       assertEquals(map.get(key(i)), -1L)
   }
   
@@ -61,9 +61,9 @@ class OffsetMapTest {
   
   def validateMap(items: Int, loadFactor: Double = 0.5): SkimpyOffsetMap = {
     val map = new SkimpyOffsetMap((items/loadFactor * 24).toInt)
-    for(i <- 0 until items)
+    for (i <- 0 until items)
       map.put(key(i), i)
-    for(i <- 0 until items)
+    for (i <- 0 until items)
       assertEquals(map.get(key(i)), i.toLong)
     map
   }
@@ -72,7 +72,7 @@ class OffsetMapTest {
 
 object OffsetMapTest {
   def main(args: Array[String]): Unit = {
-    if(args.length != 2) {
+    if (args.length != 2) {
       System.err.println("USAGE: java OffsetMapTest size load")
       Exit.exit(1)
     }
diff --git a/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala b/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
index 72bbe3cd202cf..e356ef160ec04 100644
--- a/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
@@ -39,7 +39,7 @@ class TimeIndexTest {
 
   @AfterEach
   def teardown(): Unit = {
-    if(this.idx != null)
+    if (this.idx != null)
       this.idx.file.delete()
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala b/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
index 5a476574e183f..7b92a9e2df794 100755
--- a/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
+++ b/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
@@ -80,7 +80,7 @@ class UnifiedLogTest {
   }
 
   def createEmptyLogs(dir: File, offsets: Int*): Unit = {
-    for(offset <- offsets) {
+    for (offset <- offsets) {
       Files.createFile(LogFileUtils.logFile(dir, offset).toPath)
       Files.createFile(LogFileUtils.offsetIndexFile(dir, offset).toPath)
     }
@@ -1545,10 +1545,10 @@ class UnifiedLogTest {
     val log = createLog(logDir, logConfig)
     val values = (0 until 100 by 2).map(id => id.toString.getBytes).toArray
 
-    for(value <- values)
+    for (value <- values)
       log.appendAsLeader(TestUtils.singletonRecords(value = value), leaderEpoch = 0)
 
-    for(i <- values.indices) {
+    for (i <- values.indices) {
       val read = LogTestUtils.readLog(log, i, 1).records.batches.iterator.next()
       assertEquals(i, read.lastOffset, "Offset read should match order appended.")
       val actual = read.iterator.next()
@@ -1571,9 +1571,9 @@ class UnifiedLogTest {
     val records = messageIds.map(id => new SimpleRecord(id.toString.getBytes))
 
     // now test the case that we give the offsets and use non-sequential offsets
-    for(i <- records.indices)
+    for (i <- records.indices)
       log.appendAsFollower(MemoryRecords.withRecords(messageIds(i), CompressionType.NONE, 0, records(i)))
-    for(i <- 50 until messageIds.max) {
+    for (i <- 50 until messageIds.max) {
       val idx = messageIds.indexWhere(_ >= i)
       val read = LogTestUtils.readLog(log, i, 100).records.records.iterator.next()
       assertEquals(messageIds(idx), read.offset, "Offset read should match message id.")
@@ -1593,7 +1593,7 @@ class UnifiedLogTest {
     val log = createLog(logDir, logConfig)
 
     // keep appending until we have two segments with only a single message in the second segment
-    while(log.numberOfSegments == 1)
+    while (log.numberOfSegments == 1)
       log.appendAsLeader(TestUtils.singletonRecords(value = "42".getBytes), leaderEpoch = 0)
 
     // now manually truncate off all but one message from the first segment to create a gap in the messages
@@ -1715,7 +1715,7 @@ class UnifiedLogTest {
 
     /* do successive reads to ensure all our messages are there */
     var offset = 0L
-    for(i <- 0 until numMessages) {
+    for (i <- 0 until numMessages) {
       val messages = LogTestUtils.readLog(log, offset, 1024*1024).records.batches
       val head = messages.iterator.next()
       assertEquals(offset, head.lastOffset, "Offsets not equal")
@@ -1731,7 +1731,7 @@ class UnifiedLogTest {
     assertEquals(0, lastRead.records.asScala.size, "Should be no more messages")
 
     // check that rolling the log forced a flushed, the flush is async so retry in case of failure
-    TestUtils.retry(1000L){
+    TestUtils.retry(1000L) {
       assertTrue(log.recoveryPoint >= log.activeSegment.baseOffset, "Log role should have forced flush")
     }
   }
@@ -1763,12 +1763,12 @@ class UnifiedLogTest {
    */
   @Test
   def testThatGarbageCollectingSegmentsDoesntChangeOffset(): Unit = {
-    for(messagesToAppend <- List(0, 1, 25)) {
+    for (messagesToAppend <- List(0, 1, 25)) {
       logDir.mkdirs()
       // first test a log segment starting at 0
       val logConfig = LogTestUtils.createLogConfig(segmentBytes = 100, retentionMs = 0)
       val log = createLog(logDir, logConfig)
-      for(i <- 0 until messagesToAppend)
+      for (i <- 0 until messagesToAppend)
         log.appendAsLeader(TestUtils.singletonRecords(value = i.toString.getBytes, timestamp = mockTime.milliseconds - 10), leaderEpoch = 0)
 
       val currOffset = log.logEndOffset
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
index 06244de7dee68..80790ab905fdb 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
@@ -1112,12 +1112,12 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
                                  resource: ResourcePattern = resource): Set[AccessControlEntry] = {
     var acls = originalAcls
 
-    if(addedAcls.nonEmpty) {
+    if (addedAcls.nonEmpty) {
       addAcls(authorizer1, addedAcls, resource)
       acls ++= addedAcls
     }
 
-    if(removedAcls.nonEmpty) {
+    if (removedAcls.nonEmpty) {
       removeAcls(authorizer1, removedAcls, resource)
       acls --=removedAcls
     }
diff --git a/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala
index 376cd40c6940d..df049b68a8762 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala
@@ -111,7 +111,7 @@ abstract class AbstractApiVersionsRequestTest(cluster: ClusterInstance) {
 
     val defaultApiVersionsResponse = if (!cluster.isKRaftTest) {
       TestUtils.defaultApiVersionsResponse(0, ListenerType.ZK_BROKER, enableUnstableLastVersion)
-    } else if(cluster.controllerListenerName().asScala.contains(listenerName)) {
+    } else if (cluster.controllerListenerName().asScala.contains(listenerName)) {
       TestUtils.defaultApiVersionsResponse(0, ListenerType.CONTROLLER, enableUnstableLastVersion)
     } else {
       TestUtils.createApiVersionsResponse(0, expectedApis)
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
index 9533f4529f3f9..1384ae3901ed9 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
@@ -958,7 +958,7 @@ class AbstractFetcherThreadTest {
     val mockTierStateMachine = new MockTierStateMachine(mockLeaderEndpoint)
     val fetcherForTruncation = new MockFetcherThread(mockLeaderEndpoint, mockTierStateMachine, failedPartitions = failedPartitions) {
       override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
-        if(topicPartition == partition1)
+        if (topicPartition == partition1)
           throw new Exception()
         else {
           super.truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState)
diff --git a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
index ff14323b019f7..344d912d58b34 100644
--- a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
@@ -237,7 +237,7 @@ class DeleteTopicsRequestTest extends BaseRequestTest with Logging {
   private def validateTopicIsDeleted(topic: String): Unit = {
     val metadata = connectAndReceive[MetadataResponse](new MetadataRequest.Builder(
       List(topic).asJava, true).build).topicMetadata.asScala
-    TestUtils.waitUntilTrue (() => !metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE),
+    TestUtils.waitUntilTrue(() => !metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE),
       s"The topic $topic should not exist")
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala b/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
index 360614e93a459..3d3d5732247ee 100644
--- a/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
@@ -67,7 +67,7 @@ class ThrottledChannelExpirationTest {
       delayQueue.add(channel4)
       assertEquals(4, numCallbacksForStartThrottling)
 
-      for(itr <- 1 to 3) {
+      for (itr <- 1 to 3) {
         time.sleep(10)
         reaper.doWork()
         assertEquals(itr, numCallbacksForEndThrottling)
@@ -91,7 +91,7 @@ class ThrottledChannelExpirationTest {
     assertEquals(20, t2.throttleTimeMs)
     assertEquals(20, t3.throttleTimeMs)
 
-    for(itr <- 0 to 2) {
+    for (itr <- 0 to 2) {
       assertEquals(10 - 10*itr, t1.getDelay(TimeUnit.MILLISECONDS))
       assertEquals(20 - 10*itr, t2.getDelay(TimeUnit.MILLISECONDS))
       assertEquals(20 - 10*itr, t3.getDelay(TimeUnit.MILLISECONDS))
diff --git a/core/src/test/scala/unit/kafka/server/TopicIdWithOldInterBrokerProtocolTest.scala b/core/src/test/scala/unit/kafka/server/TopicIdWithOldInterBrokerProtocolTest.scala
index 25c7c7bd04c4a..2ca248b1500ba 100644
--- a/core/src/test/scala/unit/kafka/server/TopicIdWithOldInterBrokerProtocolTest.scala
+++ b/core/src/test/scala/unit/kafka/server/TopicIdWithOldInterBrokerProtocolTest.scala
@@ -172,7 +172,7 @@ class TopicIdWithOldInterBrokerProtocolTest extends BaseRequestTest {
   private def validateTopicIsDeleted(topic: String): Unit = {
     val metadata = connectAndReceive[MetadataResponse](new MetadataRequest.Builder(
       List(topic).asJava, true).build).topicMetadata.asScala
-    TestUtils.waitUntilTrue (() => !metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE),
+    TestUtils.waitUntilTrue(() => !metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE),
       s"The topic $topic should not exist")
   }
 
diff --git a/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
index 11f2752384a53..a75041b49e161 100644
--- a/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
@@ -416,7 +416,7 @@ class DumpLogSegmentsTest {
     var batchesBytes = 0
     var batchesCounter = 0
     while (lines.hasNext) {
-      if (batchesCounter >= limit){
+      if (batchesCounter >= limit) {
         return batchesBytes
       }
       val line = lines.next()
diff --git a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
index f1c0dd0371f0a..8c360c49990db 100755
--- a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
@@ -45,7 +45,7 @@ class CoreUtilsTest extends Logging {
 
   @Test
   def testReadBytes(): Unit = {
-    for(testCase <- List("", "a", "abcd")) {
+    for (testCase <- List("", "a", "abcd")) {
       val bytes = testCase.getBytes
       assertTrue(Arrays.equals(bytes, Utils.readBytes(ByteBuffer.wrap(bytes))))
     }
diff --git a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
index 0bf8fe347628d..6faec45010712 100644
--- a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
@@ -105,7 +105,7 @@ class SchedulerTest {
   @Test
   def testPeriodicTask(): Unit = {
     scheduler.schedule("test", () => counter1.getAndIncrement(), 0, 5)
-    retry(30000){
+    retry(30000) {
       assertTrue(counter1.get >= 20, "Should count to 20")
     }
   }
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 19eac0bcfdb8d..7b1ce80cbed3b 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -752,7 +752,7 @@ object TestUtils extends Logging {
    */
   def checkEquals(b1: ByteBuffer, b2: ByteBuffer): Unit = {
     assertEquals(b1.limit() - b1.position(), b2.limit() - b2.position(), "Buffers should have equal length")
-    for(i <- 0 until b1.limit() - b1.position())
+    for (i <- 0 until b1.limit() - b1.position())
       assertEquals(b1.get(b1.position() + i), b2.get(b1.position() + i), "byte " + i + " byte not equal.")
   }
 
@@ -774,7 +774,7 @@ object TestUtils extends Logging {
    * different messages on their Nth element
    */
   def checkEquals[T](s1: java.util.Iterator[T], s2: java.util.Iterator[T]): Unit = {
-    while(s1.hasNext && s2.hasNext)
+    while (s1.hasNext && s2.hasNext)
       assertEquals(s1.next, s2.next)
     assertFalse(s1.hasNext, "Iterators have uneven length--first has more")
     assertFalse(s2.hasNext, "Iterators have uneven length--second has more")
@@ -815,7 +815,7 @@ object TestUtils extends Logging {
    */
   def hexString(buffer: ByteBuffer): String = {
     val builder = new StringBuilder("0x")
-    for(i <- 0 until buffer.limit())
+    for (i <- 0 until buffer.limit())
       builder.append(String.format("%x", Integer.valueOf(buffer.get(buffer.position() + i))))
     builder.toString
   }
@@ -1086,7 +1086,7 @@ object TestUtils extends Logging {
   def retry(maxWaitMs: Long)(block: => Unit): Unit = {
     var wait = 1L
     val startTime = System.currentTimeMillis()
-    while(true) {
+    while (true) {
       try {
         block
         return
diff --git a/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
index 0ac2e22477760..18147683a30fb 100644
--- a/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/AdminZkClientTest.scala
@@ -123,7 +123,7 @@ class AdminZkClientTest extends QuorumTestHarness with Logging with RackAwareTes
     TestUtils.makeLeaderForPartition(zkClient, topic, leaderForPartitionMap, 1)
     val actualReplicaMap = leaderForPartitionMap.keys.map(p => p -> zkClient.getReplicasForPartition(new TopicPartition(topic, p))).toMap
     assertEquals(expectedReplicaAssignment.size, actualReplicaMap.size)
-    for(i <- 0 until actualReplicaMap.size)
+    for (i <- 0 until actualReplicaMap.size)
       assertEquals(expectedReplicaAssignment.get(i).get, actualReplicaMap(i))
 
     // shouldn't be able to create a topic that already exists
diff --git a/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java b/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
index 0885f4e4fbc4c..241e2d3aae99f 100644
--- a/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
+++ b/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
@@ -51,8 +51,6 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
-import static org.apache.kafka.metadata.AssignmentsHelper.buildRequestData;
-
 public class AssignmentsManager {
 
     private static final Logger log = LoggerFactory.getLogger(AssignmentsManager.class);

From 9865d54c425d53610f2ad5d442a107f1882673ac Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Fri, 9 Feb 2024 04:29:04 -0800
Subject: [PATCH 017/258] MINOR: EventAccumulator should signal one thread when
 key becomes available (#15340)

`signalAll` was mistakenly used instead of `signal` when a key become available in the `EventAccumulator`. The fix relies on existing tests.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../kafka/coordinator/group/runtime/EventAccumulator.java       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
index 35ea0f909837b..f46e8b8a8bfe3 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
@@ -235,7 +235,7 @@ public void close() {
      */
     private void addAvailableKey(K key) {
         availableKeys.add(key);
-        condition.signalAll();
+        condition.signal();
     }
 
     /**

From 092dc7fc467ed7d354ec504d6939b3fcd7b80632 Mon Sep 17 00:00:00 2001
From: Mickael Maison <mimaison@users.noreply.github.com>
Date: Fri, 9 Feb 2024 13:36:56 +0100
Subject: [PATCH 018/258] KAFKA-16238: Fix ConnectRestApiTest system test
 (#15346)

Reviewers: Yash Mayya <yash.mayya@gmail.com>
---
 tests/kafkatest/tests/connect/connect_rest_test.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/kafkatest/tests/connect/connect_rest_test.py b/tests/kafkatest/tests/connect/connect_rest_test.py
index e3d023c9dfb6a..da3e5be534bb9 100644
--- a/tests/kafkatest/tests/connect/connect_rest_test.py
+++ b/tests/kafkatest/tests/connect/connect_rest_test.py
@@ -36,11 +36,13 @@ class ConnectRestApiTest(KafkaTest):
     FILE_SOURCE_CONFIGS = {'name', 'connector.class', 'tasks.max', 'key.converter', 'value.converter', 'header.converter', 'batch.size',
                            'topic', 'file', 'transforms', 'config.action.reload', 'errors.retry.timeout', 'errors.retry.delay.max.ms',
                            'errors.tolerance', 'errors.log.enable', 'errors.log.include.messages', 'predicates', 'topic.creation.groups',
-                           'exactly.once.support', 'transaction.boundary', 'transaction.boundary.interval.ms', 'offsets.storage.topic'}
+                           'exactly.once.support', 'transaction.boundary', 'transaction.boundary.interval.ms', 'offsets.storage.topic',
+                           'tasks.max.enforce'}
     FILE_SINK_CONFIGS = {'name', 'connector.class', 'tasks.max', 'key.converter', 'value.converter', 'header.converter', 'topics',
                          'file', 'transforms', 'topics.regex', 'config.action.reload', 'errors.retry.timeout', 'errors.retry.delay.max.ms',
                          'errors.tolerance', 'errors.log.enable', 'errors.log.include.messages', 'errors.deadletterqueue.topic.name',
-                         'errors.deadletterqueue.topic.replication.factor', 'errors.deadletterqueue.context.headers.enable', 'predicates'}
+                         'errors.deadletterqueue.topic.replication.factor', 'errors.deadletterqueue.context.headers.enable', 'predicates',
+                         'tasks.max.enforce'}
 
     INPUT_FILE = "/mnt/connect.input"
     INPUT_FILE2 = "/mnt/connect.input2"

From ec4a8aaadbc95cfcf0de2f5e1385373f095298ca Mon Sep 17 00:00:00 2001
From: Vedarth Sharma <142404391+VedarthConfluent@users.noreply.github.com>
Date: Fri, 9 Feb 2024 18:30:20 +0530
Subject: [PATCH 019/258] [MINOR] Fix docker image build by introducing bash
 (#15347)

The base eclipse-temuring:21-jre-alpine image got modified and had `bash` removed from it. This broke our build, since downstream steps utilizing bash scripts depended on it. This patch explicitly installs bash
---
 docker/jvm/Dockerfile | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docker/jvm/Dockerfile b/docker/jvm/Dockerfile
index f16084b43f5d8..df552ad019abd 100644
--- a/docker/jvm/Dockerfile
+++ b/docker/jvm/Dockerfile
@@ -28,7 +28,7 @@ COPY jsa_launch /etc/kafka/docker/jsa_launch
 RUN set -eux ; \
     apk update ; \
     apk upgrade ; \
-    apk add --no-cache wget gcompat gpg gpg-agent procps netcat-openbsd; \
+    apk add --no-cache wget gcompat gpg gpg-agent procps netcat-openbsd bash; \
     mkdir opt/kafka; \
     wget -nv -O kafka.tgz "$kafka_url"; \
     wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \
@@ -62,7 +62,7 @@ LABEL org.label-schema.name="kafka" \
 RUN set -eux ; \
     apk update ; \
     apk upgrade ; \
-    apk add --no-cache wget gpg gpg-agent gcompat; \
+    apk add --no-cache wget gpg gpg-agent gcompat bash; \
     mkdir opt/kafka; \
     wget -nv -O kafka.tgz "$kafka_url"; \
     wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \

From b25c96a91599ea029eae77c55cffaeafcc87b8a4 Mon Sep 17 00:00:00 2001
From: Jorge Esteban Quilcate Otoya <jorge.quilcate@aiven.io>
Date: Fri, 9 Feb 2024 20:17:17 -0500
Subject: [PATCH 020/258] KAFKA-16229: Fix slow expired producer id deletion
 (#15324)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Expiration of ProducerIds is implemented with a slow removal of map keys:
        producers.keySet().removeAll(keys);
Unnecessarily going through all producer ids and then throw all expired keys to be removed.
This leads to exponential time on worst case when most/all keys need to be removed:

Benchmark                                        (numProducerIds)  Mode  Cnt           Score            Error  Units
ProducerStateManagerBench.testDeleteExpiringIds               100  avgt    3        9164.043 ±      10647.877  ns/op
ProducerStateManagerBench.testDeleteExpiringIds              1000  avgt    3      341561.093 ±      20283.211  ns/op
ProducerStateManagerBench.testDeleteExpiringIds             10000  avgt    3    44957983.550 ±    9389011.290  ns/op
ProducerStateManagerBench.testDeleteExpiringIds            100000  avgt    3  5683374164.167 ± 1446242131.466  ns/op
A simple fix is to use map#remove(key) instead, leading to a more linear growth:

Benchmark                                        (numProducerIds)  Mode  Cnt        Score         Error  Units
ProducerStateManagerBench.testDeleteExpiringIds               100  avgt    3     5779.056 ±     651.389  ns/op
ProducerStateManagerBench.testDeleteExpiringIds              1000  avgt    3    61430.530 ±   21875.644  ns/op
ProducerStateManagerBench.testDeleteExpiringIds             10000  avgt    3   643887.031 ±  600475.302  ns/op
ProducerStateManagerBench.testDeleteExpiringIds            100000  avgt    3  7741689.539 ± 3218317.079  ns/op
Flamegraph of the CPU usage at dealing with expiration when producers ids ~1Million:

Reviewers: Justine Olshan <jolshan@confluent.io>
---
 .../storage/ProducerStateManagerBench.java    | 99 +++++++++++++++++++
 .../internals/log/ProducerStateManager.java   |  2 +-
 2 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 jmh-benchmarks/src/main/java/org/apache/kafka/jmh/storage/ProducerStateManagerBench.java

diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/storage/ProducerStateManagerBench.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/storage/ProducerStateManagerBench.java
new file mode 100644
index 0000000000000..291c78d72adb0
--- /dev/null
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/storage/ProducerStateManagerBench.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.jmh.storage;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.storage.internals.log.ProducerStateEntry;
+import org.apache.kafka.storage.internals.log.ProducerStateManager;
+import org.apache.kafka.storage.internals.log.ProducerStateManagerConfig;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Optional;
+import java.util.OptionalLong;
+import java.util.concurrent.TimeUnit;
+
+@Warmup(iterations = 2)
+@Measurement(iterations = 3)
+@Fork(1)
+@BenchmarkMode(Mode.AverageTime)
+@State(value = Scope.Benchmark)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+public class ProducerStateManagerBench {
+    Time time = new MockTime();
+    final int producerIdExpirationMs = 1000;
+
+    ProducerStateManager manager;
+    Path tempDirectory;
+
+    @Param({"100", "1000", "10000", "100000"})
+    public int numProducerIds;
+
+    @Setup(Level.Trial)
+    public void setup() throws IOException {
+        tempDirectory = Files.createTempDirectory("kafka-logs");
+        manager = new ProducerStateManager(
+            new TopicPartition("t1", 0),
+            tempDirectory.toFile(),
+            Integer.MAX_VALUE,
+            new ProducerStateManagerConfig(producerIdExpirationMs, false),
+            time
+        );
+    }
+
+
+    @TearDown(Level.Trial)
+    public void tearDown() throws Exception {
+        Files.deleteIfExists(tempDirectory);
+    }
+
+    @Benchmark
+    @Threads(1)
+    public void testDeleteExpiringIds() {
+        short epoch = 0;
+        for (long i = 0L; i < numProducerIds; i++) {
+            final ProducerStateEntry entry = new ProducerStateEntry(
+                i,
+                epoch,
+                0,
+                time.milliseconds(),
+                OptionalLong.empty(),
+                Optional.empty()
+            );
+            manager.loadProducerEntry(entry);
+        }
+
+        manager.removeExpiredProducers(time.milliseconds() + producerIdExpirationMs + 1);
+    }
+}
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
index 6bcafd2d60763..270aa0a42f94c 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
@@ -177,7 +177,7 @@ private void addProducerId(long producerId, ProducerStateEntry entry) {
     }
 
     private void removeProducerIds(List<Long> keys) {
-        producers.keySet().removeAll(keys);
+        keys.forEach(producers::remove);
         producerIdCount = producers.size();
     }
 

From d233eb98f7c7e55fe0dd673dbc058ddf619663a7 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Sun, 11 Feb 2024 04:46:51 +0800
Subject: [PATCH 021/258] KAFKA-14957: Update-Description-String (#13909)

HTML code for configs is auto-generated and for Kafka Streams config `state.dir` produces a confusing default value.
This PR adds a new property `alternativeString` to set a "default" value which should be rendered in HTML instead of the actual default value.

Reviewers: Manyanda Chitimbo <manyanda.chitimbo@gmail.com>, @eziosudo <eziosudo@gmail.com>, Matthias J. Sax <matthias@confluent.io>
---
 .../apache/kafka/common/config/ConfigDef.java | 55 +++++++++++++++++--
 .../kafka/common/config/ConfigDefTest.java    |  2 +-
 .../connect/runtime/AbstractHerderTest.java   |  2 +-
 .../apache/kafka/streams/StreamsConfig.java   |  5 +-
 4 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/config/ConfigDef.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigDef.java
index 6f2e7413ccc67..57df493347b48 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/ConfigDef.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigDef.java
@@ -152,7 +152,30 @@ public ConfigDef define(ConfigKey key) {
      */
     public ConfigDef define(String name, Type type, Object defaultValue, Validator validator, Importance importance, String documentation,
                             String group, int orderInGroup, Width width, String displayName, List<String> dependents, Recommender recommender) {
-        return define(new ConfigKey(name, type, defaultValue, validator, importance, documentation, group, orderInGroup, width, displayName, dependents, recommender, false));
+        return define(new ConfigKey(name, type, defaultValue, validator, importance, documentation, group, orderInGroup, width, displayName, dependents, recommender, false, null));
+    }
+
+    /**
+     * Define a new configuration
+     * @param name               the name of the config parameter
+     * @param type               the type of the config
+     * @param defaultValue       the default value to use if this config isn't present
+     * @param validator          the validator to use in checking the correctness of the config
+     * @param importance         the importance of this config
+     * @param documentation      the documentation string for the config
+     * @param group              the group this config belongs to
+     * @param orderInGroup       the order of this config in the group
+     * @param width              the width of the config
+     * @param displayName        the name suitable for display
+     * @param dependents         the configurations that are dependents of this configuration
+     * @param recommender        the recommender provides valid values given the parent configuration values
+     * @param alternativeString  the string which will be used to override the string of defaultValue
+     * @return This ConfigDef so you can chain calls
+     */
+    public ConfigDef define(String name, Type type, Object defaultValue, Validator validator, Importance importance, String documentation,
+                            String group, int orderInGroup, Width width, String displayName, List<String> dependents, Recommender recommender,
+                            String alternativeString) {
+        return define(new ConfigKey(name, type, defaultValue, validator, importance, documentation, group, orderInGroup, width, displayName, dependents, recommender, false, alternativeString));
     }
 
     /**
@@ -389,6 +412,21 @@ public ConfigDef define(String name, Type type, Object defaultValue, Importance
         return define(name, type, defaultValue, null, importance, documentation);
     }
 
+    /**
+     * Define a new configuration with no special validation logic
+     * @param name              The name of the config parameter
+     * @param type              The type of the config
+     * @param defaultValue      The default value to use if this config isn't present
+     * @param importance        The importance of this config: is this something you will likely need to change.
+     * @param documentation     The documentation string for the config
+     * @param alternativeString The string which will be used to override the string of defaultValue
+     * @return This ConfigDef so you can chain calls
+     */
+    public ConfigDef define(String name, Type type, Object defaultValue, Importance importance, String documentation, String alternativeString) {
+        return define(name, type, defaultValue, null, importance, documentation, null, -1, Width.NONE,
+                name, Collections.emptyList(), null, alternativeString);
+    }
+
     /**
      * Define a new configuration with no default value and no special validation logic
      * @param name          The name of the config parameter
@@ -411,7 +449,7 @@ public ConfigDef define(String name, Type type, Importance importance, String do
      * @return This ConfigDef so you can chain calls
      */
     public ConfigDef defineInternal(final String name, final Type type, final Object defaultValue, final Importance importance) {
-        return define(new ConfigKey(name, type, defaultValue, null, importance, "", "", -1, Width.NONE, name, Collections.emptyList(), null, true));
+        return define(new ConfigKey(name, type, defaultValue, null, importance, "", "", -1, Width.NONE, name, Collections.emptyList(), null, true, null));
     }
 
     /**
@@ -426,7 +464,7 @@ public ConfigDef defineInternal(final String name, final Type type, final Object
      * @return This ConfigDef so you can chain calls
      */
     public ConfigDef defineInternal(final String name, final Type type, final Object defaultValue, final Validator validator, final Importance importance, final String documentation) {
-        return define(new ConfigKey(name, type, defaultValue, validator, importance, documentation, "", -1, Width.NONE, name, Collections.emptyList(), null, true));
+        return define(new ConfigKey(name, type, defaultValue, validator, importance, documentation, "", -1, Width.NONE, name, Collections.emptyList(), null, true, null));
     }
 
     /**
@@ -1216,12 +1254,13 @@ public static class ConfigKey {
         public final List<String> dependents;
         public final Recommender recommender;
         public final boolean internalConfig;
+        public final String alternativeString;
 
         public ConfigKey(String name, Type type, Object defaultValue, Validator validator,
                          Importance importance, String documentation, String group,
                          int orderInGroup, Width width, String displayName,
                          List<String> dependents, Recommender recommender,
-                         boolean internalConfig) {
+                         boolean internalConfig, String alternativeString) {
             this.name = name;
             this.type = type;
             boolean hasDefault = !NO_DEFAULT_VALUE.equals(defaultValue);
@@ -1238,6 +1277,7 @@ public ConfigKey(String name, Type type, Object defaultValue, Validator validato
             this.displayName = displayName;
             this.recommender = recommender;
             this.internalConfig = internalConfig;
+            this.alternativeString = alternativeString;
         }
 
         public boolean hasDefault() {
@@ -1530,7 +1570,8 @@ public void embed(final String keyPrefix, final String groupPrefix, final int st
                     key.displayName,
                     embeddedDependents(keyPrefix, key.dependents),
                     embeddedRecommender(keyPrefix, key.recommender),
-                    key.internalConfig));
+                    key.internalConfig,
+                    key.alternativeString));
         }
     }
 
@@ -1652,6 +1693,10 @@ public String toHtml(int headerDepth, Function<String, String> idGenerator,
                     "<tbody>\n");
             for (String detail : headers()) {
                 if (detail.equals("Name") || detail.equals("Description")) continue;
+                if (detail.equals("Default") && key.alternativeString != null) {
+                    addConfigDetail(b, detail, key.alternativeString);
+                    continue;
+                }
                 addConfigDetail(b, detail, getConfigValue(key, detail));
             }
             if (hasUpdateModes) {
diff --git a/clients/src/test/java/org/apache/kafka/common/config/ConfigDefTest.java b/clients/src/test/java/org/apache/kafka/common/config/ConfigDefTest.java
index d2c572f06f4af..890cbb52a66b1 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/ConfigDefTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/ConfigDefTest.java
@@ -95,7 +95,7 @@ public void testInvalidDefault() {
 
     @Test
     public void testNullDefault() {
-        ConfigDef def = new ConfigDef().define("a", Type.INT, null, null, null, "docs");
+        ConfigDef def = new ConfigDef().define("a", Type.INT, null, null, "docs");
         Map<String, Object> vals = def.parse(new Properties());
 
         assertNull(vals.get("a"));
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/AbstractHerderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/AbstractHerderTest.java
index 7c39ce536be68..a11b0232d8992 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/AbstractHerderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/AbstractHerderTest.java
@@ -1119,7 +1119,7 @@ public void testConnectorOffsets() throws Exception {
     protected void addConfigKey(Map<String, ConfigDef.ConfigKey> keys, String name, String group) {
         keys.put(name, new ConfigDef.ConfigKey(name, ConfigDef.Type.STRING, null, null,
                 ConfigDef.Importance.HIGH, "doc", group, 10,
-                ConfigDef.Width.MEDIUM, "display name", Collections.emptyList(), null, false));
+                ConfigDef.Width.MEDIUM, "display name", Collections.emptyList(), null, false, null));
     }
 
     protected void addValue(List<ConfigValue> values, String name, String value, String...errors) {
diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 2c8a1b59ff99d..70520c1eee26e 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -749,7 +749,7 @@ public class StreamsConfig extends AbstractConfig {
     /** {@code state.dir} */
     @SuppressWarnings("WeakerAccess")
     public static final String STATE_DIR_CONFIG = "state.dir";
-    private static final String STATE_DIR_DOC = "Directory location for state store. This path must be unique for each streams instance sharing the same underlying filesystem.";
+    private static final String STATE_DIR_DOC = "Directory location for state store. This path must be unique for each streams instance sharing the same underlying filesystem. Note that if not configured, then the default location will be different in each environment as it is computed using System.getProperty(\"java.io.tmpdir\")";
 
     /** {@code task.timeout.ms} */
     public static final String TASK_TIMEOUT_MS_CONFIG = "task.timeout.ms";
@@ -852,7 +852,8 @@ public class StreamsConfig extends AbstractConfig {
                     Type.STRING,
                     System.getProperty("java.io.tmpdir") + File.separator + "kafka-streams",
                     Importance.HIGH,
-                    STATE_DIR_DOC)
+                    STATE_DIR_DOC,
+                    "${java.io.tmpdir}")
 
             // MEDIUM
 

From 6a4078ddde7888ca12582b6dff13f9ece0cf4d59 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Mon, 12 Feb 2024 02:16:57 -0500
Subject: [PATCH 022/258] KAFKA-16215; KAFKA-16178; Fix member not rejoining
 after error (#15311)

This fixes a bug that was causing that members wouldn't rejoin the group after receiving an error in the heartbeat response (ex. fenced, not coordinator, as reported in KAFKA-16215 and KAFKA-16178). The issue was that when receiving a response with error, the response receive time was not being updated, so following heartbeat would be skipped, considering that there was already a previous one inflight.

Reviewers: Andrew Schofield <aschofield@confluent.io>, David Jacot <djacot@confluent.io>
---
 .../internals/HeartbeatRequestManager.java    | 17 ++++++++++----
 .../HeartbeatRequestManagerTest.java          | 23 ++++++++++++++++---
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 03e11ddfa02b4..246ea05b220d8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -267,11 +267,12 @@ private NetworkClientDelegate.UnsentRequest makeHeartbeatRequest(final boolean i
             return logResponse(request);
         else
             return request.whenComplete((response, exception) -> {
+                long completionTimeMs = request.handler().completionTimeMs();
                 if (response != null) {
                     metricsManager.recordRequestLatency(response.requestLatencyMs());
-                    onResponse((ConsumerGroupHeartbeatResponse) response.responseBody(), request.handler().completionTimeMs());
+                    onResponse((ConsumerGroupHeartbeatResponse) response.responseBody(), completionTimeMs);
                 } else {
-                    onFailure(exception, request.handler().completionTimeMs());
+                    onFailure(exception, completionTimeMs);
                 }
             });
     }
@@ -341,9 +342,8 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
         String message;
 
         this.heartbeatState.reset();
+        this.heartbeatRequestState.onFailedAttempt(currentTimeMs);
 
-        // TODO: upon encountering a fatal/fenced error, trigger onPartitionLost logic to give up the current
-        //  assignments.
         switch (error) {
             case NOT_COORDINATOR:
                 // the manager should retry immediately when the coordinator node becomes available again
@@ -352,6 +352,8 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                         coordinatorRequestManager.coordinator());
                 logInfo(message, response, currentTimeMs);
                 coordinatorRequestManager.markCoordinatorUnknown(errorMessage, currentTimeMs);
+                // Skip backoff so that the next HB is sent as soon as the new coordinator is discovered
+                heartbeatRequestState.reset();
                 break;
 
             case COORDINATOR_NOT_AVAILABLE:
@@ -360,6 +362,8 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                         coordinatorRequestManager.coordinator());
                 logInfo(message, response, currentTimeMs);
                 coordinatorRequestManager.markCoordinatorUnknown(errorMessage, currentTimeMs);
+                // Skip backoff so that the next HB is sent as soon as the new coordinator is discovered
+                heartbeatRequestState.reset();
                 break;
 
             case COORDINATOR_LOAD_IN_PROGRESS:
@@ -368,7 +372,6 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                                 "Will retry",
                         coordinatorRequestManager.coordinator());
                 logInfo(message, response, currentTimeMs);
-                heartbeatRequestState.onFailedAttempt(currentTimeMs);
                 break;
 
             case GROUP_AUTHORIZATION_FAILED:
@@ -397,6 +400,8 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                         membershipManager.memberId(), membershipManager.memberEpoch());
                 logInfo(message, response, currentTimeMs);
                 membershipManager.transitionToFenced();
+                // Skip backoff so that a next HB to rejoin is sent as soon as the fenced member releases its assignment
+                heartbeatRequestState.reset();
                 break;
 
             case UNKNOWN_MEMBER_ID:
@@ -404,6 +409,8 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                         membershipManager.memberId());
                 logInfo(message, response, currentTimeMs);
                 membershipManager.transitionToFenced();
+                // Skip backoff so that a next HB to rejoin is sent as soon as the fenced member releases its assignment
+                heartbeatRequestState.reset();
                 break;
 
             default:
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 0df6fa94c39cb..bc014efa7730b 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -445,16 +445,26 @@ public void testHeartbeatResponseOnErrorHandling(final Errors error, final boole
 
             case COORDINATOR_LOAD_IN_PROGRESS:
                 verify(backgroundEventHandler, never()).add(any());
-                assertEquals(DEFAULT_RETRY_BACKOFF_MS, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()));
+                assertEquals(DEFAULT_RETRY_BACKOFF_MS,
+                    heartbeatRequestState.nextHeartbeatMs(time.milliseconds()), "Request should " +
+                        "backoff after receiving a coordinator load in progress error. ");
                 break;
 
             case COORDINATOR_NOT_AVAILABLE:
             case NOT_COORDINATOR:
                 verify(backgroundEventHandler, never()).add(any());
                 verify(coordinatorRequestManager).markCoordinatorUnknown(any(), anyLong());
-                assertEquals(0, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()));
+                assertEquals(0, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()),
+                    "Request should not apply backoff so that the next heartbeat is sent " +
+                        "as soon as the new coordinator is discovered.");
+                break;
+            case UNKNOWN_MEMBER_ID:
+            case FENCED_MEMBER_EPOCH:
+                verify(backgroundEventHandler, never()).add(any());
+                assertEquals(0, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()),
+                    "Request should not apply backoff so that the next heartbeat to rejoin is " +
+                        "sent as soon as the fenced member releases its assignment.");
                 break;
-
             default:
                 if (isFatal) {
                     // The memberStateManager should have stopped heartbeat at this point
@@ -465,6 +475,13 @@ public void testHeartbeatResponseOnErrorHandling(final Errors error, final boole
                 }
                 break;
         }
+
+        if (!isFatal) {
+            // Make sure a next heartbeat is sent for all non-fatal errors (to retry or rejoin)
+            time.sleep(DEFAULT_HEARTBEAT_INTERVAL_MS);
+            result = heartbeatRequestManager.poll(time.milliseconds());
+            assertEquals(1, result.unsentRequests.size());
+        }
     }
 
     @Test

From 11f37b4d7316498bc863cbf2dff4bb853df241c6 Mon Sep 17 00:00:00 2001
From: Andrew Schofield <aschofield@confluent.io>
Date: Mon, 12 Feb 2024 11:30:14 +0000
Subject: [PATCH 023/258] KAFKA-14041: Avoid the keyword var for a variable
 declaration (#15351)

Reviewers: Divij Vaidya <diviv@amazon.com>, Andras Katona <41361962+akatona84@users.noreply.github.com>
---
 .../apache/kafka/common/config/ConfigTransformer.java  | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformer.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformer.java
index 4f078b15a32f3..dbf6c7bbfcec1 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformer.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformer.java
@@ -81,11 +81,11 @@ public ConfigTransformerResult transform(Map<String, String> configs) {
         // Collect the variables from the given configs that need transformation
         for (Map.Entry<String, String> config : configs.entrySet()) {
             if (config.getValue() != null) {
-                List<ConfigVariable> vars = getVars(config.getValue(), DEFAULT_PATTERN);
-                for (ConfigVariable var : vars) {
-                    Map<String, Set<String>> keysByPath = keysByProvider.computeIfAbsent(var.providerName, k -> new HashMap<>());
-                    Set<String> keys = keysByPath.computeIfAbsent(var.path, k -> new HashSet<>());
-                    keys.add(var.variable);
+                List<ConfigVariable> configVars = getVars(config.getValue(), DEFAULT_PATTERN);
+                for (ConfigVariable configVar : configVars) {
+                    Map<String, Set<String>> keysByPath = keysByProvider.computeIfAbsent(configVar.providerName, k -> new HashMap<>());
+                    Set<String> keys = keysByPath.computeIfAbsent(configVar.path, k -> new HashSet<>());
+                    keys.add(configVar.variable);
                 }
             }
         }

From b9f1d59268bd449f05ada2b2a14ebdd1226fc4db Mon Sep 17 00:00:00 2001
From: Matthias Berndt <matthias_berndt@gmx.de>
Date: Mon, 12 Feb 2024 12:34:57 +0100
Subject: [PATCH 024/258] MINOR: fix syntax error in release.py (#15350)

Co-authored-by: Matthias Berndt <matthias.berndt@ttmzero.com>

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 release.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/release.py b/release.py
index 6b5b538656173..e00773c3da7e5 100755
--- a/release.py
+++ b/release.py
@@ -749,7 +749,7 @@ def select_gpg_key():
         print(f"In order to restart the workflow, you will have to manually switch back to the original branch and delete the branch {release_version} and tag {rc_tag}")
         sys.exit(1)
 except Exception as e:
-    print(f"Failed when trying to git push {rc_tag}. Error: {e}"
+    print(f"Failed when trying to git push {rc_tag}. Error: {e}")
     print("You may need to clean up branches/tags yourself before retrying.")
     print("Due the failure of git push, the program will exit here. Please note that: ")
     print(f"1) You are still at branch {release_version}, not {starting_branch}")

From fc8b644e5603d174e50e724789dad095cf856411 Mon Sep 17 00:00:00 2001
From: Satish Duggana <satishd@apache.org>
Date: Mon, 12 Feb 2024 17:35:01 +0530
Subject: [PATCH 025/258] MINOR Removed unused CommittedOffsetsFile class.
 (#15209)

`CommittedOffsetsFile` can be introduced when it is required for enhancing TBRLMM to consume from a specific offset when snapshots are implemented.

Reviewers: Kamal Chandraprakash<kamal.chandraprakash@gmail.com>, Divij Vaidya <diviv@amazon.com>
---
 .../storage/CommittedOffsetsFile.java         | 86 -------------------
 1 file changed, 86 deletions(-)
 delete mode 100644 storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/CommittedOffsetsFile.java

diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/CommittedOffsetsFile.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/CommittedOffsetsFile.java
deleted file mode 100644
index a08e0f30507bf..0000000000000
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/CommittedOffsetsFile.java
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.kafka.server.log.remote.metadata.storage;
-
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.server.common.CheckpointFile;
-
-import java.io.File;
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.regex.Pattern;
-
-/**
- * This class represents a file containing the committed offsets of remote log metadata partitions.
- */
-public class CommittedOffsetsFile {
-    private static final int CURRENT_VERSION = 0;
-    private static final String SEPARATOR = " ";
-
-    private static final Pattern MINIMUM_ONE_WHITESPACE = Pattern.compile("\\s+");
-    private final CheckpointFile<Map.Entry<Integer, Long>> checkpointFile;
-
-    CommittedOffsetsFile(File offsetsFile) throws IOException {
-        CheckpointFile.EntryFormatter<Map.Entry<Integer, Long>> formatter = new EntryFormatter();
-        checkpointFile = new CheckpointFile<>(offsetsFile, CURRENT_VERSION, formatter);
-    }
-
-    private static class EntryFormatter implements CheckpointFile.EntryFormatter<Map.Entry<Integer, Long>> {
-
-        @Override
-        public String toString(Map.Entry<Integer, Long> entry) {
-            // Each entry is stored in a new line as <partition-num offset>
-            return entry.getKey() + SEPARATOR + entry.getValue();
-        }
-
-        @Override
-        public Optional<Map.Entry<Integer, Long>> fromString(String line) {
-            String[] strings = MINIMUM_ONE_WHITESPACE.split(line);
-            if (strings.length != 2) {
-                return Optional.empty();
-            }
-
-            try {
-                return Optional.of(Utils.mkEntry(Integer.parseInt(strings[0]), Long.parseLong(strings[1])));
-            } catch (NumberFormatException e) {
-                return Optional.empty();
-            }
-
-        }
-    }
-
-    public synchronized void writeEntries(Map<Integer, Long> committedOffsets) throws IOException {
-        checkpointFile.write(committedOffsets.entrySet(), true);
-    }
-
-    public synchronized Map<Integer, Long> readEntries() throws IOException {
-        List<Map.Entry<Integer, Long>> entries = checkpointFile.read();
-        Map<Integer, Long> partitionToOffsets = new HashMap<>(entries.size());
-        for (Map.Entry<Integer, Long> entry : entries) {
-            Long existingValue = partitionToOffsets.put(entry.getKey(), entry.getValue());
-            if (existingValue != null) {
-                throw new IOException("Multiple entries exist for key: " + entry.getKey());
-            }
-        }
-
-        return partitionToOffsets;
-    }
-}

From 5cfcc52fb3fce4a43ca77df311382d7a02a40ed2 Mon Sep 17 00:00:00 2001
From: ghostspiders <30502753+ghostspiders@users.noreply.github.com>
Date: Mon, 12 Feb 2024 20:27:47 +0800
Subject: [PATCH 026/258] KAFKA-16239: Clean up references to non-existent
 IntegrationTestHelper (#15352)

Co-authored-by: ghostspiders <yufeng.gao@seres.cn>

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 checkstyle/import-control-core.xml              |  1 -
 core/src/test/java/kafka/test/junit/README.md   | 17 -----------------
 .../junit/RaftClusterInvocationContext.java     |  1 -
 .../test/junit/ZkClusterInvocationContext.java  |  1 -
 4 files changed, 20 deletions(-)

diff --git a/checkstyle/import-control-core.xml b/checkstyle/import-control-core.xml
index 4430b8ec9ddb2..3f9a21fffc6c8 100644
--- a/checkstyle/import-control-core.xml
+++ b/checkstyle/import-control-core.xml
@@ -105,7 +105,6 @@
     <allow pkg="kafka.server"/>
     <allow pkg="kafka.zk" />
     <allow pkg="org.apache.kafka.clients.admin"/>
-    <allow pkg="integration.kafka.server" class="IntegrationTestHelper"/>
     <subpackage name="annotation">
       <allow pkg="kafka.test"/>
     </subpackage>
diff --git a/core/src/test/java/kafka/test/junit/README.md b/core/src/test/java/kafka/test/junit/README.md
index 6df7a26a9c107..6973ca78ed3c5 100644
--- a/core/src/test/java/kafka/test/junit/README.md
+++ b/core/src/test/java/kafka/test/junit/README.md
@@ -107,7 +107,6 @@ previously garnered from the test hierarchy.
 
 * ClusterConfig: a mutable cluster configuration, includes cluster type, number of brokers, properties, etc
 * ClusterInstance: a shim to the underlying class that actually runs the cluster, provides access to things like SocketServers
-* IntegrationTestHelper: connection related functions taken from IntegrationTestHarness and BaseRequestTest
 
 In order to have one of these objects injected, simply add it as a parameter to your test class, `@BeforeEach` method, or test method.
 
@@ -115,23 +114,7 @@ In order to have one of these objects injected, simply add it as a parameter to
 | --- | --- | --- | --- | --- |
 | ClusterConfig | yes | yes | yes* | Once in the test, changing config has no effect |
 | ClusterInstance | yes* | no | yes | Injectable at class level for convenience, can only be accessed inside test |
-| IntegrationTestHelper | yes | yes | yes | - |
 
-```scala
-@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
-class SomeTestClass(helper: IntegrationTestHelper) {
- 
-  @BeforeEach
-  def setup(config: ClusterConfig): Unit = {
-    config.serverProperties().put("foo", "bar")
-  }
-
-  @ClusterTest
-  def testSomething(cluster: ClusterInstance): Unit = {
-    val topics = cluster.createAdminClient().listTopics()
-  }
-}
-```
 
 # Gotchas
 * Test methods annotated with JUnit's `@Test` will still be run, but no cluster will be started and no dependency 
diff --git a/core/src/test/java/kafka/test/junit/RaftClusterInvocationContext.java b/core/src/test/java/kafka/test/junit/RaftClusterInvocationContext.java
index 4aa152055af20..aa1cd8d136c6e 100644
--- a/core/src/test/java/kafka/test/junit/RaftClusterInvocationContext.java
+++ b/core/src/test/java/kafka/test/junit/RaftClusterInvocationContext.java
@@ -59,7 +59,6 @@
  * <ul>
  *     <li>ClusterConfig (the same instance passed to the constructor)</li>
  *     <li>ClusterInstance (includes methods to expose underlying SocketServer-s)</li>
- *     <li>IntegrationTestHelper (helper methods)</li>
  * </ul>
  */
 public class RaftClusterInvocationContext implements TestTemplateInvocationContext {
diff --git a/core/src/test/java/kafka/test/junit/ZkClusterInvocationContext.java b/core/src/test/java/kafka/test/junit/ZkClusterInvocationContext.java
index dd0098f7868ca..ccfceade7bf36 100644
--- a/core/src/test/java/kafka/test/junit/ZkClusterInvocationContext.java
+++ b/core/src/test/java/kafka/test/junit/ZkClusterInvocationContext.java
@@ -61,7 +61,6 @@
  * <ul>
  *     <li>ClusterConfig (the same instance passed to the constructor)</li>
  *     <li>ClusterInstance (includes methods to expose underlying SocketServer-s)</li>
- *     <li>IntegrationTestHelper (helper methods)</li>
  * </ul>
  */
 public class ZkClusterInvocationContext implements TestTemplateInvocationContext {

From c6f4c604d8e50ad9e182eeb66f0d1650aa44f277 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Mon, 12 Feb 2024 15:43:21 +0100
Subject: [PATCH 027/258] KAFKA-15832: Trigger client reconciliation based on
 manager poll (#15275)

Currently, the reconciliation logic on the client is triggered when a new target assignment is received and resolved, or when new unresolved target assignments are discovered in metadata.

This change improves this by triggering the reconciliation logic on each poll iteration, to reconcile whatever is ready to be reconciled. This would require changes to support poll on the MembershipManager, and integrate it with the current polling logic in the background thread. Receiving a new target assignment from the broker, or resolving new topic names via a metadata update could only ensure that the #assignmentReadyToReconcile is properly updated (currently done), but wouldn't trigger the #reconcile() logic, leaving that to the #poll() operation.

Reviewers: David Jacot <djacot@confluent.io>, Lianet Magrans <lianetmr@gmail.com>
---
 .../consumer/internals/MembershipManager.java |   2 +-
 .../internals/MembershipManagerImpl.java      |  76 +++----
 .../consumer/internals/RequestManagers.java   |   9 +-
 .../internals/ConsumerTestBuilder.java        |   4 +-
 .../HeartbeatRequestManagerTest.java          |   4 +-
 .../internals/MembershipManagerImplTest.java  | 200 ++++++++++++++++--
 .../events/ApplicationEventProcessorTest.java |   4 +-
 7 files changed, 225 insertions(+), 74 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
index f2e1291a52e58..7f0975b5aa968 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
@@ -34,7 +34,7 @@
  * <li>Keeping assignment for the member</li>
  * <li>Computing assignment for the group if the member is required to do so<li/>
  */
-public interface MembershipManager {
+public interface MembershipManager extends RequestManager {
 
     /**
      * @return Group ID of the consumer group the member is part of (or wants to be part of).
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 50f9a920a561d..a5a28da049ec4 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -19,6 +19,7 @@
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
 import org.apache.kafka.clients.consumer.internals.Utils.TopicIdPartitionComparator;
 import org.apache.kafka.clients.consumer.internals.Utils.TopicPartitionComparator;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
@@ -26,8 +27,6 @@
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
 import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
-import org.apache.kafka.common.ClusterResource;
-import org.apache.kafka.common.ClusterResourceListener;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.TopicPartition;
@@ -112,7 +111,7 @@
  * the user if the callbacks fail. This manager is only concerned about the callbacks completion to
  * know that it can proceed with the reconciliation.
  */
-public class MembershipManagerImpl implements MembershipManager, ClusterResourceListener {
+public class MembershipManagerImpl implements MembershipManager {
 
     /**
      * TopicPartition comparator based on topic name and partition id.
@@ -215,9 +214,9 @@ public class MembershipManagerImpl implements MembershipManager, ClusterResource
      */
     private final Map<Uuid, SortedSet<Integer>> currentTargetAssignment;
 
-    /**
+    /**     
      * If there is a reconciliation running (triggering commit, callbacks) for the
-     * assignmentReadyToReconcile. This will be true if {@link #reconcile()} has been triggered
+     * assignmentReadyToReconcile. This will be true if {@link #maybeReconcile()} has been triggered
      * after receiving a heartbeat response, or a metadata update.
      */
     private boolean reconciliationInProgress;
@@ -237,14 +236,6 @@ public class MembershipManagerImpl implements MembershipManager, ClusterResource
      */
     private Optional<CompletableFuture<Void>> leaveGroupInProgress = Optional.empty();
 
-    /**
-     * True if the member has registered to be notified when the cluster metadata is updated.
-     * This is initially false, as the member that is not part of a consumer group does not
-     * require metadata updated. This becomes true the first time the member joins on the
-     * {@link #transitionToJoining()}
-     */
-    private boolean isRegisteredForMetadataUpdates;
-
     /**
      * Registered listeners that will be notified whenever the memberID/epoch gets updated (valid
      * values received from the broker, or values cleared due to member leaving the group, getting
@@ -393,9 +384,9 @@ public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData respo
 
     /**
      * This will process the assignment received if it is different from the member's current
-     * assignment. If a new assignment is received, this will try to resolve the topic names from
-     * metadata, reconcile the resolved assignment, and keep the unresolved to be reconciled when
-     * metadata is discovered.
+     * assignment. If a new assignment is received, this will make sure reconciliation is attempted
+     * on the next call of `poll`. If another reconciliation is currently in process, the first `poll`
+     * after that reconciliation will trigger the new reconciliation.
      *
      * @param assignment Assignment received from the broker.
      */
@@ -406,7 +397,6 @@ private void processAssignmentReceived(ConsumerGroupHeartbeatResponseData.Assign
             // assignment from the broker, different from the current assignment. Note that the
             // reconciliation might not be triggered just yet because of missing metadata.
             transitionTo(MemberState.RECONCILING);
-            reconcile();
         } else {
             // Same assignment received, nothing to reconcile.
             log.debug("Target assignment {} received from the broker is equals to the member " +
@@ -567,18 +557,6 @@ public void transitionToJoining() {
         resetEpoch();
         transitionTo(MemberState.JOINING);
         clearPendingAssignmentsAndLocalNamesCache();
-        registerForMetadataUpdates();
-    }
-
-    /**
-     * Register to get notified when the cluster metadata is updated, via the
-     * {@link #onUpdate(ClusterResource)}. Register only if the manager is not register already.
-     */
-    private void registerForMetadataUpdates() {
-        if (!isRegisteredForMetadataUpdates) {
-            this.metadata.addClusterUpdateListener(this);
-            isRegisteredForMetadataUpdates = true;
-        }
     }
 
     /**
@@ -770,12 +748,15 @@ public void transitionToStale() {
      * Reconcile the assignment that has been received from the server. If for some topics, the
      * topic ID cannot be matched to a topic name, a metadata update will be triggered and only
      * the subset of topics that are resolvable will be reconciled. Reconciliation will trigger the
-     * callbacks and update the subscription state. Note that only one reconciliation
-     * can be in progress at a time. If there is already another one in progress when this is
-     * triggered, it will be no-op, and the assignment will be reconciled on the next
-     * reconciliation loop.
+     * callbacks and update the subscription state.
+     *
+     * There are three conditions under which no reconciliation will be triggered:
+     *  - We have already reconciled the assignment (the target assignment is the same as the current assignment).
+     *  - Another reconciliation is already in progress.
+     *  - There are topics that haven't been added to the current assignment yet, but all their topic IDs
+     *    are missing from the target assignment.
      */
-    void reconcile() {
+    void maybeReconcile() {
         if (targetAssignmentReconciled()) {
             log.debug("Ignoring reconciliation attempt. Target assignment is equal to the " +
                     "current assignment.");
@@ -1356,29 +1337,12 @@ Map<Uuid, SortedSet<Integer>> topicPartitionsAwaitingReconciliation() {
 
     /**
      * @return If there is a reconciliation in process now. Note that reconciliation is triggered
-     * by a call to {@link #reconcile()}. Visible for testing.
+     * by a call to {@link #maybeReconcile()}. Visible for testing.
      */
     boolean reconciliationInProgress() {
         return reconciliationInProgress;
     }
 
-    /**
-     * When cluster metadata is updated, try to resolve topic names for topic IDs received in
-     * assignment that hasn't been resolved yet.
-     * <ul>
-     *     <li>Try to find topic names for all assignments</li>
-     *     <li>Add discovered topic names to the local topic names cache</li>
-     *     <li>If any topics are resolved, trigger a reconciliation process</li>
-     *     <li>If some topics still remain unresolved, request another metadata update</li>
-     * </ul>
-     */
-    @Override
-    public void onUpdate(ClusterResource clusterResource) {
-        if (state == MemberState.RECONCILING) {
-            reconcile();
-        }
-    }
-
     /**
      * Register a new listener that will be invoked whenever the member state changes, or a new
      * member ID or epoch is received.
@@ -1392,4 +1356,12 @@ public void registerStateListener(MemberStateListener listener) {
         }
         this.stateUpdatesListeners.add(listener);
     }
+
+    @Override
+    public PollResult poll(final long currentTimeMs) {
+        if (state == MemberState.RECONCILING) {
+            maybeReconcile();
+        }
+        return PollResult.EMPTY;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
index b3c0c5b63b58e..2d90a3ad7082e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
@@ -48,6 +48,7 @@ public class RequestManagers implements Closeable {
     public final Optional<CoordinatorRequestManager> coordinatorRequestManager;
     public final Optional<CommitRequestManager> commitRequestManager;
     public final Optional<HeartbeatRequestManager> heartbeatRequestManager;
+    public final Optional<MembershipManager> membershipManager;
     public final OffsetsRequestManager offsetsRequestManager;
     public final TopicMetadataRequestManager topicMetadataRequestManager;
     public final FetchRequestManager fetchRequestManager;
@@ -60,7 +61,8 @@ public RequestManagers(LogContext logContext,
                            FetchRequestManager fetchRequestManager,
                            Optional<CoordinatorRequestManager> coordinatorRequestManager,
                            Optional<CommitRequestManager> commitRequestManager,
-                           Optional<HeartbeatRequestManager> heartbeatRequestManager) {
+                           Optional<HeartbeatRequestManager> heartbeatRequestManager,
+                           Optional<MembershipManager> membershipManager) {
         this.log = logContext.logger(RequestManagers.class);
         this.offsetsRequestManager = requireNonNull(offsetsRequestManager, "OffsetsRequestManager cannot be null");
         this.coordinatorRequestManager = coordinatorRequestManager;
@@ -68,11 +70,13 @@ public RequestManagers(LogContext logContext,
         this.topicMetadataRequestManager = topicMetadataRequestManager;
         this.fetchRequestManager = fetchRequestManager;
         this.heartbeatRequestManager = heartbeatRequestManager;
+        this.membershipManager = membershipManager;
 
         List<Optional<? extends RequestManager>> list = new ArrayList<>();
         list.add(coordinatorRequestManager);
         list.add(commitRequestManager);
         list.add(heartbeatRequestManager);
+        list.add(membershipManager);
         list.add(Optional.of(offsetsRequestManager));
         list.add(Optional.of(topicMetadataRequestManager));
         list.add(Optional.of(fetchRequestManager));
@@ -205,7 +209,8 @@ protected RequestManagers create() {
                         fetch,
                         Optional.ofNullable(coordinator),
                         Optional.ofNullable(commit),
-                        Optional.ofNullable(heartbeatRequestManager)
+                        Optional.ofNullable(heartbeatRequestManager),
+                        Optional.ofNullable(membershipManager)
                 );
             }
         };
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
index f6f71f3f7a95b..d6ae62950608c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
@@ -263,7 +263,9 @@ public ConsumerTestBuilder(Optional<GroupInformation> groupInfo, boolean enableA
                 fetchRequestManager,
                 coordinatorRequestManager,
                 commitRequestManager,
-                heartbeatRequestManager);
+                heartbeatRequestManager,
+                membershipManager
+            );
         this.applicationEventProcessor = spy(new ApplicationEventProcessor(
                 logContext,
                 applicationEventQueue,
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index bc014efa7730b..4040d6823314d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -546,7 +546,9 @@ public void testHeartbeatState() {
                 .setAssignment(assignmentTopic1));
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, "topic1"));
         membershipManager.onHeartbeatResponseReceived(rs1.data());
-        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+
+        // We remain in RECONCILING state, as the assignment will be reconciled on the next poll
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 234b374d33289..bb441fe7e649a 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -57,6 +57,7 @@
 import static org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
+import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.apache.kafka.common.utils.Utils.mkSortedSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -165,7 +166,6 @@ public void testMembershipManagerRegistersForClusterMetadataUpdatesOnFirstJoin()
                 subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
                 backgroundEventHandler, time);
         manager.transitionToJoining();
-        verify(metadata).addClusterUpdateListener(manager);
         clearInvocations(metadata);
 
         // Following joins should not register again.
@@ -176,7 +176,6 @@ public void testMembershipManagerRegistersForClusterMetadataUpdatesOnFirstJoin()
         manager.onHeartbeatRequestSent();
         assertEquals(MemberState.UNSUBSCRIBED, manager.state());
         manager.transitionToJoining();
-        verify(metadata, never()).addClusterUpdateListener(manager);
     }
 
     @Test
@@ -507,12 +506,10 @@ public void testDelayedReconciliationResultDiscardedIfMemberRejoins() {
     }
 
     /**
-     * This is the case where a member is stuck reconciling an assignment A (waiting on
-     * metadata, commit or callbacks), and the target assignment changes (due to new topics added
-     * to metadata, or new assignment received from broker). If the reconciliation of A completes
-     * it should be applied (should update the assignment on the member and send ack), and then
-     * the reconciliation of assignment B will be processed and applied in the next
-     * reconciliation loop.
+     * This is the case where a member is stuck reconciling an assignment A (waiting on metadata, commit or callbacks), and the target
+     * assignment changes due to new topics. If the reconciliation of A completes it should be applied (should update the assignment
+     * on the member and send ack), and then the reconciliation of assignment B will be processed and applied in the next reconciliation
+     * loop.
      */
     @Test
     public void testDelayedReconciliationResultAppliedWhenTargetChangedWithMetadataUpdate() {
@@ -539,8 +536,8 @@ public void testDelayedReconciliationResultAppliedWhenTargetChangedWithMetadataU
 
         // Metadata discovered for topic2 while reconciliation in progress to revoke topic1.
         // Should not trigger a new reconciliation because there is one already in progress.
-        mockTopicNameInMetadataCache(Collections.singletonMap(topicId2, topic2), true);
-        membershipManager.onUpdate(null);
+        final Map<Uuid, String> topic2Metadata = Collections.singletonMap(topicId2, topic2);
+        mockTopicNameInMetadataCache(topic2Metadata, true);
         assertEquals(Collections.singleton(topicId2), membershipManager.topicsAwaitingReconciliation());
         verifyReconciliationNotTriggered(membershipManager);
 
@@ -557,6 +554,86 @@ public void testDelayedReconciliationResultAppliedWhenTargetChangedWithMetadataU
         // next reconciliation loop.
         Map<Uuid, SortedSet<Integer>> topic2Assignment = topicIdPartitionsMap(topicId2,  1, 2);
         assertEquals(topic2Assignment, membershipManager.topicPartitionsAwaitingReconciliation());
+
+        // After acknowledging the assignment, we should be back to RECONCILING, because we have not
+        // yet reached the target assignment.
+        membershipManager.onHeartbeatRequestSent();
+
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        assertEquals(topic2Assignment, membershipManager.topicPartitionsAwaitingReconciliation());
+
+        // Next reconciliation triggered in poll
+        membershipManager.poll(time.milliseconds());
+
+        assertEquals(Collections.emptySet(), membershipManager.topicsAwaitingReconciliation());
+        verify(subscriptionState).assignFromSubscribed(topicPartitions(topic2Assignment, topic2Metadata));
+    }
+
+    /**
+     * This is the case where a member is stuck reconciling an assignment A (waiting on metadata, commit or callbacks), and the target
+     * assignment changes due to a new assignment received from broker. If the reconciliation of A completes it should be applied (should
+     * update the assignment on the member and send ack), and then the reconciliation of assignment B will be processed and applied in the
+     * next reconciliation loop.
+     */
+    @Test
+    public void testDelayedReconciliationResultAppliedWhenTargetChangedWithNewAssignment() {
+
+        Uuid topicId1 = Uuid.randomUuid();
+        String topic1 = "topic1";
+        final TopicIdPartition topicId1Partition0 = new TopicIdPartition(topicId1, new TopicPartition(topic1, 0));
+
+        Uuid topicId2 = Uuid.randomUuid();
+        String topic2 = "topic2";
+        final TopicIdPartition topicId2Partition0 = new TopicIdPartition(topicId2, new TopicPartition(topic2, 0));
+
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(metadata.topicNames()).thenReturn(
+            mkMap(
+                mkEntry(topicId1, topic1),
+                mkEntry(topicId2, topic2)
+            )
+        );
+
+        // Receive assignment with only topic1-0, getting stuck during commit.
+        final CompletableFuture<Void> commitFuture = mockNewAssignmentAndRevocationStuckOnCommit(membershipManager, topicId1,
+            topic1, Collections.singletonList(0), false);
+
+        // New assignment adding a new topic2-0 (not in metadata).
+        // No reconciliation triggered, because another reconciliation is in progress.
+        Map<Uuid, SortedSet<Integer>> newAssignment =
+            mkMap(
+                mkEntry(topicId1, mkSortedSet(0)),
+                mkEntry(topicId2, mkSortedSet(0))
+            );
+
+        receiveAssignment(newAssignment, membershipManager);
+        membershipManager.poll(time.milliseconds());
+
+        verifyReconciliationNotTriggered(membershipManager);
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        assertEquals(mkSet(topicId1, topicId2), membershipManager.topicsAwaitingReconciliation());
+        clearInvocations(membershipManager, commitRequestManager);
+
+        // First reconciliation completes. Should trigger follow-up reconciliation to complete the assignment,
+        // with membership manager entering ACKNOWLEDGING state.
+
+        commitFuture.complete(null);
+
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        assertEquals(mkSet(topicId2), membershipManager.topicsAwaitingReconciliation());
+
+        // After acknowledging the assignment, we should be back to RECONCILING, because we have not
+        // yet reached the target assignment.
+        membershipManager.onHeartbeatRequestSent();
+
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        clearInvocations(membershipManager, commitRequestManager);
+
+        // Next poll should trigger final reconciliation
+        membershipManager.poll(time.milliseconds());
+
+        verifyReconciliationTriggeredAndCompleted(membershipManager, Arrays.asList(topicId1Partition0, topicId2Partition0));
     }
 
     // Tests the case where topic metadata is not available at the time of the assignment,
@@ -592,6 +669,7 @@ public void testDelayedMetadataUsedToCompleteAssignment() {
             );
 
         receiveAssignment(newAssignment, membershipManager);
+        membershipManager.poll(time.milliseconds());
 
         verifyReconciliationNotTriggered(membershipManager);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
@@ -608,7 +686,7 @@ public void testDelayedMetadataUsedToCompleteAssignment() {
         );
         when(metadata.topicNames()).thenReturn(fullTopicMetadata);
 
-        membershipManager.onUpdate(null);
+        membershipManager.poll(time.milliseconds());
 
         verifyReconciliationTriggeredAndCompleted(membershipManager, Arrays.asList(topicId1Partition0, topicId2Partition0));
     }
@@ -637,6 +715,9 @@ public void testLeaveGroupWhenMemberOwnsAssignment() {
 
         receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         List<TopicIdPartition> assignedPartitions = Arrays.asList(
             new TopicIdPartition(topicId, new TopicPartition(topicName, 0)),
             new TopicIdPartition(topicId, new TopicPartition(topicName, 1)));
@@ -818,6 +899,10 @@ public void testNewAssignmentReplacesPreviousOneWaitingOnMetadata() {
         String topicName = "topic1";
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         receiveAssignment(topicId, Collections.singletonList(0), membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         Set<TopicPartition> expectedAssignment = Collections.singleton(new TopicPartition(topicName, 0));
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         verify(subscriptionState).assignFromSubscribed(expectedAssignment);
@@ -855,6 +940,10 @@ public void testNewEmptyAssignmentReplacesPreviousOneWaitingOnMetadata() {
         String topicName = "topic1";
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         receiveEmptyAssignment(membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.STABLE, membershipManager.state());
         verify(subscriptionState, never()).assignFromSubscribed(any());
     }
@@ -872,6 +961,10 @@ public void testNewAssignmentNotInMetadataReplacesPreviousOneWaitingOnMetadata()
         Uuid topicId = Uuid.randomUuid();
         when(metadata.topicNames()).thenReturn(Collections.emptyMap());
         receiveAssignment(topicId, Collections.singletonList(0), membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertFalse(membershipManager.topicsAwaitingReconciliation().isEmpty());
         assertEquals(topicId, membershipManager.topicsAwaitingReconciliation().iterator().next());
@@ -896,7 +989,8 @@ public void testUnresolvedTargetAssignmentIsReconciledWhenMetadataReceived() {
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty());
-        membershipManager.onUpdate(null);
+
+        membershipManager.poll(time.milliseconds());
 
         // Assignment should have been reconciled.
         Set<TopicPartition> expectedAssignment = Collections.singleton(new TopicPartition(topicName, 1));
@@ -957,6 +1051,9 @@ public void testReconcileNewPartitionsAssignedWhenNoPartitionOwned() {
 
         receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         List<TopicIdPartition> assignedPartitions = topicIdPartitions(topicId, topicName, 0, 1);
         verifyReconciliationTriggeredAndCompleted(membershipManager, assignedPartitions);
     }
@@ -973,6 +1070,9 @@ public void testReconcileNewPartitionsAssignedWhenOtherPartitionsOwned() {
         // New assignment received, adding partitions 1 and 2 to the previously owned partition 0.
         receiveAssignment(topicId, Arrays.asList(0, 1, 2), membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         List<TopicIdPartition> assignedPartitions = new ArrayList<>();
         assignedPartitions.add(ownedPartition);
         assignedPartitions.addAll(topicIdPartitions(topicId, topicName, 1, 2));
@@ -990,6 +1090,10 @@ public void testReconciliationSkippedWhenSameAssignmentReceived() {
         mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId, topicName, Collections.emptyList());
         List<TopicIdPartition> expectedAssignmentReconciled = topicIdPartitions(topicId, topicName, 0, 1);
         receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         verifyReconciliationTriggeredAndCompleted(membershipManager, expectedAssignmentReconciled);
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         clearInvocations(subscriptionState, membershipManager);
@@ -1017,6 +1121,9 @@ public void testReconcilePartitionsRevokedNoAutoCommitNoCallbacks() {
 
         receiveEmptyAssignment(membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         testRevocationOfAllPartitionsCompleted(membershipManager);
     }
 
@@ -1029,6 +1136,9 @@ public void testReconcilePartitionsRevokedWithSuccessfulAutoCommitNoCallbacks()
 
         receiveEmptyAssignment(membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         // Member stays in RECONCILING while the commit request hasn't completed.
         assertEquals(MemberState.RECONCILING, membershipManager.state());
 
@@ -1050,6 +1160,9 @@ public void testReconcilePartitionsRevokedWithFailedAutoCommitCompletesRevocatio
 
         receiveEmptyAssignment(membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         // Member stays in RECONCILING while the commit request hasn't completed.
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         // Partitions should be still owned by the member
@@ -1077,6 +1190,9 @@ public void testReconcileNewPartitionsAssignedAndRevoked() {
         // New assignment received, revoking partition 0, and assigning new partitions 1 and 2.
         receiveAssignment(topicId, Arrays.asList(1, 2), membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         assertEquals(topicIdPartitionsMap(topicId, 1, 2), membershipManager.currentAssignment());
         assertFalse(membershipManager.reconciliationInProgress());
@@ -1105,8 +1221,8 @@ public void testMetadataUpdatesReconcilesUnresolvedAssignments() {
         String topicName = "topic1";
         mockTopicNameInMetadataCache(Collections.singletonMap(topicId, topicName), true);
 
-        // When metadata is updated, the member should re-trigger reconciliation
-        membershipManager.onUpdate(null);
+        // When the next poll is run, the member should re-trigger reconciliation
+        membershipManager.poll(time.milliseconds());
         List<TopicIdPartition> expectedAssignmentReconciled = topicIdPartitions(topicId, topicName, 0, 1);
         verifyReconciliationTriggeredAndCompleted(membershipManager, expectedAssignmentReconciled);
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
@@ -1131,10 +1247,10 @@ public void testMetadataUpdatesRequestsAnotherUpdateIfNeeded() {
         assertEquals(Collections.singleton(topicId), membershipManager.topicsAwaitingReconciliation());
         verify(metadata).requestUpdate(anyBoolean());
 
-        // Metadata update received, but still without the unresolved topic in it. Should keep
+        // Next poll is run, but metadata still without the unresolved topic in it. Should keep
         // the unresolved and request update again.
         when(metadata.topicNames()).thenReturn(Collections.emptyMap());
-        membershipManager.onUpdate(null);
+        membershipManager.poll(time.milliseconds());
         verifyReconciliationNotTriggered(membershipManager);
         assertEquals(Collections.singleton(topicId), membershipManager.topicsAwaitingReconciliation());
         verify(metadata, times(2)).requestUpdate(anyBoolean());
@@ -1152,6 +1268,9 @@ public void testRevokePartitionsUsesTopicNamesLocalCacheWhenMetadataNotAvailable
 
         receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         // Member should complete reconciliation
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         List<Integer> partitions = Arrays.asList(0, 1);
@@ -1173,6 +1292,8 @@ public void testRevokePartitionsUsesTopicNamesLocalCacheWhenMetadataNotAvailable
         // Revoke one of the 2 partitions
         receiveAssignment(topicId, Collections.singletonList(1), membershipManager);
 
+        membershipManager.poll(time.milliseconds());
+
         // Revocation should complete without requesting any metadata update given that the topic
         // received in target assignment should exist in local topic name cache.
         verify(metadata, never()).requestUpdate(anyBoolean());
@@ -1209,6 +1330,10 @@ public void testListenerCallbacksBasic() {
         // Step 2: put the state machine into the appropriate... state
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertTrue(membershipManager.reconciliationInProgress());
         assertEquals(0, listener.revokedCount());
@@ -1240,6 +1365,8 @@ public void testListenerCallbacksBasic() {
         // Step 5: receive an empty assignment, which means we should call revoke
         when(subscriptionState.assignedPartitions()).thenReturn(topicPartitions(topicName, 0, 1));
         receiveEmptyAssignment(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertTrue(membershipManager.reconciliationInProgress());
 
@@ -1293,6 +1420,10 @@ public void testListenerCallbacksThrowsErrorOnPartitionsRevoked() {
 
         // Step 2: put the state machine into the appropriate... state
         receiveEmptyAssignment(membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment());
         assertTrue(membershipManager.reconciliationInProgress());
@@ -1342,6 +1473,10 @@ public void testListenerCallbacksThrowsErrorOnPartitionsAssigned() {
 
         // Step 2: put the state machine into the appropriate... state
         receiveEmptyAssignment(membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment());
         assertTrue(membershipManager.reconciliationInProgress());
@@ -1396,6 +1531,8 @@ public void testAddedPartitionsTemporarilyDisabledAwaitingOnPartitionsAssignedCa
         mockPartitionOwnedAndNewPartitionAdded(topicName, partitionOwned, partitionAdded,
             new CounterConsumerRebalanceListener(), membershipManager);
 
+        membershipManager.poll(time.milliseconds());
+
         verify(subscriptionState).assignFromSubscribedAwaitingCallback(assignedPartitions, addedPartitions);
 
         performCallback(
@@ -1426,6 +1563,8 @@ public void testAddedPartitionsNotEnabledAfterFailedOnPartitionsAssignedCallback
         mockPartitionOwnedAndNewPartitionAdded(topicName, partitionOwned, partitionAdded,
             listener, membershipManager);
 
+        membershipManager.poll(time.milliseconds());
+
         verify(subscriptionState).assignFromSubscribedAwaitingCallback(assignedPartitions, addedPartitions);
 
         performCallback(
@@ -1521,6 +1660,18 @@ private ConsumerRebalanceListenerInvoker consumerRebalanceListenerInvoker() {
         );
     }
 
+    private static SortedSet<TopicPartition> topicPartitions(Map<Uuid, SortedSet<Integer>> topicIdMap, Map<Uuid, String> topicIdNames) {
+        SortedSet<TopicPartition> topicPartitions = new TreeSet<>(new Utils.TopicPartitionComparator());
+
+        for (Uuid topicId : topicIdMap.keySet()) {
+            for (int partition : topicIdMap.get(topicId)) {
+                topicPartitions.add(new TopicPartition(topicIdNames.get(topicId), partition));
+            }
+        }
+
+        return topicPartitions;
+    }
+
     private SortedSet<TopicPartition> topicPartitions(String topicName, int... partitions) {
         SortedSet<TopicPartition> topicPartitions = new TreeSet<>(new Utils.TopicPartitionComparator());
 
@@ -1613,6 +1764,10 @@ public void testMemberJoiningTransitionsToStableWhenReceivingEmptyAssignment() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(null);
         assertEquals(MemberState.JOINING, membershipManager.state());
         receiveEmptyAssignment(membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         assertEquals(MemberState.STABLE, membershipManager.state());
     }
 
@@ -1624,6 +1779,9 @@ private MembershipManagerImpl mockMemberSuccessfullyReceivesAndAcksAssignment(
 
         receiveAssignment(topicId, partitions, membershipManager);
 
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         List<TopicIdPartition> assignedPartitions =
             partitions.stream().map(tp -> new TopicIdPartition(topicId,
                 new TopicPartition(topicName, tp))).collect(Collectors.toList());
@@ -1635,6 +1793,10 @@ private CompletableFuture<Void> mockEmptyAssignmentAndRevocationStuckOnCommit(
             MembershipManagerImpl membershipManager) {
         CompletableFuture<Void> commitResult = mockRevocationNoCallbacks(true);
         receiveEmptyAssignment(membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         verifyReconciliationTriggered(membershipManager);
         clearInvocations(membershipManager);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
@@ -1650,6 +1812,10 @@ private CompletableFuture<Void> mockNewAssignmentAndRevocationStuckOnCommit(
             when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         }
         receiveAssignment(topicId, partitions, membershipManager);
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+
         verifyReconciliationTriggered(membershipManager);
         clearInvocations(membershipManager);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
@@ -1784,6 +1950,8 @@ private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean expectSubscri
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty()).thenReturn(Optional.empty());
 
         membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.poll(time.milliseconds());
+
         if (expectSubscriptionUpdated) {
             verify(subscriptionState).assignFromSubscribed(anyCollection());
         } else {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
index 075552fcd5558..8ea8cb7a729f5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
@@ -77,7 +77,9 @@ public void setup() {
             fetchRequestManager,
             Optional.of(coordinatorRequestManager),
             Optional.of(commitRequestManager),
-            Optional.of(heartbeatRequestManager));
+            Optional.of(heartbeatRequestManager),
+            Optional.of(membershipManager)
+        );
         processor = new ApplicationEventProcessor(
             new LogContext(),
             applicationEventQueue,

From 794c52802c0bcc2676d750613fefe14006d03a44 Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Mon, 12 Feb 2024 14:15:27 -0500
Subject: [PATCH 028/258] KAFKA-14683: Migrate WorkerSinkTaskTest to Mockito
 (2/3) (#15313)

Reviewers: Greg Harris <greg.harris@aiven.io>
---
 .../kafka/connect/runtime/WorkerSinkTask.java |   8 +-
 .../runtime/WorkerSinkTaskMockitoTest.java    | 430 +++++++++++++++-
 .../connect/runtime/WorkerSinkTaskTest.java   | 460 ------------------
 3 files changed, 428 insertions(+), 470 deletions(-)

diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
index 5c2bd3e39cab4..78a61e755f4ba 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
@@ -60,6 +60,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -343,11 +344,16 @@ protected void poll(long timeoutMs) {
         deliverMessages();
     }
 
-    // Visible for testing
+    // VisibleForTesting
     boolean isCommitting() {
         return committing;
     }
 
+    //VisibleForTesting
+    Map<TopicPartition, OffsetAndMetadata> lastCommittedOffsets() {
+        return Collections.unmodifiableMap(lastCommittedOffsets);
+    }
+
     private void doCommitSync(Map<TopicPartition, OffsetAndMetadata> offsets, int seqno) {
         log.debug("{} Committing offsets synchronously using sequence number {}: {}", this, seqno, offsets);
         try {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
index b41f571d6357e..f39dce6646070 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
@@ -20,7 +20,10 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
@@ -33,6 +36,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.WorkerSinkTask.SinkTaskMetricsGroup;
 import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
@@ -49,6 +53,7 @@
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.junit.After;
 import org.junit.Before;
@@ -67,19 +72,25 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
+import java.util.regex.Pattern;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyList;
@@ -300,7 +311,7 @@ public void testPause() {
         verify(statusListener).onResume(taskId);
         verify(consumer, times(2)).wakeup();
         INITIAL_ASSIGNMENT.forEach(tp -> {
-            verify(consumer).resume(Collections.singleton(tp));
+            verify(consumer).resume(singleton(tp));
         });
         verify(sinkTask, times(4)).put(anyList());
     }
@@ -343,6 +354,99 @@ public void testShutdown() throws Exception {
         verify(headerConverter).close();
     }
 
+    @Test
+    public void testPollRedelivery() {
+        createTask(initialState);
+        expectTaskGetTopic();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
+                .thenAnswer(expectConsumerPoll(1))
+                // Retry delivery should succeed
+                .thenAnswer(expectConsumerPoll(0))
+                .thenAnswer(expectConsumerPoll(0));
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        doAnswer(invocation -> null)
+                .doThrow(new RetriableException("retry"))
+                .doAnswer(invocation -> null)
+                .when(sinkTask).put(anyList());
+
+        workerTask.iteration();
+        time.sleep(10000L);
+
+        verifyPollInitialAssignment();
+        verify(sinkTask).put(anyList());
+
+        assertSinkMetricValue("partition-count", 2);
+        assertSinkMetricValue("sink-record-read-total", 0.0);
+        assertSinkMetricValue("sink-record-send-total", 0.0);
+        assertSinkMetricValue("sink-record-active-count", 0.0);
+        assertSinkMetricValue("sink-record-active-count-max", 0.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.0);
+        assertSinkMetricValue("offset-commit-seq-no", 0.0);
+        assertSinkMetricValue("offset-commit-completion-rate", 0.0);
+        assertSinkMetricValue("offset-commit-completion-total", 0.0);
+        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 0.0);
+        assertTaskMetricValue("batch-size-avg", 0.0);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
+
+        // Pause
+        workerTask.iteration();
+
+        verify(consumer, times(3)).assignment();
+        verify(consumer).pause(INITIAL_ASSIGNMENT);
+
+        // Retry delivery should succeed
+        workerTask.iteration();
+        time.sleep(30000L);
+
+        verify(sinkTask, times(3)).put(anyList());
+        INITIAL_ASSIGNMENT.forEach(tp -> {
+            verify(consumer).resume(Collections.singleton(tp));
+        });
+
+        assertSinkMetricValue("sink-record-read-total", 1.0);
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        assertSinkMetricValue("sink-record-active-count", 1.0);
+        assertSinkMetricValue("sink-record-active-count-max", 1.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.5);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("batch-size-max", 1.0);
+        assertTaskMetricValue("batch-size-avg", 0.5);
+
+        // Expect commit
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        // Commit advance by one
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        // Nothing polled for this partition
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(workerCurrentOffsets);
+
+        sinkTaskContext.getValue().requestCommit();
+        time.sleep(10000L);
+        workerTask.iteration();
+
+        final ArgumentCaptor<OffsetCommitCallback> callback = ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        verify(consumer).commitAsync(eq(workerCurrentOffsets), callback.capture());
+        callback.getValue().onComplete(workerCurrentOffsets, null);
+
+        verify(sinkTask, times(4)).put(anyList());
+        assertSinkMetricValue("offset-commit-completion-total", 1.0);
+    }
+
     @Test
     public void testErrorInRebalancePartitionLoss() {
         RuntimeException exception = new RuntimeException("Revocation error");
@@ -455,18 +559,18 @@ public void testPartialRevocationAndAssignment() {
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsRevoked(singleton(TOPIC_PARTITION));
                     rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
                     rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
-                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsAssigned(singleton(TOPIC_PARTITION3));
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsLost(Collections.singleton(TOPIC_PARTITION3));
-                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsLost(singleton(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsAssigned(singleton(TOPIC_PARTITION));
                     return ConsumerRecords.empty();
                 });
 
@@ -482,21 +586,188 @@ public void testPartialRevocationAndAssignment() {
 
         // Second iteration--second call to poll, partial consumer revocation
         workerTask.iteration();
-        verify(sinkTask).close(Collections.singleton(TOPIC_PARTITION));
+        verify(sinkTask).close(singleton(TOPIC_PARTITION));
         verify(sinkTask, times(2)).put(Collections.emptyList());
 
         // Third iteration--third call to poll, partial consumer assignment
         workerTask.iteration();
-        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION3));
+        verify(sinkTask).open(singleton(TOPIC_PARTITION3));
         verify(sinkTask, times(3)).put(Collections.emptyList());
 
         // Fourth iteration--fourth call to poll, one partition lost; can't commit offsets for it, one new partition assigned
         workerTask.iteration();
-        verify(sinkTask).close(Collections.singleton(TOPIC_PARTITION3));
-        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION));
+        verify(sinkTask).close(singleton(TOPIC_PARTITION3));
+        verify(sinkTask).open(singleton(TOPIC_PARTITION));
         verify(sinkTask, times(4)).put(Collections.emptyList());
     }
 
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testTaskCancelPreventsFinalOffsetCommit() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                // Put one message through the task to get some offsets to commit
+                .thenAnswer(expectConsumerPoll(1))
+                // the second put will return after the task is stopped and cancelled (asynchronously)
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        doAnswer(invocation -> null)
+                .doAnswer(invocation -> null)
+                .doAnswer(invocation -> {
+                    workerTask.stop();
+                    workerTask.cancel();
+                    return null;
+                })
+                .when(sinkTask).put(anyList());
+
+        // task performs normal steps in advance of committing offsets
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        workerTask.execute();
+
+        // stop wakes up the consumer
+        verify(consumer).wakeup();
+
+        verify(sinkTask).close(any());
+    }
+
+    @Test
+    public void testDeliveryWithMutatingTransform() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(expectConsumerPoll(0));
+
+        expectConversionAndTransformation("newtopic_", new RecordHeaders());
+
+        workerTask.iteration(); // initial assignment
+
+        workerTask.iteration(); // first record delivered
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        sinkTaskContext.getValue().requestCommit();
+        assertTrue(sinkTaskContext.getValue().isCommitRequested());
+
+        assertNotEquals(offsets, workerTask.lastCommittedOffsets());
+        workerTask.iteration(); // triggers the commit
+
+        ArgumentCaptor<OffsetCommitCallback> callback = ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        verify(consumer).commitAsync(eq(offsets), callback.capture());
+
+        callback.getValue().onComplete(offsets, null);
+
+        assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
+        assertEquals(offsets, workerTask.lastCommittedOffsets());
+        assertEquals(0, workerTask.commitFailures());
+        assertEquals(1.0, metrics.currentMetricValueAsDouble(workerTask.taskMetricsGroup().metricGroup(), "batch-size-max"), 0.0001);
+    }
+
+    @Test
+    public void testMissingTimestampPropagation() {
+        createTask(initialState);
+        expectTaskGetTopic();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1, RecordBatch.NO_TIMESTAMP, TimestampType.CREATE_TIME, new RecordHeaders()));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver 1 record
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Collection<SinkRecord>> records = ArgumentCaptor.forClass(Collection.class);
+        verify(sinkTask, times(2)).put(records.capture());
+
+        SinkRecord record = records.getValue().iterator().next();
+
+        // we expect null for missing timestamp, the sentinel value of Record.NO_TIMESTAMP is Kafka's API
+        assertNull(record.timestamp());
+        assertEquals(TimestampType.CREATE_TIME, record.timestampType());
+    }
+
+    @Test
+    public void testTimestampPropagation() {
+        final Long timestamp = System.currentTimeMillis();
+        final TimestampType timestampType = TimestampType.CREATE_TIME;
+
+        createTask(initialState);
+        expectTaskGetTopic();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1, timestamp, timestampType, new RecordHeaders()));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver 1 record
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Collection<SinkRecord>> records = ArgumentCaptor.forClass(Collection.class);
+        verify(sinkTask, times(2)).put(records.capture());
+
+        SinkRecord record = records.getValue().iterator().next();
+
+        assertEquals(timestamp, record.timestamp());
+        assertEquals(timestampType, record.timestampType());
+    }
+
+    @Test
+    public void testTopicsRegex() {
+        Map<String, String> props = new HashMap<>(TASK_PROPS);
+        props.remove("topics");
+        props.put("topics.regex", "te.*");
+        TaskConfig taskConfig = new TaskConfig(props);
+
+        createTask(TargetState.PAUSED);
+
+        workerTask.initialize(taskConfig);
+        workerTask.initializeAndStart();
+
+        ArgumentCaptor<Pattern> topicsRegex = ArgumentCaptor.forClass(Pattern.class);
+
+        verify(consumer).subscribe(topicsRegex.capture(), rebalanceListener.capture());
+        assertEquals("te.*", topicsRegex.getValue().pattern());
+        verify(sinkTask).initialize(sinkTaskContext.capture());
+        verify(sinkTask).start(props);
+
+        expectPollInitialAssignment();
+
+        workerTask.iteration();
+        time.sleep(10000L);
+
+        verify(consumer).pause(INITIAL_ASSIGNMENT);
+    }
+
     @Test
     public void testMetricsGroup() {
         SinkTaskMetricsGroup group = new SinkTaskMetricsGroup(taskId, metrics);
@@ -558,6 +829,143 @@ public void testMetricsGroup() {
         assertEquals(30, metrics.currentMetricValueAsDouble(group1.metricGroup(), "put-batch-max-time-ms"), 0.001d);
     }
 
+    @Test
+    public void testHeaders() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        Headers headers = new RecordHeaders();
+        headers.add("header_key", "header_value".getBytes());
+
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1, headers));
+
+        expectConversionAndTransformation(null, headers);
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver 1 record
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Collection<SinkRecord>> recordCapture = ArgumentCaptor.forClass(Collection.class);
+        verify(sinkTask, times(2)).put(recordCapture.capture());
+
+        assertEquals(1, recordCapture.getValue().size());
+        SinkRecord record = recordCapture.getValue().iterator().next();
+
+        assertEquals("header_value", record.headers().lastWithName("header_key").value());
+    }
+
+    @Test
+    public void testHeadersWithCustomConverter() {
+        StringConverter stringConverter = new StringConverter();
+        SampleConverterWithHeaders testConverter = new SampleConverterWithHeaders();
+
+        createTask(initialState, stringConverter, testConverter, stringConverter);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        String keyA = "a";
+        String valueA = "Árvíztűrő tükörfúrógép";
+        Headers headersA = new RecordHeaders();
+        String encodingA = "latin2";
+        headersA.add("encoding", encodingA.getBytes());
+
+        String keyB = "b";
+        String valueB = "Тестовое сообщение";
+        Headers headersB = new RecordHeaders();
+        String encodingB = "koi8_r";
+        headersB.add("encoding", encodingB.getBytes());
+
+        expectPollInitialAssignment()
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
+                            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
+                                    0, 0, keyA.getBytes(), valueA.getBytes(encodingA), headersA, Optional.empty()),
+                            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 2, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
+                                    0, 0, keyB.getBytes(), valueB.getBytes(encodingB), headersB, Optional.empty())
+                    );
+                    return new ConsumerRecords<>(Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records));
+                });
+
+        expectTransformation(null);
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver records
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Collection<SinkRecord>> records = ArgumentCaptor.forClass(Collection.class);
+        verify(sinkTask, times(2)).put(records.capture());
+
+        Iterator<SinkRecord> iterator = records.getValue().iterator();
+
+        SinkRecord recordA = iterator.next();
+        assertEquals(keyA, recordA.key());
+        assertEquals(valueA, recordA.value());
+
+        SinkRecord recordB = iterator.next();
+        assertEquals(keyB, recordB.key());
+        assertEquals(valueB, recordB.value());
+    }
+
+    @Test
+    public void testOriginalTopicWithTopicMutatingTransformations() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation("newtopic_", new RecordHeaders());
+
+        workerTask.iteration(); // initial assignment
+        workerTask.iteration(); // first record delivered
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<Collection<SinkRecord>> recordCapture = ArgumentCaptor.forClass(Collection.class);
+        verify(sinkTask, times(2)).put(recordCapture.capture());
+
+        assertEquals(1, recordCapture.getValue().size());
+        SinkRecord record = recordCapture.getValue().iterator().next();
+        assertEquals(TOPIC, record.originalTopic());
+        assertEquals("newtopic_" + TOPIC, record.topic());
+    }
+
+    @Test
+    public void testPartitionCountInCaseOfPartitionRevocation() {
+        MockConsumer<byte[], byte[]> mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
+        // Setting up Worker Sink Task to check metrics
+        workerTask = new WorkerSinkTask(
+                taskId, sinkTask, statusListener, TargetState.PAUSED, workerConfig, ClusterConfigState.EMPTY, metrics,
+                keyConverter, valueConverter, errorHandlingMetrics, headerConverter,
+                transformationChain, mockConsumer, pluginLoader, time,
+                RetryWithToleranceOperatorTest.noopOperator(), null, statusBackingStore, Collections::emptyList);
+        mockConsumer.updateBeginningOffsets(
+                new HashMap<TopicPartition, Long>() {{
+                    put(TOPIC_PARTITION, 0L);
+                    put(TOPIC_PARTITION2, 0L);
+                }}
+        );
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        // Initial Re-balance to assign INITIAL_ASSIGNMENT which is "TOPIC_PARTITION" and "TOPIC_PARTITION2"
+        mockConsumer.rebalance(INITIAL_ASSIGNMENT);
+        assertSinkMetricValue("partition-count", 2);
+        // Revoked "TOPIC_PARTITION" and second re-balance with "TOPIC_PARTITION2"
+        mockConsumer.rebalance(Collections.singleton(TOPIC_PARTITION2));
+        assertSinkMetricValue("partition-count", 1);
+        // Closing the Worker Sink Task which will update the partition count as 0.
+        workerTask.close();
+        assertSinkMetricValue("partition-count", 0);
+    }
+
     private void expectRebalanceRevocationError(RuntimeException e) {
         when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
         doThrow(e).when(sinkTask).close(INITIAL_ASSIGNMENT);
@@ -599,6 +1007,10 @@ private Answer<ConsumerRecords<byte[], byte[]>> expectConsumerPoll(final int num
         return expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, new RecordHeaders());
     }
 
+    private Answer<ConsumerRecords<byte[], byte[]>> expectConsumerPoll(final int numMessages,  Headers headers) {
+        return expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, headers);
+    }
+
     private Answer<ConsumerRecords<byte[], byte[]>> expectConsumerPoll(final int numMessages, final long timestamp, final TimestampType timestampType, Headers headers) {
         return invocation -> {
             List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
index 66edf6a07170c..e103c30157b9d 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
@@ -35,7 +35,6 @@
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.RetriableException;
-import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
 import org.apache.kafka.connect.runtime.errors.ErrorReporter;
@@ -50,13 +49,10 @@
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.StatusBackingStore;
-import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.easymock.Capture;
-import org.easymock.CaptureType;
 import org.easymock.EasyMock;
 import org.easymock.IExpectationSetters;
-import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -71,11 +67,9 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -88,15 +82,12 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.function.Supplier;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singleton;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -134,7 +125,6 @@ public class WorkerSinkTaskTest {
     private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
 
     private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
-    private ConnectorTaskId taskId1 = new ConnectorTaskId("job", 1);
     private TargetState initialState = TargetState.STARTED;
     private MockTime time;
     private WorkerSinkTask workerTask;
@@ -162,7 +152,6 @@ public class WorkerSinkTaskTest {
     @Mock
     private ErrorHandlingMetrics errorHandlingMetrics;
     private Capture<ConsumerRebalanceListener> rebalanceListener = EasyMock.newCapture();
-    private Capture<Pattern> topicsRegex = EasyMock.newCapture();
 
     private long recordsReturnedTp1;
     private long recordsReturnedTp3;
@@ -203,103 +192,6 @@ public void tearDown() {
         if (metrics != null) metrics.stop();
     }
 
-    @Test
-    public void testPollRedelivery() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);
-        sinkTask.put(EasyMock.capture(records));
-        EasyMock.expectLastCall().andThrow(new RetriableException("retry"));
-        // Pause
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        consumer.pause(INITIAL_ASSIGNMENT);
-        PowerMock.expectLastCall();
-
-        // Retry delivery should succeed
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.capture(records));
-        EasyMock.expectLastCall();
-        // And unpause
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            consumer.resume(singleton(tp));
-            PowerMock.expectLastCall();
-        });
-
-        // Expect commit
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        // Commit advance by one
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        // Nothing polled for this partition
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        EasyMock.expect(sinkTask.preCommit(workerCurrentOffsets)).andReturn(workerCurrentOffsets);
-        final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
-        consumer.commitAsync(EasyMock.eq(workerCurrentOffsets), EasyMock.capture(callback));
-        EasyMock.expectLastCall().andAnswer(() -> {
-            callback.getValue().onComplete(workerCurrentOffsets, null);
-            return null;
-        });
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.eq(Collections.emptyList()));
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        time.sleep(10000L);
-
-        assertSinkMetricValue("partition-count", 2);
-        assertSinkMetricValue("sink-record-read-total", 0.0);
-        assertSinkMetricValue("sink-record-send-total", 0.0);
-        assertSinkMetricValue("sink-record-active-count", 0.0);
-        assertSinkMetricValue("sink-record-active-count-max", 0.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.0);
-        assertSinkMetricValue("offset-commit-seq-no", 0.0);
-        assertSinkMetricValue("offset-commit-completion-rate", 0.0);
-        assertSinkMetricValue("offset-commit-completion-total", 0.0);
-        assertSinkMetricValue("offset-commit-skip-rate", 0.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 0.0);
-        assertTaskMetricValue("batch-size-avg", 0.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
-
-        workerTask.iteration();
-        workerTask.iteration();
-        time.sleep(30000L);
-
-        assertSinkMetricValue("sink-record-read-total", 1.0);
-        assertSinkMetricValue("sink-record-send-total", 1.0);
-        assertSinkMetricValue("sink-record-active-count", 1.0);
-        assertSinkMetricValue("sink-record-active-count-max", 1.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.5);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("batch-size-max", 1.0);
-        assertTaskMetricValue("batch-size-avg", 0.5);
-        
-        sinkTaskContext.getValue().requestCommit();
-        time.sleep(10000L);
-        workerTask.iteration();
-        assertSinkMetricValue("offset-commit-completion-total", 1.0);
-
-        PowerMock.verifyAll();
-    }
-
     @Test
     public void testPollRedeliveryWithConsumerRebalance() throws Exception {
         createTask(initialState);
@@ -1054,53 +946,6 @@ public void testSuppressCloseErrors() throws Exception {
         }
     }
 
-    @Test
-    public void testTaskCancelPreventsFinalOffsetCommit() throws Exception {
-        createTask(initialState);
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        expectPollInitialAssignment();
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(1);
-
-        // Put one message through the task to get some offsets to commit
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall();
-
-        // the second put will return after the task is stopped and cancelled (asynchronously)
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall().andAnswer(() -> {
-            workerTask.stop();
-            workerTask.cancel();
-            return null;
-        });
-
-        // stop wakes up the consumer
-        consumer.wakeup();
-        EasyMock.expectLastCall();
-
-        // task performs normal steps in advance of committing offsets
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
-        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-        sinkTask.close(EasyMock.anyObject());
-        PowerMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.execute();
-
-        PowerMock.verifyAll();
-    }
-
     // Verify that when commitAsync is called but the supplied callback is not called by the consumer before a
     // rebalance occurs, the async callback does not reset the last committed offset from the rebalance.
     // See KAFKA-5731 for more information.
@@ -1318,294 +1163,6 @@ public void testCommitWithOutOfOrderCallback() throws Exception {
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testDeliveryWithMutatingTransform() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1, "newtopic_");
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-
-        final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
-        consumer.commitAsync(EasyMock.eq(offsets), EasyMock.capture(callback));
-        EasyMock.expectLastCall().andAnswer(() -> {
-            callback.getValue().onComplete(offsets, null);
-            return null;
-        });
-
-        expectConsumerPoll(0);
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-
-        workerTask.iteration(); // initial assignment
-
-        workerTask.iteration(); // first record delivered
-
-        sinkTaskContext.getValue().requestCommit();
-        assertTrue(sinkTaskContext.getValue().isCommitRequested());
-        assertNotEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-        workerTask.iteration(); // triggers the commit
-        assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
-        assertEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-        assertEquals(0, workerTask.commitFailures());
-        assertEquals(1.0, metrics.currentMetricValueAsDouble(workerTask.taskMetricsGroup().metricGroup(), "batch-size-max"), 0.0001);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testMissingTimestampPropagation() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-        expectConsumerPoll(1, RecordBatch.NO_TIMESTAMP, TimestampType.CREATE_TIME);
-        expectConversionAndTransformation(1);
-
-        Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);
-
-        sinkTask.put(EasyMock.capture(records));
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-        workerTask.iteration(); // iter 2 -- deliver 1 record
-
-        SinkRecord record = records.getValue().iterator().next();
-
-        // we expect null for missing timestamp, the sentinel value of Record.NO_TIMESTAMP is Kafka's API
-        assertNull(record.timestamp());
-        assertEquals(TimestampType.CREATE_TIME, record.timestampType());
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testTimestampPropagation() throws Exception {
-        final Long timestamp = System.currentTimeMillis();
-        final TimestampType timestampType = TimestampType.CREATE_TIME;
-
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-        expectConsumerPoll(1, timestamp, timestampType);
-        expectConversionAndTransformation(1);
-
-        Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);
-        sinkTask.put(EasyMock.capture(records));
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-        workerTask.iteration(); // iter 2 -- deliver 1 record
-
-        SinkRecord record = records.getValue().iterator().next();
-
-        assertEquals(timestamp, record.timestamp());
-        assertEquals(timestampType, record.timestampType());
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testTopicsRegex() {
-        Map<String, String> props = new HashMap<>(TASK_PROPS);
-        props.remove("topics");
-        props.put("topics.regex", "te.*");
-        TaskConfig taskConfig = new TaskConfig(props);
-
-        createTask(TargetState.PAUSED);
-
-        consumer.subscribe(EasyMock.capture(topicsRegex), EasyMock.capture(rebalanceListener));
-        PowerMock.expectLastCall();
-
-        sinkTask.initialize(EasyMock.capture(sinkTaskContext));
-        PowerMock.expectLastCall();
-        sinkTask.start(props);
-        PowerMock.expectLastCall();
-
-        expectPollInitialAssignment();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        consumer.pause(INITIAL_ASSIGNMENT);
-        PowerMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(taskConfig);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        time.sleep(10000L);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testHeaders() throws Exception {
-        Headers headers = new RecordHeaders();
-        headers.add("header_key", "header_value".getBytes());
-
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1, headers);
-        expectConversionAndTransformation(1, null, headers);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-        workerTask.iteration(); // iter 2 -- deliver 1 record
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testHeadersWithCustomConverter() throws Exception {
-        StringConverter stringConverter = new StringConverter();
-        SampleConverterWithHeaders testConverter = new SampleConverterWithHeaders();
-
-        createTask(initialState, stringConverter, testConverter, stringConverter);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        String keyA = "a";
-        String valueA = "Árvíztűrő tükörfúrógép";
-        Headers headersA = new RecordHeaders();
-        String encodingA = "latin2";
-        headersA.add("encoding", encodingA.getBytes());
-
-        String keyB = "b";
-        String valueB = "Тестовое сообщение";
-        Headers headersB = new RecordHeaders();
-        String encodingB = "koi8_r";
-        headersB.add("encoding", encodingB.getBytes());
-
-        expectConsumerPoll(Arrays.asList(
-            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
-                0, 0, keyA.getBytes(), valueA.getBytes(encodingA), headersA, Optional.empty()),
-            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 2, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
-                0, 0, keyB.getBytes(), valueB.getBytes(encodingB), headersB, Optional.empty())
-        ));
-
-        expectTransformation(2, null);
-
-        Capture<Collection<SinkRecord>> records = EasyMock.newCapture(CaptureType.ALL);
-        sinkTask.put(EasyMock.capture(records));
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-        workerTask.iteration(); // iter 2 -- deliver 1 record
-
-        Iterator<SinkRecord> iterator = records.getValue().iterator();
-
-        SinkRecord recordA = iterator.next();
-        assertEquals(keyA, recordA.key());
-        assertEquals(valueA, recordA.value());
-
-        SinkRecord recordB = iterator.next();
-        assertEquals(keyB, recordB.key());
-        assertEquals(valueB, recordB.value());
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testOriginalTopicWithTopicMutatingTransformations() {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1, "newtopic_");
-        Capture<Collection<SinkRecord>> recordCapture = EasyMock.newCapture();
-        sinkTask.put(EasyMock.capture(recordCapture));
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-
-        workerTask.iteration(); // initial assignment
-
-        workerTask.iteration(); // first record delivered
-
-        assertTrue(recordCapture.hasCaptured());
-        assertEquals(1, recordCapture.getValue().size());
-        SinkRecord record = recordCapture.getValue().iterator().next();
-        assertEquals(TOPIC, record.originalTopic());
-        assertEquals("newtopic_" + TOPIC, record.topic());
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPartitionCountInCaseOfPartitionRevocation() {
-        MockConsumer<byte[], byte[]> mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
-        // Setting up Worker Sink Task to check metrics
-        workerTask = new WorkerSinkTask(
-                taskId, sinkTask, statusListener, TargetState.PAUSED, workerConfig, ClusterConfigState.EMPTY, metrics,
-                keyConverter, valueConverter, errorHandlingMetrics, headerConverter,
-                transformationChain, mockConsumer, pluginLoader, time,
-                RetryWithToleranceOperatorTest.noopOperator(), null, statusBackingStore, Collections::emptyList);
-        mockConsumer.updateBeginningOffsets(new HashMap<TopicPartition, Long>() {{
-                put(TOPIC_PARTITION, 0 * 1L);
-                put(TOPIC_PARTITION2, 0 * 1L);
-            }});
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        // Initial Re-balance to assign INITIAL_ASSIGNMENT which is "TOPIC_PARTITION" and "TOPIC_PARTITION2"
-        mockConsumer.rebalance(INITIAL_ASSIGNMENT);
-        assertSinkMetricValue("partition-count", 2);
-        // Revoked "TOPIC_PARTITION" and second re-balance with "TOPIC_PARTITION2"
-        mockConsumer.rebalance(Collections.singleton(TOPIC_PARTITION2));
-        assertSinkMetricValue("partition-count", 1);
-        // Closing the Worker Sink Task which will update the partition count as 0.
-        workerTask.close();
-        assertSinkMetricValue("partition-count", 0);
-    }
-
     private void expectInitializeTask() {
         consumer.subscribe(EasyMock.eq(asList(TOPIC)), EasyMock.capture(rebalanceListener));
         PowerMock.expectLastCall();
@@ -1636,14 +1193,6 @@ private void expectConsumerPoll(final int numMessages) {
         expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, emptyHeaders());
     }
 
-    private void expectConsumerPoll(final int numMessages, Headers headers) {
-        expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, headers);
-    }
-
-    private void expectConsumerPoll(final int numMessages, final long timestamp, final TimestampType timestampType) {
-        expectConsumerPoll(numMessages, timestamp, timestampType, emptyHeaders());
-    }
-
     private void expectConsumerPoll(final int numMessages, final long timestamp, final TimestampType timestampType, Headers headers) {
         EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
             () -> {
@@ -1660,15 +1209,6 @@ private void expectConsumerPoll(final int numMessages, final long timestamp, fin
             });
     }
 
-    private void expectConsumerPoll(List<ConsumerRecord<byte[], byte[]>> records) {
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> new ConsumerRecords<>(
-                records.isEmpty() ?
-                    Collections.emptyMap() :
-                    Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records)
-            ));
-    }
-
     private void expectConversionAndTransformation(final int numMessages) {
         expectConversionAndTransformation(numMessages, null);
     }

From 88c5543ccfb517ceba8dc56d22576efdf2540e0e Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Tue, 13 Feb 2024 06:02:36 +0300
Subject: [PATCH 029/258] KAFKA-14589: [1/3] Tests of ConsoleGroupCommand
 rewritten in java (#15256)

This PR is part of #14471
Is contains some of ConsoleGroupCommand tests rewritten in java.
Intention of separate PR is to reduce changes and simplify review.

Reviewers: Luke Chen <showuon@gmail.com>
---
 build.gradle                                  |   2 +
 checkstyle/import-control.xml                 |   2 +
 .../kafka/admin/ConsumerGroupCommand.scala    |   2 +-
 .../admin/DeleteConsumerGroupsTest.scala      | 258 -------------
 ...sConsumerGroupCommandIntegrationTest.scala | 209 ----------
 .../unit/kafka/admin/DeleteTopicTest.scala    |   2 +
 .../kafka/admin/ListConsumerGroupTest.scala   | 147 -------
 .../kafka/tools/{reassign => }/Tuple2.java    |   2 +-
 .../reassign/ReassignPartitionsCommand.java   |   1 +
 .../group/ConsumerGroupCommandTest.java       | 359 ++++++++++++++++++
 .../group/DeleteConsumerGroupsTest.java       | 310 +++++++++++++++
 ...tsConsumerGroupCommandIntegrationTest.java | 216 +++++++++++
 .../consumer/group/ListConsumerGroupTest.java | 164 ++++++++
 .../ReassignPartitionsIntegrationTest.java    |   1 +
 .../reassign/ReassignPartitionsUnitTest.java  |   1 +
 15 files changed, 1060 insertions(+), 616 deletions(-)
 delete mode 100644 core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
 delete mode 100644 core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala
 delete mode 100644 core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
 rename tools/src/main/java/org/apache/kafka/tools/{reassign => }/Tuple2.java (97%)
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java

diff --git a/build.gradle b/build.gradle
index d365233465108..e715144d54fb1 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2012,6 +2012,8 @@ project(':tools') {
 
     testImplementation project(':clients')
     testImplementation project(':clients').sourceSets.test.output
+    testImplementation project(':server')
+    testImplementation project(':server').sourceSets.test.output
     testImplementation project(':core')
     testImplementation project(':core').sourceSets.test.output
     testImplementation project(':server-common')
diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index b6fa73b68caf6..2f9c396048309 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -332,7 +332,9 @@
     </subpackage>
 
     <subpackage name="consumer.group">
+      <allow pkg="kafka.api"/>
       <allow pkg="org.apache.kafka.tools"/>
+      <allow pkg="org.apache.kafka.server.config" />
       <allow pkg="scala"/>
     </subpackage>
 
diff --git a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
index 26c5be3b9818b..4187274a22d6e 100755
--- a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
@@ -135,7 +135,7 @@ object ConsumerGroupCommand extends Logging {
   private[admin] case class MemberAssignmentState(group: String, consumerId: String, host: String, clientId: String, groupInstanceId: String,
                                              numPartitions: Int, assignment: List[TopicPartition])
 
-  private[admin] case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
+  case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
 
   private[admin] sealed trait CsvRecord
   private[admin] case class CsvRecordWithGroup(group: String, topic: String, partition: Int, offset: Long) extends CsvRecord
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
deleted file mode 100644
index a5bdda2d02529..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
+++ /dev/null
@@ -1,258 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.admin
-
-import joptsimple.OptionException
-import kafka.utils.{TestInfoUtils, TestUtils}
-import org.apache.kafka.common.errors.{GroupIdNotFoundException, GroupNotEmptyException}
-import org.apache.kafka.common.protocol.Errors
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.ValueSource
-
-class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteWithTopicOption(quorum: String): Unit = {
-    createOffsetsTopic()
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group, "--topic")
-    assertThrows(classOf[OptionException], () => getConsumerGroupService(cgcArgs))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteCmdNonExistingGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-    val missingGroup = "missing.group"
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", missingGroup)
-    val service = getConsumerGroupService(cgcArgs)
-
-    val output = TestUtils.grabConsoleOutput(service.deleteGroups())
-    assertTrue(output.contains(s"Group '$missingGroup' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message),
-      s"The expected error (${Errors.GROUP_ID_NOT_FOUND}) was not detected while deleting consumer group")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteNonExistingGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-    val missingGroup = "missing.group"
-
-    // note the group to be deleted is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", missingGroup)
-    val service = getConsumerGroupService(cgcArgs)
-
-    val result = service.deleteGroups()
-    assertTrue(result.size == 1 && result.keySet.contains(missingGroup) && result(missingGroup).getCause.isInstanceOf[GroupIdNotFoundException],
-      s"The expected error (${Errors.GROUP_ID_NOT_FOUND}) was not detected while deleting consumer group")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteCmdNonEmptyGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group
-    addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupMembers(group, false)._2.get.size == 1
-    }, "The group did not initialize as expected.")
-
-    val output = TestUtils.grabConsoleOutput(service.deleteGroups())
-    assertTrue(output.contains(s"Group '$group' could not be deleted due to:") && output.contains(Errors.NON_EMPTY_GROUP.message),
-      s"The expected error (${Errors.NON_EMPTY_GROUP}) was not detected while deleting consumer group. Output was: (${output})")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteNonEmptyGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group
-    addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupMembers(group, false)._2.get.size == 1
-    }, "The group did not initialize as expected.")
-
-    val result = service.deleteGroups()
-    assertNotNull(result(group),
-      s"Group was deleted successfully, but it shouldn't have been. Result was:(${result})")
-    assertTrue(result.size == 1 && result.keySet.contains(group) && result(group).getCause.isInstanceOf[GroupNotEmptyException],
-      s"The expected error (${Errors.NON_EMPTY_GROUP}) was not detected while deleting consumer group. Result was:(${result})")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteCmdEmptyGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group
-    val executor = addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
-    }, "The group did not initialize as expected.")
-
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupState(group).state == "Empty"
-    }, "The group did not become empty as expected.")
-
-    val output = TestUtils.grabConsoleOutput(service.deleteGroups())
-    assertTrue(output.contains(s"Deletion of requested consumer groups ('$group') was successful."),
-      s"The consumer group could not be deleted as expected")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteCmdAllGroups(quorum: String): Unit = {
-    createOffsetsTopic()
-
-    // Create 3 groups with 1 consumer per each
-    val groups =
-      (for (i <- 1 to 3) yield {
-        val group = this.group + i
-        val executor = addConsumerGroupExecutor(numConsumers = 1, group = group)
-        group -> executor
-      }).toMap
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--all-groups")
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.listConsumerGroups().toSet == groups.keySet &&
-        groups.keySet.forall(groupId => service.collectGroupState(groupId).state == "Stable")
-    }, "The group did not initialize as expected.")
-
-    // Shutdown consumers to empty out groups
-    groups.values.foreach(executor => executor.shutdown())
-
-    TestUtils.waitUntilTrue(() => {
-      groups.keySet.forall(groupId => service.collectGroupState(groupId).state == "Empty")
-    }, "The group did not become empty as expected.")
-
-    val output = TestUtils.grabConsoleOutput(service.deleteGroups()).trim
-    val expectedGroupsForDeletion = groups.keySet
-    val deletedGroupsGrepped = output.substring(output.indexOf('(') + 1, output.indexOf(')')).split(',')
-      .map(_.replaceAll("'", "").trim).toSet
-
-    assertTrue(output.matches(s"Deletion of requested consumer groups (.*) was successful.")
-        && deletedGroupsGrepped == expectedGroupsForDeletion, s"The consumer group(s) could not be deleted as expected"
-    )
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteEmptyGroup(quorum: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group
-    val executor = addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
-    }, "The group did not initialize as expected.")
-
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupState(group).state == "Empty"
-    }, "The group did not become empty as expected.")
-
-    val result = service.deleteGroups()
-    assertTrue(result.size == 1 && result.keySet.contains(group) && result(group) == null,
-      s"The consumer group could not be deleted as expected")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteCmdWithMixOfSuccessAndError(quorum: String): Unit = {
-    createOffsetsTopic()
-    val missingGroup = "missing.group"
-
-    // run one consumer in the group
-    val executor = addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
-    }, "The group did not initialize as expected.")
-
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupState(group).state == "Empty"
-    }, "The group did not become empty as expected.")
-
-    val service2 = getConsumerGroupService(cgcArgs ++ Array("--group", missingGroup))
-    val output = TestUtils.grabConsoleOutput(service2.deleteGroups())
-    assertTrue(output.contains(s"Group '$missingGroup' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message) &&
-        output.contains(s"These consumer groups were deleted successfully: '$group'"), s"The consumer group deletion did not work as expected")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteWithMixOfSuccessAndError(quorum: String): Unit = {
-    createOffsetsTopic()
-    val missingGroup = "missing.group"
-
-    // run one consumer in the group
-    val executor = addConsumerGroupExecutor(numConsumers = 1)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
-    }, "The group did not initialize as expected.")
-
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      service.collectGroupState(group).state == "Empty"
-    }, "The group did not become empty as expected.")
-
-    val service2 = getConsumerGroupService(cgcArgs ++ Array("--group", missingGroup))
-    val result = service2.deleteGroups()
-    assertTrue(result.size == 2 &&
-        result.keySet.contains(group) && result(group) == null &&
-        result.keySet.contains(missingGroup) &&
-        result(missingGroup).getMessage.contains(Errors.GROUP_ID_NOT_FOUND.message),
-      s"The consumer group deletion did not work as expected")
-  }
-
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteWithUnrecognizedNewConsumerOption(quorum: String): Unit = {
-    val cgcArgs = Array("--new-consumer", "--bootstrap-server", bootstrapServers(), "--delete", "--group", group)
-    assertThrows(classOf[OptionException], () => getConsumerGroupService(cgcArgs))
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala
deleted file mode 100644
index e7e73a260e165..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/DeleteOffsetsConsumerGroupCommandIntegrationTest.scala
+++ /dev/null
@@ -1,209 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.admin
-
-import java.util.Properties
-import kafka.utils.TestUtils
-import kafka.utils.TestInfoUtils
-import org.apache.kafka.clients.consumer.Consumer
-import org.apache.kafka.clients.consumer.ConsumerConfig
-import org.apache.kafka.clients.consumer.KafkaConsumer
-import org.apache.kafka.clients.producer.KafkaProducer
-import org.apache.kafka.clients.producer.ProducerConfig
-import org.apache.kafka.clients.producer.ProducerRecord
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.serialization.ByteArrayDeserializer
-import org.apache.kafka.common.serialization.ByteArraySerializer
-import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.server.config.Defaults
-
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.ValueSource
-
-class DeleteOffsetsConsumerGroupCommandIntegrationTest extends ConsumerGroupCommandTest {
-
-  def getArgs(group: String, topic: String): Array[String] = {
-    Array(
-      "--bootstrap-server", bootstrapServers(),
-      "--delete-offsets",
-      "--group", group,
-      "--topic", topic
-    )
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsNonExistingGroup(quorum: String): Unit = {
-    val group = "missing.group"
-    val topic = "foo:1"
-    val service = getConsumerGroupService(getArgs(group, topic))
-
-    val (error, _) = service.deleteOffsets(group, List(topic))
-    assertEquals(Errors.GROUP_ID_NOT_FOUND, error)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfStableConsumerGroupWithTopicPartition(quorum: String): Unit = {
-    testWithStableConsumerGroup(topic, 0, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfStableConsumerGroupWithTopicOnly(quorum: String): Unit = {
-    testWithStableConsumerGroup(topic, -1, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicPartition(quorum: String): Unit = {
-    testWithStableConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicOnly(quorum: String): Unit = {
-    testWithStableConsumerGroup("foobar", -1, -1, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfEmptyConsumerGroupWithTopicPartition(quorum: String): Unit = {
-    testWithEmptyConsumerGroup(topic, 0, 0, Errors.NONE)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfEmptyConsumerGroupWithTopicOnly(quorum: String): Unit = {
-    testWithEmptyConsumerGroup(topic, -1, 0, Errors.NONE)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicPartition(quorum: String): Unit = {
-    testWithEmptyConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicOnly(quorum: String): Unit = {
-    testWithEmptyConsumerGroup("foobar", -1, -1, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-  }
-
-  private def testWithStableConsumerGroup(inputTopic: String,
-                                          inputPartition: Int,
-                                          expectedPartition: Int,
-                                          expectedError: Errors): Unit = {
-    testWithConsumerGroup(
-      withStableConsumerGroup,
-      inputTopic,
-      inputPartition,
-      expectedPartition,
-      expectedError)
-  }
-
-  private def testWithEmptyConsumerGroup(inputTopic: String,
-                                         inputPartition: Int,
-                                         expectedPartition: Int,
-                                         expectedError: Errors): Unit = {
-    testWithConsumerGroup(
-      withEmptyConsumerGroup,
-      inputTopic,
-      inputPartition,
-      expectedPartition,
-      expectedError)
-  }
-
-  private def testWithConsumerGroup(withConsumerGroup: (=> Unit) => Unit,
-                                    inputTopic: String,
-                                    inputPartition: Int,
-                                    expectedPartition: Int,
-                                    expectedError: Errors): Unit = {
-    produceRecord()
-    withConsumerGroup {
-      val topic = if (inputPartition >= 0) inputTopic + ":" + inputPartition else inputTopic
-      val service = getConsumerGroupService(getArgs(group, topic))
-      val (topLevelError, partitions) = service.deleteOffsets(group, List(topic))
-      val tp = new TopicPartition(inputTopic, expectedPartition)
-      // Partition level error should propagate to top level, unless this is due to a missed partition attempt.
-      if (inputPartition >= 0) {
-        assertEquals(expectedError, topLevelError)
-      }
-      if (expectedError == Errors.NONE)
-        assertNull(partitions(tp))
-      else
-        assertEquals(expectedError.exception, partitions(tp).getCause)
-    }
-  }
-
-  private def produceRecord(): Unit = {
-    val producer = createProducer()
-    try {
-      producer.send(new ProducerRecord(topic, 0, null, null)).get()
-    } finally {
-      Utils.closeQuietly(producer, "producer")
-    }
-  }
-
-  private def withStableConsumerGroup(body: => Unit): Unit = {
-    val consumer = createConsumer()
-    try {
-      TestUtils.subscribeAndWaitForRecords(this.topic, consumer)
-      consumer.commitSync()
-      body
-    } finally {
-      Utils.closeQuietly(consumer, "consumer")
-    }
-  }
-
-  private def withEmptyConsumerGroup(body: => Unit): Unit = {
-    val consumer = createConsumer()
-    try {
-      TestUtils.subscribeAndWaitForRecords(this.topic, consumer)
-      consumer.commitSync()
-    } finally {
-      Utils.closeQuietly(consumer, "consumer")
-    }
-    body
-  }
-
-  private def createProducer(config: Properties = new Properties()): KafkaProducer[Array[Byte], Array[Byte]] = {
-    config.putIfAbsent(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers())
-    config.putIfAbsent(ProducerConfig.ACKS_CONFIG, "-1")
-    config.putIfAbsent(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
-    config.putIfAbsent(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[ByteArraySerializer].getName)
-
-    new KafkaProducer(config)
-  }
-
-  private def createConsumer(config: Properties = new Properties()): Consumer[Array[Byte], Array[Byte]] = {
-    config.putIfAbsent(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers())
-    config.putIfAbsent(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
-    config.putIfAbsent(ConsumerConfig.GROUP_ID_CONFIG, group)
-    config.putIfAbsent(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer].getName)
-    config.putIfAbsent(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer].getName)
-    // Increase timeouts to avoid having a rebalance during the test
-    config.putIfAbsent(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, Integer.MAX_VALUE.toString)
-    config.putIfAbsent(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, Defaults.GROUP_MAX_SESSION_TIMEOUT_MS.toString)
-
-    new KafkaConsumer(config)
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
index e02266e0a48e7..894a47d5608bc 100644
--- a/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
@@ -488,6 +488,8 @@ class DeleteTopicTest extends QuorumTestHarness {
       // Restart KRaft quorum with the updated config
       val overridingProps = new Properties()
       overridingProps.put(KafkaConfig.DeleteTopicEnableProp, false.toString)
+      if (implementation != null)
+        implementation.shutdown()
       implementation = newKRaftQuorum(overridingProps)
     }
 
diff --git a/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
deleted file mode 100644
index 1aa56d51297b9..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
+++ /dev/null
@@ -1,147 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.admin
-
-import joptsimple.OptionException
-import org.junit.jupiter.api.Assertions._
-import kafka.utils.TestUtils
-import org.apache.kafka.common.ConsumerGroupState
-import org.apache.kafka.clients.admin.ConsumerGroupListing
-import java.util.Optional
-
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.ValueSource
-
-class ListConsumerGroupTest extends ConsumerGroupCommandTest {
-
-  @ParameterizedTest
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testListConsumerGroups(quorum: String): Unit = {
-    val simpleGroup = "simple-group"
-    addSimpleGroupExecutor(group = simpleGroup)
-    addConsumerGroupExecutor(numConsumers = 1)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list")
-    val service = getConsumerGroupService(cgcArgs)
-
-    val expectedGroups = Set(group, simpleGroup)
-    var foundGroups = Set.empty[String]
-    TestUtils.waitUntilTrue(() => {
-      foundGroups = service.listConsumerGroups().toSet
-      expectedGroups == foundGroups
-    }, s"Expected --list to show groups $expectedGroups, but found $foundGroups.")
-  }
-
-  @ParameterizedTest
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testListWithUnrecognizedNewConsumerOption(): Unit = {
-    val cgcArgs = Array("--new-consumer", "--bootstrap-server", bootstrapServers(), "--list")
-    assertThrows(classOf[OptionException], () => getConsumerGroupService(cgcArgs))
-  }
-
-  @ParameterizedTest
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testListConsumerGroupsWithStates(): Unit = {
-    val simpleGroup = "simple-group"
-    addSimpleGroupExecutor(group = simpleGroup)
-    addConsumerGroupExecutor(numConsumers = 1)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list", "--state")
-    val service = getConsumerGroupService(cgcArgs)
-
-    val expectedListing = Set(
-      new ConsumerGroupListing(simpleGroup, true, Optional.of(ConsumerGroupState.EMPTY)),
-      new ConsumerGroupListing(group, false, Optional.of(ConsumerGroupState.STABLE)))
-
-    var foundListing = Set.empty[ConsumerGroupListing]
-    TestUtils.waitUntilTrue(() => {
-      foundListing = service.listConsumerGroupsWithState(ConsumerGroupState.values.toSet).toSet
-      expectedListing == foundListing
-    }, s"Expected to show groups $expectedListing, but found $foundListing")
-
-    val expectedListingStable = Set(
-      new ConsumerGroupListing(group, false, Optional.of(ConsumerGroupState.STABLE)))
-
-    foundListing = Set.empty[ConsumerGroupListing]
-    TestUtils.waitUntilTrue(() => {
-      foundListing = service.listConsumerGroupsWithState(Set(ConsumerGroupState.STABLE)).toSet
-      expectedListingStable == foundListing
-    }, s"Expected to show groups $expectedListingStable, but found $foundListing")
-  }
-
-  @ParameterizedTest
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testConsumerGroupStatesFromString(quorum: String): Unit = {
-    var result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable")
-    assertEquals(Set(ConsumerGroupState.STABLE), result)
-
-    result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable, PreparingRebalance")
-    assertEquals(Set(ConsumerGroupState.STABLE, ConsumerGroupState.PREPARING_REBALANCE), result)
-
-    result = ConsumerGroupCommand.consumerGroupStatesFromString("Dead,CompletingRebalance,")
-    assertEquals(Set(ConsumerGroupState.DEAD, ConsumerGroupState.COMPLETING_REBALANCE), result)
-
-    result = ConsumerGroupCommand.consumerGroupStatesFromString("stable")
-    assertEquals(Set(ConsumerGroupState.STABLE), result)
-
-    result = ConsumerGroupCommand.consumerGroupStatesFromString("stable, assigning")
-    assertEquals(Set(ConsumerGroupState.STABLE, ConsumerGroupState.ASSIGNING), result)
-
-    result = ConsumerGroupCommand.consumerGroupStatesFromString("dead,reconciling,")
-    assertEquals(Set(ConsumerGroupState.DEAD, ConsumerGroupState.RECONCILING), result)
-
-    assertThrows(classOf[IllegalArgumentException], () => ConsumerGroupCommand.consumerGroupStatesFromString("bad, wrong"))
-
-    assertThrows(classOf[IllegalArgumentException], () => ConsumerGroupCommand.consumerGroupStatesFromString("  bad, Stable"))
-
-    assertThrows(classOf[IllegalArgumentException], () => ConsumerGroupCommand.consumerGroupStatesFromString("   ,   ,"))
-  }
-
-  @ParameterizedTest
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testListGroupCommand(quorum: String): Unit = {
-    val simpleGroup = "simple-group"
-    addSimpleGroupExecutor(group = simpleGroup)
-    addConsumerGroupExecutor(numConsumers = 1)
-    var out = ""
-
-    var cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list")
-    TestUtils.waitUntilTrue(() => {
-      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
-      !out.contains("STATE") && out.contains(simpleGroup) && out.contains(group)
-    }, s"Expected to find $simpleGroup, $group and no header, but found $out")
-
-    cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list", "--state")
-    TestUtils.waitUntilTrue(() => {
-      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
-      out.contains("STATE") && out.contains(simpleGroup) && out.contains(group)
-    }, s"Expected to find $simpleGroup, $group and the header, but found $out")
-
-    cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list", "--state", "Stable")
-    TestUtils.waitUntilTrue(() => {
-      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
-      out.contains("STATE") && out.contains(group) && out.contains("Stable")
-    }, s"Expected to find $group in state Stable and the header, but found $out")
-
-    cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--list", "--state", "stable")
-    TestUtils.waitUntilTrue(() => {
-      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
-      out.contains("STATE") && out.contains(group) && out.contains("Stable")
-    }, s"Expected to find $group in state Stable and the header, but found $out")
-  }
-
-}
diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/Tuple2.java b/tools/src/main/java/org/apache/kafka/tools/Tuple2.java
similarity index 97%
rename from tools/src/main/java/org/apache/kafka/tools/reassign/Tuple2.java
rename to tools/src/main/java/org/apache/kafka/tools/Tuple2.java
index a9e84317cf718..02dd4bf3981a8 100644
--- a/tools/src/main/java/org/apache/kafka/tools/reassign/Tuple2.java
+++ b/tools/src/main/java/org/apache/kafka/tools/Tuple2.java
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.tools.reassign;
+package org.apache.kafka.tools;
 
 import java.util.Objects;
 
diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
index 6a2d9b9346863..3623346980743 100644
--- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
@@ -49,6 +49,7 @@
 import org.apache.kafka.server.util.json.JsonValue;
 import org.apache.kafka.tools.TerseException;
 import org.apache.kafka.tools.ToolsUtils;
+import org.apache.kafka.tools.Tuple2;
 
 import java.io.IOException;
 import java.util.ArrayList;
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
new file mode 100644
index 0000000000000..b78054cb4adaa
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
@@ -0,0 +1,359 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import kafka.api.BaseConsumerTest;
+import kafka.admin.ConsumerGroupCommand;
+import kafka.server.KafkaConfig;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.GroupProtocol;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.RangeAssignor;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestInfo;
+import org.junit.jupiter.params.provider.Arguments;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+public class ConsumerGroupCommandTest extends kafka.integration.KafkaServerTestHarness {
+    public static final String TOPIC = "foo";
+    public static final String GROUP = "test.group";
+
+    List<ConsumerGroupCommand.ConsumerGroupService> consumerGroupService = new ArrayList<>();
+    List<AbstractConsumerGroupExecutor> consumerGroupExecutors = new ArrayList<>();
+
+    @Override
+    public Seq<kafka.server.KafkaConfig> generateConfigs() {
+        List<KafkaConfig> cfgs = new ArrayList<>();
+
+        TestUtils.createBrokerConfigs(
+            1,
+            zkConnectOrNull(),
+            false,
+            true,
+            scala.None$.empty(),
+            scala.None$.empty(),
+            scala.None$.empty(),
+            true,
+            false,
+            false,
+            false,
+            scala.collection.immutable.Map$.MODULE$.empty(),
+            1,
+            false,
+            1,
+            (short) 1,
+            0,
+            false
+        ).foreach(props -> {
+            if (isNewGroupCoordinatorEnabled()) {
+                props.setProperty(KafkaConfig.NewGroupCoordinatorEnableProp(), "true");
+            }
+
+            cfgs.add(KafkaConfig.fromProps(props));
+            return null;
+        });
+
+        return seq(cfgs);
+    }
+
+    @BeforeEach
+    @Override
+    public void setUp(TestInfo testInfo) {
+        super.setUp(testInfo);
+        createTopic(TOPIC, 1, 1, new Properties(), listenerName(), new Properties());
+    }
+
+    @AfterEach
+    @Override
+    public void tearDown() {
+        consumerGroupService.forEach(ConsumerGroupCommand.ConsumerGroupService::close);
+        consumerGroupExecutors.forEach(AbstractConsumerGroupExecutor::shutdown);
+        super.tearDown();
+    }
+
+    Map<TopicPartition, Long> committedOffsets(String topic, String group) {
+        try (Consumer<String, String> consumer = createNoAutoCommitConsumer(group)) {
+            Set<TopicPartition> partitions = consumer.partitionsFor(topic).stream()
+                .map(partitionInfo -> new TopicPartition(partitionInfo.topic(), partitionInfo.partition()))
+                .collect(Collectors.toSet());
+            return consumer.committed(partitions).entrySet().stream()
+                .filter(e -> e.getValue() != null)
+                .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().offset()));
+        }
+    }
+
+    Consumer<String, String> createNoAutoCommitConsumer(String group) {
+        Properties props = new Properties();
+        props.put("bootstrap.servers", bootstrapServers(listenerName()));
+        props.put("group.id", group);
+        props.put("enable.auto.commit", "false");
+        return new KafkaConsumer<>(props, new StringDeserializer(), new StringDeserializer());
+    }
+
+    ConsumerGroupCommand.ConsumerGroupService getConsumerGroupService(String[] args) {
+        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(args);
+        ConsumerGroupCommand.ConsumerGroupService service = new ConsumerGroupCommand.ConsumerGroupService(
+            opts,
+            asScala(Collections.singletonMap(AdminClientConfig.RETRIES_CONFIG, Integer.toString(Integer.MAX_VALUE)))
+        );
+
+        consumerGroupService.add(0, service);
+        return service;
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers) {
+        return addConsumerGroupExecutor(numConsumers, TOPIC, GROUP, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, GroupProtocol.CLASSIC.name);
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String groupProtocol) {
+        return addConsumerGroupExecutor(numConsumers, TOPIC, GROUP, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, groupProtocol);
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String groupProtocol, Optional<String> remoteAssignor) {
+        return addConsumerGroupExecutor(numConsumers, TOPIC, GROUP, RangeAssignor.class.getName(), remoteAssignor, Optional.empty(), false, groupProtocol);
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String topic, String group) {
+        return addConsumerGroupExecutor(numConsumers, topic, group, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, GroupProtocol.CLASSIC.name);
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String topic, String group, String groupProtocol) {
+        return addConsumerGroupExecutor(numConsumers, topic, group, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, groupProtocol);
+    }
+
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String topic, String group, String strategy, Optional<String> remoteAssignor,
+                                                   Optional<Properties> customPropsOpt, boolean syncCommit, String groupProtocol) {
+        ConsumerGroupExecutor executor = new ConsumerGroupExecutor(bootstrapServers(listenerName()), numConsumers, group, groupProtocol,
+            topic, strategy, remoteAssignor, customPropsOpt, syncCommit);
+        addExecutor(executor);
+        return executor;
+    }
+
+    SimpleConsumerGroupExecutor addSimpleGroupExecutor(String group) {
+        return addSimpleGroupExecutor(Arrays.asList(new TopicPartition(TOPIC, 0)), group);
+    }
+
+    SimpleConsumerGroupExecutor addSimpleGroupExecutor(Collection<TopicPartition> partitions, String group) {
+        SimpleConsumerGroupExecutor executor = new SimpleConsumerGroupExecutor(bootstrapServers(listenerName()), group, partitions);
+        addExecutor(executor);
+        return executor;
+    }
+
+    private AbstractConsumerGroupExecutor addExecutor(AbstractConsumerGroupExecutor executor) {
+        consumerGroupExecutors.add(0, executor);
+        return executor;
+    }
+
+    abstract class AbstractConsumerRunnable implements Runnable {
+        final String broker;
+        final String groupId;
+        final Optional<Properties> customPropsOpt;
+        final boolean syncCommit;
+
+        final Properties props = new Properties();
+        KafkaConsumer<String, String> consumer;
+
+        boolean configured = false;
+
+        public AbstractConsumerRunnable(String broker, String groupId, Optional<Properties> customPropsOpt, boolean syncCommit) {
+            this.broker = broker;
+            this.groupId = groupId;
+            this.customPropsOpt = customPropsOpt;
+            this.syncCommit = syncCommit;
+        }
+
+        void configure() {
+            configured = true;
+            configure(props);
+            customPropsOpt.ifPresent(props::putAll);
+            consumer = new KafkaConsumer<>(props);
+        }
+
+        void configure(Properties props) {
+            props.put("bootstrap.servers", broker);
+            props.put("group.id", groupId);
+            props.put("key.deserializer", StringDeserializer.class.getName());
+            props.put("value.deserializer", StringDeserializer.class.getName());
+        }
+
+        abstract void subscribe();
+
+        @Override
+        public void run() {
+            assert configured : "Must call configure before use";
+            try {
+                subscribe();
+                while (true) {
+                    consumer.poll(Duration.ofMillis(Long.MAX_VALUE));
+                    if (syncCommit)
+                        consumer.commitSync();
+                }
+            } catch (WakeupException e) {
+                // OK
+            } finally {
+                consumer.close();
+            }
+        }
+
+        void shutdown() {
+            consumer.wakeup();
+        }
+    }
+
+    class ConsumerRunnable extends AbstractConsumerRunnable {
+        final String topic;
+        final String groupProtocol;
+        final String strategy;
+        final Optional<String> remoteAssignor;
+
+        public ConsumerRunnable(String broker, String groupId, String groupProtocol, String topic, String strategy,
+                                Optional<String> remoteAssignor, Optional<Properties> customPropsOpt, boolean syncCommit) {
+            super(broker, groupId, customPropsOpt, syncCommit);
+
+            this.topic = topic;
+            this.groupProtocol = groupProtocol;
+            this.strategy = strategy;
+            this.remoteAssignor = remoteAssignor;
+        }
+
+        @Override
+        void configure(Properties props) {
+            super.configure(props);
+            props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);
+            if (groupProtocol.toUpperCase(Locale.ROOT).equals(GroupProtocol.CONSUMER.toString())) {
+                remoteAssignor.ifPresent(assignor -> props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, assignor));
+            } else {
+                props.put("partition.assignment.strategy", strategy);
+            }
+        }
+
+        @Override
+        void subscribe() {
+            consumer.subscribe(Collections.singleton(topic));
+        }
+    }
+
+    class SimpleConsumerRunnable extends AbstractConsumerRunnable {
+        final Collection<TopicPartition> partitions;
+
+        public SimpleConsumerRunnable(String broker, String groupId, Collection<TopicPartition> partitions) {
+            super(broker, groupId, Optional.empty(), false);
+
+            this.partitions = partitions;
+        }
+
+        @Override
+        void subscribe() {
+            consumer.assign(partitions);
+        }
+    }
+
+    class AbstractConsumerGroupExecutor {
+        final int numThreads;
+        final ExecutorService executor;
+        final List<AbstractConsumerRunnable> consumers = new ArrayList<>();
+
+        public AbstractConsumerGroupExecutor(int numThreads) {
+            this.numThreads = numThreads;
+            this.executor = Executors.newFixedThreadPool(numThreads);
+        }
+
+        void submit(AbstractConsumerRunnable consumerThread) {
+            consumers.add(consumerThread);
+            executor.submit(consumerThread);
+        }
+
+        void shutdown() {
+            consumers.forEach(AbstractConsumerRunnable::shutdown);
+            executor.shutdown();
+            try {
+                executor.awaitTermination(5000, TimeUnit.MILLISECONDS);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    class ConsumerGroupExecutor extends AbstractConsumerGroupExecutor {
+        public ConsumerGroupExecutor(String broker, int numConsumers, String groupId, String groupProtocol, String topic, String strategy,
+                                     Optional<String> remoteAssignor, Optional<Properties> customPropsOpt, boolean syncCommit) {
+            super(numConsumers);
+            IntStream.rangeClosed(1, numConsumers).forEach(i -> {
+                ConsumerRunnable th = new ConsumerRunnable(broker, groupId, groupProtocol, topic, strategy, remoteAssignor, customPropsOpt, syncCommit);
+                th.configure();
+                submit(th);
+            });
+        }
+    }
+
+    class SimpleConsumerGroupExecutor extends AbstractConsumerGroupExecutor {
+        public SimpleConsumerGroupExecutor(String broker, String groupId, Collection<TopicPartition> partitions) {
+            super(1);
+
+            SimpleConsumerRunnable th = new SimpleConsumerRunnable(broker, groupId, partitions);
+            th.configure();
+            submit(th);
+        }
+    }
+
+
+    public static Stream<Arguments> getTestQuorumAndGroupProtocolParametersAll() {
+        return BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll();
+    }
+
+    @SuppressWarnings({"deprecation"})
+    static <T> Seq<T> seq(Collection<T> seq) {
+        return JavaConverters.asScalaIteratorConverter(seq.iterator()).asScala().toSeq();
+    }
+
+    @SuppressWarnings("deprecation")
+    static <K, V> scala.collection.Map<K, V> asScala(Map<K, V> jmap) {
+        return JavaConverters.mapAsScalaMap(jmap);
+    }
+
+    @SuppressWarnings({"deprecation"})
+    static <T> scala.collection.immutable.Set<T> set(final Collection<T> set) {
+        return JavaConverters.asScalaSet(new HashSet<>(set)).toSet();
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
new file mode 100644
index 0000000000000..0a1bcb2daef28
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
@@ -0,0 +1,310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import joptsimple.OptionException;
+import kafka.admin.ConsumerGroupCommand;
+import org.apache.kafka.clients.consumer.GroupProtocol;
+import org.apache.kafka.clients.consumer.RangeAssignor;
+import org.apache.kafka.common.errors.GroupIdNotFoundException;
+import org.apache.kafka.common.errors.GroupNotEmptyException;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteWithTopicOption(String quorum) {
+        createOffsetsTopic(listenerName(), new Properties());
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP, "--topic"};
+        assertThrows(OptionException.class, () -> getConsumerGroupService(cgcArgs));
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteCmdNonExistingGroup(String quorum) {
+        createOffsetsTopic(listenerName(), new Properties());
+        String missingGroup = "missing.group";
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", missingGroup};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            service.deleteGroups();
+            return null;
+        });
+        assertTrue(output.contains("Group '" + missingGroup + "' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message()),
+            "The expected error (" + Errors.GROUP_ID_NOT_FOUND + ") was not detected while deleting consumer group");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteNonExistingGroup(String quorum) {
+        createOffsetsTopic(listenerName(), new Properties());
+        String missingGroup = "missing.group";
+
+        // note the group to be deleted is a different (non-existing) group
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", missingGroup};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        scala.collection.Map<String, Throwable> result = service.deleteGroups();
+        assertTrue(result.size() == 1 && result.contains(missingGroup) && result.get(missingGroup).get().getCause() instanceof GroupIdNotFoundException,
+            "The expected error (" + Errors.GROUP_ID_NOT_FOUND + ") was not detected while deleting consumer group");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteCmdNonEmptyGroup(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group
+        addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.collectGroupMembers(GROUP, false)._2.get().size() == 1,
+            "The group did not initialize as expected."
+        );
+
+        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            service.deleteGroups();
+            return null;
+        });
+        assertTrue(output.contains("Group '" + GROUP + "' could not be deleted due to:") && output.contains(Errors.NON_EMPTY_GROUP.message()),
+            "The expected error (" + Errors.NON_EMPTY_GROUP + ") was not detected while deleting consumer group. Output was: (" + output + ")");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteNonEmptyGroup(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group
+        addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.collectGroupMembers(GROUP, false)._2.get().size() == 1,
+            "The group did not initialize as expected."
+        );
+
+        scala.collection.Map<String, Throwable> result = service.deleteGroups();
+        assertNotNull(result.get(GROUP).get(),
+            "Group was deleted successfully, but it shouldn't have been. Result was:(" + result + ")");
+        assertTrue(result.size() == 1 && result.contains(GROUP) && result.get(GROUP).get().getCause() instanceof GroupNotEmptyException,
+            "The expected error (" + Errors.NON_EMPTY_GROUP + ") was not detected while deleting consumer group. Result was:(" + result + ")");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteCmdEmptyGroup(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            "The group did not initialize as expected."
+        );
+
+        executor.shutdown();
+
+        TestUtils.waitForCondition(
+            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            "The group did not become empty as expected."
+        );
+
+        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            service.deleteGroups();
+            return null;
+        });
+        assertTrue(output.contains("Deletion of requested consumer groups ('" + GROUP + "') was successful."),
+            "The consumer group could not be deleted as expected");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteCmdAllGroups(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // Create 3 groups with 1 consumer per each
+        Map<String, ConsumerGroupExecutor> groups = IntStream.rangeClosed(1, 3).mapToObj(i -> GROUP + i).collect(Collectors.toMap(
+            Function.identity(),
+            group -> addConsumerGroupExecutor(1, TOPIC, group, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, GroupProtocol.CLASSIC.name)
+        ));
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--all-groups"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() ->
+            Objects.equals(service.listConsumerGroups().toSet(), set(groups.keySet())) &&
+                groups.keySet().stream().allMatch(groupId -> {
+                    try {
+                        return Objects.equals(service.collectGroupState(groupId).state(), "Stable");
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
+                }),
+            "The group did not initialize as expected.");
+
+        // Shutdown consumers to empty out groups
+        groups.values().forEach(AbstractConsumerGroupExecutor::shutdown);
+
+        TestUtils.waitForCondition(() ->
+            groups.keySet().stream().allMatch(groupId -> {
+                try {
+                    return Objects.equals(service.collectGroupState(groupId).state(), "Empty");
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            }),
+            "The group did not become empty as expected.");
+
+        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            service.deleteGroups();
+            return null;
+        }).trim();
+        Set<String> expectedGroupsForDeletion = groups.keySet();
+        Set<String> deletedGroupsGrepped = Arrays.stream(output.substring(output.indexOf('(') + 1, output.indexOf(')')).split(","))
+            .map(str -> str.replaceAll("'", "").trim()).collect(Collectors.toSet());
+
+        assertTrue(output.matches("Deletion of requested consumer groups (.*) was successful.")
+            && Objects.equals(deletedGroupsGrepped, expectedGroupsForDeletion),
+            "The consumer group(s) could not be deleted as expected");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteEmptyGroup(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            "The group did not initialize as expected.");
+
+        executor.shutdown();
+
+        TestUtils.waitForCondition(
+            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            "The group did not become empty as expected.");
+
+        scala.collection.Map<String, Throwable> result = service.deleteGroups();
+        assertTrue(result.size() == 1 && result.contains(GROUP) && result.get(GROUP).get() == null,
+            "The consumer group could not be deleted as expected");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteCmdWithMixOfSuccessAndError(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String missingGroup = "missing.group";
+
+        // run one consumer in the group
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            "The group did not initialize as expected.");
+
+        executor.shutdown();
+
+        TestUtils.waitForCondition(
+            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            "The group did not become empty as expected.");
+
+        cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP, "--group", missingGroup};
+
+        ConsumerGroupCommand.ConsumerGroupService service2 = getConsumerGroupService(cgcArgs);
+
+        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            service2.deleteGroups();
+            return null;
+        });
+        assertTrue(output.contains("Group '" + missingGroup + "' could not be deleted due to:")
+            && output.contains(Errors.GROUP_ID_NOT_FOUND.message())
+            && output.contains("These consumer groups were deleted successfully: '" + GROUP + "'"),
+            "The consumer group deletion did not work as expected");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteWithMixOfSuccessAndError(String quorum) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String missingGroup = "missing.group";
+
+        // run one consumer in the group
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            "The group did not initialize as expected.");
+
+        executor.shutdown();
+
+        TestUtils.waitForCondition(
+            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            "The group did not become empty as expected.");
+
+        cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP, "--group", missingGroup};
+
+        ConsumerGroupCommand.ConsumerGroupService service2 = getConsumerGroupService(cgcArgs);
+        scala.collection.Map<String, Throwable> result = service2.deleteGroups();
+        assertTrue(result.size() == 2 &&
+                result.contains(GROUP) && result.get(GROUP).get() == null &&
+                result.contains(missingGroup) &&
+                result.get(missingGroup).get().getMessage().contains(Errors.GROUP_ID_NOT_FOUND.message()),
+            "The consumer group deletion did not work as expected");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteWithUnrecognizedNewConsumerOption(String quorum) {
+        String[] cgcArgs = new String[]{"--new-consumer", "--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP};
+        assertThrows(OptionException.class, () -> getConsumerGroupService(cgcArgs));
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
new file mode 100644
index 0000000000000..2ac093923ce38
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import kafka.admin.ConsumerGroupCommand;
+import kafka.utils.TestUtils;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.config.Defaults;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Collections;
+import java.util.Properties;
+import java.util.concurrent.ExecutionException;
+
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+public class DeleteOffsetsConsumerGroupCommandIntegrationTest extends ConsumerGroupCommandTest {
+    String[] getArgs(String group, String topic) {
+        return new String[] {
+            "--bootstrap-server", bootstrapServers(listenerName()),
+            "--delete-offsets",
+            "--group", group,
+            "--topic", topic
+        };
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsNonExistingGroup() {
+        String group = "missing.group";
+        String topic = "foo:1";
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(group, topic));
+
+        scala.Tuple2<Errors, scala.collection.Map<TopicPartition, Throwable>> res = service.deleteOffsets(group, seq(Collections.singleton(topic)).toList());
+        assertEquals(Errors.GROUP_ID_NOT_FOUND, res._1);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfStableConsumerGroupWithTopicPartition() {
+        testWithStableConsumerGroup(TOPIC, 0, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfStableConsumerGroupWithTopicOnly() {
+        testWithStableConsumerGroup(TOPIC, -1, 0, Errors.GROUP_SUBSCRIBED_TO_TOPIC);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicPartition() {
+        testWithStableConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfStableConsumerGroupWithUnknownTopicOnly() {
+        testWithStableConsumerGroup("foobar", -1, -1, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfEmptyConsumerGroupWithTopicPartition() {
+        testWithEmptyConsumerGroup(TOPIC, 0, 0, Errors.NONE);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfEmptyConsumerGroupWithTopicOnly() {
+        testWithEmptyConsumerGroup(TOPIC, -1, 0, Errors.NONE);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicPartition() {
+        testWithEmptyConsumerGroup("foobar", 0, 0, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDeleteOffsetsOfEmptyConsumerGroupWithUnknownTopicOnly() {
+        testWithEmptyConsumerGroup("foobar", -1, -1, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+    }
+
+    private void testWithStableConsumerGroup(String inputTopic,
+                                            int inputPartition,
+                                            int expectedPartition,
+                                            Errors expectedError) {
+        testWithConsumerGroup(
+            this::withStableConsumerGroup,
+            inputTopic,
+            inputPartition,
+            expectedPartition,
+            expectedError);
+    }
+
+    private void testWithEmptyConsumerGroup(String inputTopic,
+                                           int inputPartition,
+                                           int expectedPartition,
+                                           Errors expectedError) {
+        testWithConsumerGroup(
+            this::withEmptyConsumerGroup,
+            inputTopic,
+            inputPartition,
+            expectedPartition,
+            expectedError);
+    }
+
+    private void testWithConsumerGroup(java.util.function.Consumer<Runnable> withConsumerGroup,
+                                       String inputTopic,
+                                       int inputPartition,
+                                       int expectedPartition,
+                                       Errors expectedError) {
+        produceRecord();
+        withConsumerGroup.accept(() -> {
+            String topic = inputPartition >= 0 ? inputTopic + ":" + inputPartition : inputTopic;
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(GROUP, topic));
+            scala.Tuple2<Errors, scala.collection.Map<TopicPartition, Throwable>> res = service.deleteOffsets(GROUP, seq(Collections.singletonList(topic)).toList());
+            Errors topLevelError = res._1;
+            scala.collection.Map<TopicPartition, Throwable> partitions = res._2;
+            TopicPartition tp = new TopicPartition(inputTopic, expectedPartition);
+            // Partition level error should propagate to top level, unless this is due to a missed partition attempt.
+            if (inputPartition >= 0) {
+                assertEquals(expectedError, topLevelError);
+            }
+            if (expectedError == Errors.NONE)
+                assertNull(partitions.get(tp).get());
+            else
+                assertEquals(expectedError.exception(), partitions.get(tp).get().getCause());
+        });
+    }
+
+    private void produceRecord() {
+        KafkaProducer<byte[], byte[]> producer = createProducer(new Properties());
+        try {
+            producer.send(new ProducerRecord<>(TOPIC, 0, null, null)).get();
+        } catch (ExecutionException | InterruptedException e) {
+            throw new RuntimeException(e);
+        } finally {
+            Utils.closeQuietly(producer, "producer");
+        }
+    }
+
+    private void withStableConsumerGroup(Runnable body) {
+        Consumer<byte[], byte[]> consumer = createConsumer(new Properties());
+        try {
+            TestUtils.subscribeAndWaitForRecords(TOPIC, consumer, DEFAULT_MAX_WAIT_MS);
+            consumer.commitSync();
+            body.run();
+        } finally {
+            Utils.closeQuietly(consumer, "consumer");
+        }
+    }
+
+    private void withEmptyConsumerGroup(Runnable body) {
+        Consumer<byte[], byte[]> consumer = createConsumer(new Properties());
+        try {
+            TestUtils.subscribeAndWaitForRecords(TOPIC, consumer, DEFAULT_MAX_WAIT_MS);
+            consumer.commitSync();
+        } finally {
+            Utils.closeQuietly(consumer, "consumer");
+        }
+        body.run();
+    }
+
+    private KafkaProducer<byte[], byte[]> createProducer(Properties config) {
+        config.putIfAbsent(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers(listenerName()));
+        config.putIfAbsent(ProducerConfig.ACKS_CONFIG, "-1");
+        config.putIfAbsent(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+        config.putIfAbsent(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+
+        return new KafkaProducer<>(config);
+    }
+
+    private Consumer<byte[], byte[]> createConsumer(Properties config) {
+        config.putIfAbsent(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers(listenerName()));
+        config.putIfAbsent(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        config.putIfAbsent(ConsumerConfig.GROUP_ID_CONFIG, GROUP);
+        config.putIfAbsent(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+        config.putIfAbsent(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+        // Increase timeouts to avoid having a rebalance during the test
+        config.putIfAbsent(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+        config.putIfAbsent(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, Integer.toString(Defaults.GROUP_MAX_SESSION_TIMEOUT_MS));
+
+        return new KafkaConsumer<>(config);
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
new file mode 100644
index 0000000000000..6648cfe2a2414
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import joptsimple.OptionException;
+import kafka.admin.ConsumerGroupCommand;
+import org.apache.kafka.clients.admin.ConsumerGroupListing;
+import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class ListConsumerGroupTest extends ConsumerGroupCommandTest {
+    @ParameterizedTest
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testListConsumerGroups(String quorum) throws Exception {
+        String simpleGroup = "simple-group";
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+        scala.collection.Set<String> expectedGroups = set(Arrays.asList(GROUP, simpleGroup));
+        final AtomicReference<scala.collection.Set> foundGroups = new AtomicReference<>();
+        TestUtils.waitForCondition(() -> {
+            foundGroups.set(service.listConsumerGroups().toSet());
+            return Objects.equals(expectedGroups, foundGroups.get());
+        }, "Expected --list to show groups " + expectedGroups + ", but found " + foundGroups.get() + ".");
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testListWithUnrecognizedNewConsumerOption() {
+        String[] cgcArgs = new String[]{"--new-consumer", "--bootstrap-server", bootstrapServers(listenerName()), "--list"};
+        assertThrows(OptionException.class, () -> getConsumerGroupService(cgcArgs));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testListConsumerGroupsWithStates() throws Exception {
+        String simpleGroup = "simple-group";
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        scala.collection.Set<ConsumerGroupListing> expectedListing = set(Arrays.asList(
+            new ConsumerGroupListing(simpleGroup, true, Optional.of(ConsumerGroupState.EMPTY)),
+            new ConsumerGroupListing(GROUP, false, Optional.of(ConsumerGroupState.STABLE))));
+
+        final AtomicReference<scala.collection.Set> foundListing = new AtomicReference<>();
+        TestUtils.waitForCondition(() -> {
+            foundListing.set(service.listConsumerGroupsWithState(set(Arrays.asList(ConsumerGroupState.values()))).toSet());
+            return Objects.equals(expectedListing, foundListing.get());
+        }, "Expected to show groups " + expectedListing + ", but found " + foundListing.get());
+
+        scala.collection.Set<ConsumerGroupListing> expectedListingStable = set(Collections.singleton(
+            new ConsumerGroupListing(GROUP, false, Optional.of(ConsumerGroupState.STABLE))));
+
+        foundListing.set(null);
+
+        TestUtils.waitForCondition(() -> {
+            foundListing.set(service.listConsumerGroupsWithState(set(Collections.singleton(ConsumerGroupState.STABLE))).toSet());
+            return Objects.equals(expectedListingStable, foundListing.get());
+        }, "Expected to show groups " + expectedListingStable + ", but found " + foundListing.get());
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testConsumerGroupStatesFromString(String quorum) {
+        scala.collection.Set<ConsumerGroupState> result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable");
+        assertEquals(set(Collections.singleton(ConsumerGroupState.STABLE)), result);
+
+        result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable, PreparingRebalance");
+        assertEquals(set(Arrays.asList(ConsumerGroupState.STABLE, ConsumerGroupState.PREPARING_REBALANCE)), result);
+
+        result = ConsumerGroupCommand.consumerGroupStatesFromString("Dead,CompletingRebalance,");
+        assertEquals(set(Arrays.asList(ConsumerGroupState.DEAD, ConsumerGroupState.COMPLETING_REBALANCE)), result);
+
+        result = ConsumerGroupCommand.consumerGroupStatesFromString("stable");
+        assertEquals(set(Arrays.asList(ConsumerGroupState.STABLE)), result);
+
+        result = ConsumerGroupCommand.consumerGroupStatesFromString("stable, assigning");
+        assertEquals(set(Arrays.asList(ConsumerGroupState.STABLE, ConsumerGroupState.ASSIGNING)), result);
+
+        result = ConsumerGroupCommand.consumerGroupStatesFromString("dead,reconciling,");
+        assertEquals(set(Arrays.asList(ConsumerGroupState.DEAD, ConsumerGroupState.RECONCILING)), result);
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupStatesFromString("bad, wrong"));
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupStatesFromString("  bad, Stable"));
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupStatesFromString("   ,   ,"));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testListGroupCommand(String quorum) throws Exception {
+        String simpleGroup = "simple-group";
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1);
+        final AtomicReference<String> out = new AtomicReference<>("");
+
+        String[] cgcArgs1 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
+        TestUtils.waitForCondition(() -> {
+            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
+                ConsumerGroupCommand.main(cgcArgs1);
+                return null;
+            }));
+            return !out.get().contains("STATE") && out.get().contains(simpleGroup) && out.get().contains(GROUP);
+        }, "Expected to find " + simpleGroup + ", " + GROUP + " and no header, but found " + out.get());
+
+        String[] cgcArgs2 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"};
+        TestUtils.waitForCondition(() -> {
+            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
+                ConsumerGroupCommand.main(cgcArgs2);
+                return null;
+            }));
+            return out.get().contains("STATE") && out.get().contains(simpleGroup) && out.get().contains(GROUP);
+        }, "Expected to find " + simpleGroup + ", " + GROUP + " and the header, but found " + out.get());
+
+        String[] cgcArgs3 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "Stable"};
+        TestUtils.waitForCondition(() -> {
+            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
+                ConsumerGroupCommand.main(cgcArgs3);
+                return null;
+            }));
+            return out.get().contains("STATE") && out.get().contains(GROUP) && out.get().contains("Stable");
+        }, "Expected to find " + GROUP + " in state Stable and the header, but found " + out.get());
+
+        String[] cgcArgs4 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "stable"};
+        TestUtils.waitForCondition(() -> {
+            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
+                ConsumerGroupCommand.main(cgcArgs4);
+                return null;
+            }));
+            return out.get().contains("STATE") && out.get().contains(GROUP) && out.get().contains("Stable");
+        }, "Expected to find " + GROUP + " in state Stable and the header, but found " + out.get());
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
index b71331a6b233f..83fc665e3e487 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
@@ -43,6 +43,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.tools.TerseException;
+import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
index d0e1cb4baaf69..699e048fb206e 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
@@ -31,6 +31,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.common.AdminCommandFailedException;
 import org.apache.kafka.server.common.AdminOperationException;
+import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;

From d1d0d8ff07a4ffa62f4f125cc26317848c484cfb Mon Sep 17 00:00:00 2001
From: Emma Humber <87994408+ehumber@users.noreply.github.com>
Date: Tue, 13 Feb 2024 09:22:05 +0000
Subject: [PATCH 030/258] KSECURITY-2090 update bcprov-jdk15on for the CVE in
 this jira (#1013)

---
 build.gradle | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/build.gradle b/build.gradle
index 1d6535f9906e1..c0aa847320869 100644
--- a/build.gradle
+++ b/build.gradle
@@ -1044,10 +1044,14 @@ project(':core') {
       exclude module: 'mina-core'
     }
     testImplementation libs.apachedsCoreApi
-    testImplementation libs.apachedsInterceptorKerberos
+    testImplementation (libs.apachedsInterceptorKerberos) {
+      exclude module: 'bcprov-jdk15on'
+    }
     testImplementation libs.apachedsProtocolShared
     testImplementation libs.apachedsProtocolKerberos
-    testImplementation libs.apachedsProtocolLdap
+    testImplementation (libs.apachedsProtocolLdap) {
+      exclude module: 'bcprov-jdk15on'
+    }
     testImplementation libs.apachedsLdifPartition
     testImplementation libs.apachedsMavibotPartition
     testImplementation libs.apachedsJdbmPartition

From 011d2382689df7b850b05281f4e19387c42c1f4d Mon Sep 17 00:00:00 2001
From: Divij Vaidya <diviv@amazon.com>
Date: Tue, 13 Feb 2024 12:10:49 +0100
Subject: [PATCH 031/258] MINOR: Fix package name for
 FetchFromFollowerIntegrationTest (#15353)

Reviewers: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Josep Prat <josep.prat@aiven.io>
---
 .../kafka/server/FetchFromFollowerIntegrationTest.scala        | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/core/src/test/scala/integration/kafka/server/FetchFromFollowerIntegrationTest.scala b/core/src/test/scala/integration/kafka/server/FetchFromFollowerIntegrationTest.scala
index b5495dc2bd8f0..89b82fea31e36 100644
--- a/core/src/test/scala/integration/kafka/server/FetchFromFollowerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/FetchFromFollowerIntegrationTest.scala
@@ -14,9 +14,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package integration.kafka.server
+package kafka.server
 
-import kafka.server.{BaseFetchRequestTest, KafkaConfig}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin.NewPartitionReassignment
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer, RangeAssignor}

From fed3c3da842867510888d34ccc25d24a7a16c64e Mon Sep 17 00:00:00 2001
From: Gantigmaa Selenge <39860586+tinaselenge@users.noreply.github.com>
Date: Tue, 13 Feb 2024 17:28:28 +0000
Subject: [PATCH 032/258] KAFKA-14822: Allow restricting File and Directory
 ConfigProviders to specific paths (#14995)

Reviewers: Greg Harris <gharris1727@gmail.com>, Mickael Maison <mickael.maison@gmail.com>
---
 .../common/config/internals/AllowedPaths.java |  82 ++++++++++
 .../provider/DirectoryConfigProvider.java     |  25 ++-
 .../config/provider/FileConfigProvider.java   |  39 ++++-
 .../config/provider/AllowedPathsTest.java     | 112 +++++++++++++
 .../provider/DirectoryConfigProviderTest.java | 148 +++++++++++++-----
 .../provider/FileConfigProviderTest.java      | 114 +++++++++++++-
 .../provider/MockFileConfigProvider.java      |   5 +-
 .../provider/MockVaultConfigProvider.java     |   4 +-
 .../DynamicBrokerReconfigurationTest.scala    |   4 +-
 9 files changed, 476 insertions(+), 57 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/common/config/internals/AllowedPaths.java
 create mode 100644 clients/src/test/java/org/apache/kafka/common/config/provider/AllowedPathsTest.java

diff --git a/clients/src/main/java/org/apache/kafka/common/config/internals/AllowedPaths.java b/clients/src/main/java/org/apache/kafka/common/config/internals/AllowedPaths.java
new file mode 100644
index 0000000000000..4fb6a483f66e7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/internals/AllowedPaths.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.config.internals;
+
+import org.apache.kafka.common.config.ConfigException;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public class AllowedPaths {
+    private final List<Path> allowedPaths;
+
+    /**
+     * Constructs AllowedPaths with a list of Paths retrieved from {@code configValue}.
+     * @param configValue {@code allowed.paths} config value which is a string containing comma separated list of paths
+     * @throws ConfigException if any of the given paths is not absolute or does not exist.
+     */
+    public AllowedPaths(String configValue) {
+        this.allowedPaths = getAllowedPaths(configValue);
+    }
+
+    private List<Path> getAllowedPaths(String configValue) {
+        if (configValue != null && !configValue.isEmpty()) {
+            List<Path> allowedPaths = new ArrayList<>();
+
+            Arrays.stream(configValue.split(",")).forEach(b -> {
+                Path normalisedPath = Paths.get(b).normalize();
+
+                if (!normalisedPath.isAbsolute()) {
+                    throw new ConfigException("Path " + normalisedPath + " is not absolute");
+                } else if (!Files.exists(normalisedPath)) {
+                    throw new ConfigException("Path " + normalisedPath + " does not exist");
+                } else {
+                    allowedPaths.add(normalisedPath);
+                }
+            });
+
+            return allowedPaths;
+        }
+
+        return null;
+    }
+
+    /**
+     * Checks if the given {@code path} resides in the configured {@code allowed.paths}.
+     * If {@code allowed.paths} is not configured, the given Path is returned as allowed.
+     * @param path the Path to check if allowed
+     * @return Path that can be accessed or null if the given Path does not reside in the configured {@code allowed.paths}.
+     */
+    public Path parseUntrustedPath(String path) {
+        Path parsedPath = Paths.get(path);
+
+        if (allowedPaths != null) {
+            Path normalisedPath = parsedPath.normalize();
+            long allowed = allowedPaths.stream().filter(allowedPath -> normalisedPath.startsWith(allowedPath)).count();
+            if (allowed == 0) {
+                return null;
+            }
+            return normalisedPath;
+        }
+
+        return parsedPath;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java b/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java
index b8ca511e4f99d..2baf15a35c577 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java
@@ -18,10 +18,10 @@
 
 import org.apache.kafka.common.config.ConfigData;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.internals.AllowedPaths;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
@@ -44,8 +44,15 @@ public class DirectoryConfigProvider implements ConfigProvider {
 
     private static final Logger log = LoggerFactory.getLogger(DirectoryConfigProvider.class);
 
+    public static final String ALLOWED_PATHS_CONFIG = "allowed.paths";
+    public static final String ALLOWED_PATHS_DOC = "A comma separated list of paths that this config provider is " +
+            "allowed to access. If not set, all paths are allowed.";
+    private volatile AllowedPaths allowedPaths;
+
     @Override
-    public void configure(Map<String, ?> configs) { }
+    public void configure(Map<String, ?> configs) {
+        allowedPaths = new AllowedPaths((String) configs.getOrDefault(ALLOWED_PATHS_CONFIG, null));
+    }
 
     @Override
     public void close() throws IOException { }
@@ -75,10 +82,20 @@ public ConfigData get(String path, Set<String> keys) {
                         && keys.contains(pathname.getFileName().toString()));
     }
 
-    private static ConfigData get(String path, Predicate<Path> fileFilter) {
+    private ConfigData get(String path, Predicate<Path> fileFilter) {
+        if (allowedPaths == null) {
+            throw new IllegalStateException("The provider has not been configured yet.");
+        }
+
         Map<String, String> map = emptyMap();
+
         if (path != null && !path.isEmpty()) {
-            Path dir = new File(path).toPath();
+            Path dir = allowedPaths.parseUntrustedPath(path);
+            if (dir == null) {
+                log.warn("The path {} is not allowed to be accessed", path);
+                return new ConfigData(map);
+            }
+
             if (!Files.isDirectory(dir)) {
                 log.warn("The path {} is not a directory", path);
             } else {
diff --git a/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java b/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java
index 893ec2c49ae5d..170ebd933c8a4 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.config.ConfigData;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.internals.AllowedPaths;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -25,7 +26,7 @@
 import java.io.Reader;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
-import java.nio.file.Paths;
+import java.nio.file.Path;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.Map;
@@ -40,7 +41,13 @@ public class FileConfigProvider implements ConfigProvider {
 
     private static final Logger log = LoggerFactory.getLogger(FileConfigProvider.class);
 
+    public static final String ALLOWED_PATHS_CONFIG = "allowed.paths";
+    public static final String ALLOWED_PATHS_DOC = "A comma separated list of paths that this config provider is " +
+            "allowed to access. If not set, all paths are allowed.";
+    private volatile AllowedPaths allowedPaths;
+
     public void configure(Map<String, ?> configs) {
+        allowedPaths = new AllowedPaths((String) configs.getOrDefault(ALLOWED_PATHS_CONFIG, null));
     }
 
     /**
@@ -50,11 +57,22 @@ public void configure(Map<String, ?> configs) {
      * @return the configuration data
      */
     public ConfigData get(String path) {
+        if (allowedPaths == null) {
+            throw new IllegalStateException("The provider has not been configured yet.");
+        }
+
         Map<String, String> data = new HashMap<>();
         if (path == null || path.isEmpty()) {
             return new ConfigData(data);
         }
-        try (Reader reader = reader(path)) {
+
+        Path filePath = allowedPaths.parseUntrustedPath(path);
+        if (filePath == null) {
+            log.warn("The path {} is not allowed to be accessed", path);
+            return new ConfigData(data);
+        }
+
+        try (Reader reader = reader(filePath)) {
             Properties properties = new Properties();
             properties.load(reader);
             Enumeration<Object> keys = properties.keys();
@@ -80,11 +98,22 @@ public ConfigData get(String path) {
      * @return the configuration data
      */
     public ConfigData get(String path, Set<String> keys) {
+        if (allowedPaths == null) {
+            throw new IllegalStateException("The provider has not been configured yet.");
+        }
+
         Map<String, String> data = new HashMap<>();
         if (path == null || path.isEmpty()) {
             return new ConfigData(data);
         }
-        try (Reader reader = reader(path)) {
+
+        Path filePath = allowedPaths.parseUntrustedPath(path);
+        if (filePath == null) {
+            log.warn("The path {} is not allowed to be accessed", path);
+            return new ConfigData(data);
+        }
+
+        try (Reader reader = reader(filePath)) {
             Properties properties = new Properties();
             properties.load(reader);
             for (String key : keys) {
@@ -101,8 +130,8 @@ public ConfigData get(String path, Set<String> keys) {
     }
 
     // visible for testing
-    protected Reader reader(String path) throws IOException {
-        return Files.newBufferedReader(Paths.get(path), StandardCharsets.UTF_8);
+    protected Reader reader(Path path) throws IOException {
+        return Files.newBufferedReader(path, StandardCharsets.UTF_8);
     }
 
     public void close() {
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/AllowedPathsTest.java b/clients/src/test/java/org/apache/kafka/common/config/provider/AllowedPathsTest.java
new file mode 100644
index 0000000000000..f23d19ab7e93d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/AllowedPathsTest.java
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.config.provider;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.internals.AllowedPaths;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import org.junit.jupiter.api.io.TempDir;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class AllowedPathsTest {
+
+    private AllowedPaths allowedPaths;
+    @TempDir
+    private File parent;
+    private String dir;
+    private String myFile;
+    private String dir2;
+
+    @BeforeEach
+    public void setup() throws IOException {
+        dir = Files.createDirectory(Paths.get(parent.toString(), "dir")).toString();
+        myFile = Files.createFile(Paths.get(dir, "myFile")).toString();
+        dir2 = Files.createDirectory(Paths.get(parent.toString(), "dir2")).toString();
+    }
+
+    @Test
+    public void testAllowedPath() {
+        allowedPaths = new AllowedPaths(String.join(",", dir, dir2));
+
+        Path actual = allowedPaths.parseUntrustedPath(myFile);
+        assertEquals(myFile, actual.toString());
+    }
+
+    @Test
+    public void testNotAllowedPath() {
+        allowedPaths = new AllowedPaths(dir);
+
+        Path actual = allowedPaths.parseUntrustedPath(dir2);
+        assertNull(actual);
+    }
+
+    @Test
+    public void testNullAllowedPaths() {
+        allowedPaths = new AllowedPaths(null);
+        Path actual = allowedPaths.parseUntrustedPath(myFile);
+        assertEquals(myFile, actual.toString());
+    }
+
+    @Test
+    public void testNoTraversal() {
+        allowedPaths = new AllowedPaths(dir);
+
+        Path traversedPath = Paths.get(dir, "..", "dir2");
+        Path actual = allowedPaths.parseUntrustedPath(traversedPath.toString());
+        assertNull(actual);
+    }
+
+    @Test
+    public void testAllowedTraversal() {
+        allowedPaths = new AllowedPaths(String.join(",", dir, dir2));
+
+        Path traversedPath = Paths.get(dir, "..", "dir2");
+        Path actual = allowedPaths.parseUntrustedPath(traversedPath.toString());
+        assertEquals(traversedPath.normalize(), actual);
+    }
+
+    @Test
+    public void testNullAllowedPathsTraversal() {
+        allowedPaths = new AllowedPaths("");
+        Path traversedPath = Paths.get(dir, "..", "dir2");
+        Path actual = allowedPaths.parseUntrustedPath(traversedPath.toString());
+        // we expect non-normalised path if allowed.paths is not specified to avoid backward compatibility
+        assertEquals(traversedPath, actual);
+    }
+
+    @Test
+    public void testAllowedPathDoesNotExist() {
+        Exception e = assertThrows(ConfigException.class, () -> new AllowedPaths("/foo"));
+        assertEquals("Path /foo does not exist", e.getMessage());
+    }
+
+    @Test
+    public void testAllowedPathIsNotAbsolute() {
+        Exception e = assertThrows(ConfigException.class, () -> new AllowedPaths("foo bar "));
+        assertEquals("Path foo bar  is not absolute", e.getMessage());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java b/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java
index d162e8a6ecfba..59949e6043c3e 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java
@@ -17,93 +17,98 @@
 package org.apache.kafka.common.config.provider;
 
 import org.apache.kafka.common.config.ConfigData;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Locale;
+import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.stream.StreamSupport;
 
 import static java.util.Arrays.asList;
+
+import static org.apache.kafka.common.config.provider.DirectoryConfigProvider.ALLOWED_PATHS_CONFIG;
 import static org.apache.kafka.test.TestUtils.toSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class DirectoryConfigProviderTest {
 
     private DirectoryConfigProvider provider;
-    private File parent;
-    private File dir;
-    private File bar;
-    private File foo;
-    private File subdir;
-    private File subdirFile;
-    private File siblingDir;
-    private File siblingDirFile;
-    private File siblingFile;
-
-    private static File writeFile(File file) throws IOException {
-        Files.write(file.toPath(), file.getName().toUpperCase(Locale.ENGLISH).getBytes(StandardCharsets.UTF_8));
-        return file;
+    @TempDir
+    private Path parent;
+    private String dir;
+    private final String bar = "bar";
+    private final String foo = "foo";
+    private String subdir;
+    private final String subdirFileName = "subdirFile";
+    private String siblingDir;
+    private final String siblingDirFileName = "siblingDirFile";
+    private final String siblingFileName = "siblingFile";
+
+    private static Path writeFile(Path path) throws IOException {
+        return Files.write(path, String.valueOf(path.getFileName()).toUpperCase(Locale.ENGLISH).getBytes(StandardCharsets.UTF_8));
     }
 
     @BeforeEach
     public void setup() throws IOException {
         provider = new DirectoryConfigProvider();
         provider.configure(Collections.emptyMap());
-        parent = TestUtils.tempDirectory();
-        dir = new File(parent, "dir");
-        dir.mkdir();
-        foo = writeFile(new File(dir, "foo"));
-        bar = writeFile(new File(dir, "bar"));
-        subdir = new File(dir, "subdir");
-        subdir.mkdir();
-        subdirFile = writeFile(new File(subdir, "subdirFile"));
-        siblingDir = new File(parent, "siblingdir");
-        siblingDir.mkdir();
-        siblingDirFile = writeFile(new File(siblingDir, "siblingdirFile"));
-        siblingFile = writeFile(new File(parent, "siblingFile"));
+
+        dir = Files.createDirectory(Paths.get(parent.toString(), "dir")).toString();
+        writeFile(Files.createFile(Paths.get(dir, foo)));
+        writeFile(Files.createFile(Paths.get(dir, bar)));
+
+        subdir = Files.createDirectory(Paths.get(dir, "subdir")).toString();
+        writeFile(Files.createFile(Paths.get(subdir, subdirFileName)));
+
+        siblingDir = Files.createDirectory(Paths.get(parent.toString(), "siblingDir")).toString();
+        writeFile(Files.createFile(Paths.get(siblingDir, siblingDirFileName)));
+
+        writeFile(Files.createFile(Paths.get(parent.toString(), siblingFileName)));
     }
 
     @AfterEach
     public void close() throws IOException {
         provider.close();
-        Utils.delete(parent);
     }
 
     @Test
     public void testGetAllKeysAtPath() {
-        ConfigData configData = provider.get(dir.getAbsolutePath());
-        assertEquals(toSet(asList(foo.getName(), bar.getName())), configData.data().keySet());
-        assertEquals("FOO", configData.data().get(foo.getName()));
-        assertEquals("BAR", configData.data().get(bar.getName()));
+        ConfigData configData = provider.get(dir);
+        assertEquals(toSet(asList(foo, bar)), configData.data().keySet());
+        assertEquals("FOO", configData.data().get(foo));
+        assertEquals("BAR", configData.data().get(bar));
         assertNull(configData.ttl());
     }
 
     @Test
     public void testGetSetOfKeysAtPath() {
-        Set<String> keys = toSet(asList(foo.getName(), "baz"));
-        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
-        assertEquals(Collections.singleton(foo.getName()), configData.data().keySet());
-        assertEquals("FOO", configData.data().get(foo.getName()));
+        Set<String> keys = toSet(asList(foo, "baz"));
+        ConfigData configData = provider.get(dir, keys);
+        assertEquals(Collections.singleton(foo), configData.data().keySet());
+        assertEquals("FOO", configData.data().get(foo));
         assertNull(configData.ttl());
     }
 
     @Test
     public void testNoSubdirs() {
         // Only regular files directly in the path directory are allowed, not in subdirs
-        Set<String> keys = toSet(asList(subdir.getName(), String.join(File.separator, subdir.getName(), subdirFile.getName())));
-        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
+        Set<String> keys = toSet(asList(subdir, String.join(File.separator, subdir, subdirFileName)));
+        ConfigData configData = provider.get(dir, keys);
         assertTrue(configData.data().isEmpty());
         assertNull(configData.ttl());
     }
@@ -112,10 +117,10 @@ public void testNoSubdirs() {
     public void testNoTraversal() {
         // Check we can't escape outside the path directory
         Set<String> keys = toSet(asList(
-                String.join(File.separator, "..", siblingFile.getName()),
-                String.join(File.separator, "..", siblingDir.getName()),
-                String.join(File.separator, "..", siblingDir.getName(), siblingDirFile.getName())));
-        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
+                String.join(File.separator, "..", siblingFileName),
+                String.join(File.separator, "..", siblingDir),
+                String.join(File.separator, "..", siblingDir, siblingDirFileName)));
+        ConfigData configData = provider.get(dir, keys);
         assertTrue(configData.data().isEmpty());
         assertNull(configData.ttl());
     }
@@ -153,5 +158,64 @@ public void testServiceLoaderDiscovery() {
         ServiceLoader<ConfigProvider> serviceLoader = ServiceLoader.load(ConfigProvider.class);
         assertTrue(StreamSupport.stream(serviceLoader.spliterator(), false).anyMatch(configProvider -> configProvider instanceof DirectoryConfigProvider));
     }
+
+    @Test
+    public void testAllowedPath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, parent.toString());
+        provider.configure(configs);
+
+        ConfigData configData = provider.get(dir);
+        assertEquals(toSet(asList(foo, bar)), configData.data().keySet());
+        assertEquals("FOO", configData.data().get(foo));
+        assertEquals("BAR", configData.data().get(bar));
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testMultipleAllowedPaths() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir + "," + siblingDir);
+        provider.configure(configs);
+
+        ConfigData configData = provider.get(subdir);
+        assertEquals(toSet(asList(subdirFileName)), configData.data().keySet());
+        assertEquals("SUBDIRFILE", configData.data().get(subdirFileName));
+        assertNull(configData.ttl());
+
+        configData = provider.get(siblingDir);
+        assertEquals(toSet(asList(siblingDirFileName)), configData.data().keySet());
+        assertEquals("SIBLINGDIRFILE", configData.data().get(siblingDirFileName));
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNotAllowedPath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir);
+        provider.configure(configs);
+
+        ConfigData configData = provider.get(siblingDir);
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNoTraversalAllowedPath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir);
+        provider.configure(configs);
+
+        ConfigData configData = provider.get(Paths.get(dir, "..", "siblingDir").toString());
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNonConfiguredProvider() {
+        DirectoryConfigProvider provider2 = new DirectoryConfigProvider();
+        IllegalStateException ise = assertThrows(IllegalStateException.class, () -> provider2.get(Paths.get(dir).toString()));
+        assertEquals("The provider has not been configured yet.", ise.getMessage());
+    }
 }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java b/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java
index 5bbb60c155d57..22fb05bfc715e 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java
@@ -19,27 +19,47 @@
 import org.apache.kafka.common.config.ConfigData;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 
+import java.io.File;
 import java.io.IOException;
 import java.io.Reader;
 import java.io.StringReader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.stream.StreamSupport;
 
+import static org.apache.kafka.common.config.provider.DirectoryConfigProvider.ALLOWED_PATHS_CONFIG;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class FileConfigProviderTest {
 
     private FileConfigProvider configProvider;
+    @TempDir
+    private File parent;
+    private String dir;
+    private String dirFile;
+    private String siblingDir;
+    private String siblingDirFile;
 
     @BeforeEach
-    public void setup() {
+    public void setup() throws IOException {
         configProvider = new TestFileConfigProvider();
+        configProvider.configure(Collections.emptyMap());
+
+        dir = Files.createDirectory(Paths.get(parent.toString(), "dir")).toString();
+        dirFile = Files.createFile(Paths.get(dir, "dirFile")).toString();
+
+        siblingDir = Files.createDirectory(Paths.get(parent.toString(), "siblingDir")).toString();
+        siblingDirFile = Files.createFile(Paths.get(siblingDir, "siblingDirFile")).toString();
     }
 
     @Test
@@ -98,8 +118,98 @@ public void testServiceLoaderDiscovery() {
     public static class TestFileConfigProvider extends FileConfigProvider {
 
         @Override
-        protected Reader reader(String path) throws IOException {
+        protected Reader reader(Path path) throws IOException {
             return new StringReader("testKey=testResult\ntestKey2=testResult2");
         }
     }
+
+    @Test
+    public void testAllowedDirPath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir);
+        configProvider.configure(configs);
+
+        ConfigData configData = configProvider.get(dirFile);
+        Map<String, String> result = new HashMap<>();
+        result.put("testKey", "testResult");
+        result.put("testKey2", "testResult2");
+        assertEquals(result, configData.data());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testAllowedFilePath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dirFile);
+        configProvider.configure(configs);
+
+        ConfigData configData = configProvider.get(dirFile);
+        Map<String, String> result = new HashMap<>();
+        result.put("testKey", "testResult");
+        result.put("testKey2", "testResult2");
+        assertEquals(result, configData.data());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testMultipleAllowedPaths() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir + "," + siblingDir);
+        configProvider.configure(configs);
+
+        Map<String, String> result = new HashMap<>();
+        result.put("testKey", "testResult");
+        result.put("testKey2", "testResult2");
+
+        ConfigData configData = configProvider.get(dirFile);
+        assertEquals(result, configData.data());
+        assertNull(configData.ttl());
+
+        configData = configProvider.get(siblingDirFile);
+        assertEquals(result, configData.data());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNotAllowedDirPath() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dir);
+        configProvider.configure(configs);
+
+        ConfigData configData = configProvider.get(siblingDirFile);
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNotAllowedFilePath() throws IOException {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dirFile);
+        configProvider.configure(configs);
+
+        //another file under the same directory
+        Path dirFile2 = Files.createFile(Paths.get(dir, "dirFile2"));
+        ConfigData configData = configProvider.get(dirFile2.toString());
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNoTraversal() {
+        Map<String, String> configs = new HashMap<>();
+        configs.put(ALLOWED_PATHS_CONFIG, dirFile);
+        configProvider.configure(configs);
+
+        // Check we can't escape outside the path directory
+        ConfigData configData = configProvider.get(Paths.get(dirFile, "..", "..", "siblingDir", "siblingDirFile").toString());
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNonConfiguredProvider() {
+        FileConfigProvider provider2 = new FileConfigProvider();
+        IllegalStateException ise = assertThrows(IllegalStateException.class, () -> provider2.get(Paths.get(dirFile).toString()));
+        assertEquals("The provider has not been configured yet.", ise.getMessage());
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/MockFileConfigProvider.java b/clients/src/test/java/org/apache/kafka/common/config/provider/MockFileConfigProvider.java
index 50f60e24ab47f..22b202d9c3d54 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/provider/MockFileConfigProvider.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/MockFileConfigProvider.java
@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.io.Reader;
 import java.io.StringReader;
+import java.nio.file.Path;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -33,6 +34,8 @@ public class MockFileConfigProvider extends FileConfigProvider {
     private boolean closed = false;
 
     public void configure(Map<String, ?> configs) {
+        super.configure(configs);
+
         Object id = configs.get("testId");
         if (id == null) {
             throw new RuntimeException(getClass().getName() + " missing 'testId' config");
@@ -45,7 +48,7 @@ public void configure(Map<String, ?> configs) {
     }
 
     @Override
-    protected Reader reader(String path) throws IOException {
+    protected Reader reader(Path path) throws IOException {
         return new StringReader("key=testKey\npassword=randomPassword");
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/MockVaultConfigProvider.java b/clients/src/test/java/org/apache/kafka/common/config/provider/MockVaultConfigProvider.java
index c741798a72972..f801af40076e9 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/provider/MockVaultConfigProvider.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/MockVaultConfigProvider.java
@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.io.Reader;
 import java.io.StringReader;
+import java.nio.file.Path;
 import java.util.Map;
 
 public class MockVaultConfigProvider extends FileConfigProvider {
@@ -28,13 +29,14 @@ public class MockVaultConfigProvider extends FileConfigProvider {
     private static final String LOCATION = "location";
 
     @Override
-    protected Reader reader(String path) throws IOException {
+    protected Reader reader(Path path) throws IOException {
         String vaultLocation = (String) vaultConfigs.get(LOCATION);
         return new StringReader("truststoreKey=testTruststoreKey\ntruststorePassword=randomtruststorePassword\n" + "truststoreLocation=" + vaultLocation + "\n");
     }
 
     @Override
     public void configure(Map<String, ?> configs) {
+        super.configure(configs);
         this.vaultConfigs = configs;
         configured = true;
     }
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
index 808fe9da41997..4be668aa32113 100644
--- a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -19,7 +19,7 @@
 package kafka.server
 
 import java.io.{Closeable, File, IOException, Reader, StringReader}
-import java.nio.file.{Files, Paths, StandardCopyOption}
+import java.nio.file.{Files, Path, Paths, StandardCopyOption}
 import java.lang.management.ManagementFactory
 import java.security.KeyStore
 import java.time.Duration
@@ -2018,7 +2018,7 @@ class TestMetricsReporter extends MetricsReporter with Reconfigurable with Close
 
 class MockFileConfigProvider extends FileConfigProvider {
   @throws(classOf[IOException])
-  override def reader(path: String): Reader = {
+  override def reader(path: Path): Reader = {
     new StringReader("key=testKey\npassword=ServerPassword\ninterval=1000\nupdinterval=2000\nstoretype=JKS")
   }
 }

From 8c0488b887be4a9178563f1d72514010f83b8614 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Tue, 13 Feb 2024 19:08:13 +0100
Subject: [PATCH 033/258] MINOR: ignore heartbeat response if leaving (#15362)

When the consumer enters state LEAVING, it sets the epoch to the leave epoch,
such as -1. When the timing is right, we may get a heartbeat response after
entering the state LEAVING, which resets the epoch to the member epoch on
the server. The result is that the consumer never leaves the group.

Seems like c6f4c604d8e50ad9e182eeb66f0d1650aa44f277 changed the timing inside
the consumer to relatively frequently triggers this problem inside
`DescribeConsumerGroupTest`.

We fix it by ignoring any heartbeat responses when we are in state LEAVING.

Reviewers: David Jacot <djacot@confluent.io>
---
 .../internals/MembershipManagerImpl.java         |  5 +++++
 .../internals/MembershipManagerImplTest.java     | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index a5a28da049ec4..7132b43a869b6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -352,6 +352,11 @@ public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData respo
             );
             throw new IllegalArgumentException(errorMessage);
         }
+        if (state == MemberState.LEAVING) {
+            log.debug("Ignoring heartbeat response received from broker. Member {} with epoch {} is " +
+                    "already leaving the group.", memberId, memberEpoch);
+            return;
+        }
 
         // Update the group member id label in the client telemetry reporter if the member id has
         // changed. Initially the member id is empty, and it is updated when the member joins the
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index bb441fe7e649a..50fede2e9ef49 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -697,6 +697,22 @@ public void testLeaveGroupWhenStateIsStable() {
         testLeaveGroupReleasesAssignmentAndResetsEpochToSendLeaveGroup(membershipManager);
         verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
     }
+    @Test
+    public void testIgnoreHeartbeatWhenLeavingGroup() {
+        MembershipManager membershipManager = createMemberInStableState();
+        mockLeaveGroup();
+
+        CompletableFuture<Void> leaveResult = membershipManager.leaveGroup();
+
+        membershipManager.onHeartbeatResponseReceived(createConsumerGroupHeartbeatResponse(createAssignment(true)).data());
+
+        assertEquals(MemberState.LEAVING, membershipManager.state());
+        assertEquals(-1, membershipManager.memberEpoch());
+        assertEquals(MEMBER_ID, membershipManager.memberId());
+        assertTrue(membershipManager.currentAssignment().isEmpty());
+        assertFalse(leaveResult.isDone(), "Leave group result should not complete until the " +
+            "heartbeat request to leave is sent out.");
+    }
 
     @Test
     public void testLeaveGroupWhenStateIsReconciling() {

From 0bf830fc9c3915bc99b6e487e6083dabd593c5d3 Mon Sep 17 00:00:00 2001
From: Mickael Maison <mimaison@users.noreply.github.com>
Date: Tue, 13 Feb 2024 19:24:07 +0100
Subject: [PATCH 034/258] KAFKA-14576: Move ConsoleConsumer to tools (#15274)

Reviewers: Josep Prat <josep.prat@aiven.io>, Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
---
 bin/kafka-console-consumer.sh                 |   2 +-
 bin/windows/kafka-console-consumer.bat        |   2 +-
 checkstyle/import-control.xml                 |  18 +-
 checkstyle/suppressions.xml                   |   4 +-
 .../scala/kafka/tools/ConsoleConsumer.scala   | 647 -----------------
 .../main/scala/kafka/utils/ToolsUtils.scala   |   3 +-
 .../kafka/tools/CustomDeserializerTest.scala  |  65 --
 .../tools/DefaultMessageFormatterTest.scala   | 234 ------
 .../kafka/tools/ConsoleConsumerTest.scala     | 676 ------------------
 .../KStreamAggregationIntegrationTest.java    |   8 +-
 tests/kafkatest/services/console_consumer.py  |   7 +-
 tests/kafkatest/version.py                    |   4 +
 .../kafka/tools/consumer/ConsoleConsumer.java | 234 ++++++
 .../consumer/ConsoleConsumerOptions.java      | 414 +++++++++++
 .../consumer/DefaultMessageFormatter.java     | 210 ++++++
 .../consumer/LoggingMessageFormatter.java     |  49 ++
 .../tools/consumer/NoOpMessageFormatter.java  |  30 +
 .../apache/kafka/tools/ToolsTestUtils.java    |  10 +
 .../consumer/ConsoleConsumerOptionsTest.java  | 621 ++++++++++++++++
 .../tools/consumer/ConsoleConsumerTest.java   | 211 ++++++
 .../consumer/DefaultMessageFormatterTest.java | 147 ++++
 .../consumer/NoOpMessageFormatterTest.java    |  41 ++
 22 files changed, 1996 insertions(+), 1641 deletions(-)
 delete mode 100755 core/src/main/scala/kafka/tools/ConsoleConsumer.scala
 delete mode 100644 core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
 delete mode 100644 core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala
 delete mode 100644 core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/DefaultMessageFormatter.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/LoggingMessageFormatter.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/NoOpMessageFormatter.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/DefaultMessageFormatterTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/NoOpMessageFormatterTest.java

diff --git a/bin/kafka-console-consumer.sh b/bin/kafka-console-consumer.sh
index dbaac2b83b189..f2201462738b9 100755
--- a/bin/kafka-console-consumer.sh
+++ b/bin/kafka-console-consumer.sh
@@ -18,4 +18,4 @@ if [ "x$KAFKA_HEAP_OPTS" = "x" ]; then
     export KAFKA_HEAP_OPTS="-Xmx512M"
 fi
 
-exec $(dirname $0)/kafka-run-class.sh kafka.tools.ConsoleConsumer "$@"
+exec $(dirname $0)/kafka-run-class.sh org.apache.kafka.tools.consumer.ConsoleConsumer "$@"
diff --git a/bin/windows/kafka-console-consumer.bat b/bin/windows/kafka-console-consumer.bat
index bbbd33656ad0e..9236cfb16a38a 100644
--- a/bin/windows/kafka-console-consumer.bat
+++ b/bin/windows/kafka-console-consumer.bat
@@ -16,5 +16,5 @@ rem limitations under the License.
 
 SetLocal
 set KAFKA_HEAP_OPTS=-Xmx512M
-"%~dp0kafka-run-class.bat" kafka.tools.ConsoleConsumer %*
+"%~dp0kafka-run-class.bat" org.apache.kafka.tools.consumer.ConsoleConsumer %*
 EndLocal
diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index 2f9c396048309..caf1fe5ebe146 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -321,6 +321,17 @@
     <allow pkg="kafka.utils" />
     <allow pkg="scala.collection" />
 
+    <subpackage name="consumer">
+      <allow pkg="org.apache.kafka.tools"/>
+
+      <subpackage name="group">
+        <allow pkg="kafka.api"/>
+        <allow pkg="org.apache.kafka.tools"/>
+        <allow pkg="org.apache.kafka.server.config" />
+        <allow pkg="scala"/>
+      </subpackage>
+    </subpackage>
+
     <subpackage name="reassign">
       <allow pkg="org.apache.kafka.admin"/>
       <allow pkg="org.apache.kafka.tools"/>
@@ -331,13 +342,6 @@
       <allow pkg="scala" />
     </subpackage>
 
-    <subpackage name="consumer.group">
-      <allow pkg="kafka.api"/>
-      <allow pkg="org.apache.kafka.tools"/>
-      <allow pkg="org.apache.kafka.server.config" />
-      <allow pkg="scala"/>
-    </subpackage>
-
     <subpackage name="other">
       <allow pkg="org.apache.kafka.tools.reassign"/>
       <allow pkg="kafka.log" />
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 925c50dec860e..6a8bc0b1852eb 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -273,11 +273,11 @@
     <suppress checks="ClassDataAbstractionCoupling"
               files="VerifiableConsumer.java"/>
     <suppress checks="CyclomaticComplexity"
-              files="(StreamsResetter|ProducerPerformance|Agent).java"/>
+              files="(ConsoleConsumer|DefaultMessageFormatter|StreamsResetter|ProducerPerformance|Agent).java"/>
     <suppress checks="BooleanExpressionComplexity"
               files="StreamsResetter.java"/>
     <suppress checks="NPathComplexity"
-              files="(ProducerPerformance|StreamsResetter|Agent|TransactionalMessageCopier|ReplicaVerificationTool).java"/>
+              files="(DefaultMessageFormatter|ProducerPerformance|StreamsResetter|Agent|TransactionalMessageCopier|ReplicaVerificationTool).java"/>
     <suppress checks="ImportControl"
               files="SignalLogger.java"/>
     <suppress checks="IllegalImport"
diff --git a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala b/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
deleted file mode 100755
index 44a80dcd563ae..0000000000000
--- a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
+++ /dev/null
@@ -1,647 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.io.PrintStream
-import java.nio.charset.StandardCharsets
-import java.time.Duration
-import java.util.concurrent.CountDownLatch
-import java.util.regex.Pattern
-import java.util.{Collections, Locale, Map, Optional, Properties, Random}
-import com.typesafe.scalalogging.LazyLogging
-import joptsimple._
-import kafka.utils.Implicits._
-import kafka.utils.{Exit, _}
-import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig, ConsumerRecord, KafkaConsumer}
-import org.apache.kafka.common.{MessageFormatter, TopicPartition}
-import org.apache.kafka.common.errors.{AuthenticationException, TimeoutException, WakeupException}
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.requests.ListOffsetsRequest
-import org.apache.kafka.common.serialization.{ByteArrayDeserializer, Deserializer}
-import org.apache.kafka.common.utils.Time
-import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
-
-import scala.jdk.CollectionConverters._
-
-/**
- * Consumer that dumps messages to standard out.
- */
-object ConsoleConsumer extends Logging {
-
-  var messageCount = 0
-
-  private val shutdownLatch = new CountDownLatch(1)
-
-  def main(args: Array[String]): Unit = {
-    val conf = new ConsumerConfig(args)
-    try {
-      run(conf)
-    } catch {
-      case e: AuthenticationException =>
-        error("Authentication failed: terminating consumer process", e)
-        Exit.exit(1)
-      case e: Throwable =>
-        error("Unknown error when running consumer: ", e)
-        Exit.exit(1)
-    }
-  }
-
-  def run(conf: ConsumerConfig): Unit = {
-    val timeoutMs = if (conf.timeoutMs >= 0) conf.timeoutMs.toLong else Long.MaxValue
-    val consumer = new KafkaConsumer(consumerProps(conf), new ByteArrayDeserializer, new ByteArrayDeserializer)
-
-    val consumerWrapper =
-      if (conf.partitionArg.isDefined)
-        new ConsumerWrapper(Option(conf.topicArg), conf.partitionArg, Option(conf.offsetArg), None, consumer, timeoutMs)
-      else
-        new ConsumerWrapper(Option(conf.topicArg), None, None, Option(conf.includedTopicsArg), consumer, timeoutMs)
-
-    addShutdownHook(consumerWrapper, conf)
-
-    try process(conf.maxMessages, conf.formatter, consumerWrapper, System.out, conf.skipMessageOnError)
-    finally {
-      consumerWrapper.cleanup()
-      conf.formatter.close()
-      reportRecordCount()
-
-      shutdownLatch.countDown()
-    }
-  }
-
-  def addShutdownHook(consumer: ConsumerWrapper, conf: ConsumerConfig): Unit = {
-    Exit.addShutdownHook("consumer-shutdown-hook", {
-        consumer.wakeup()
-
-        shutdownLatch.await()
-
-        if (conf.enableSystestEventsLogging) {
-          System.out.println("shutdown_complete")
-        }
-    })
-  }
-
-  def process(maxMessages: Integer, formatter: MessageFormatter, consumer: ConsumerWrapper, output: PrintStream,
-              skipMessageOnError: Boolean): Unit = {
-    while (messageCount < maxMessages || maxMessages == -1) {
-      val msg: ConsumerRecord[Array[Byte], Array[Byte]] = try {
-        consumer.receive()
-      } catch {
-        case _: WakeupException =>
-          trace("Caught WakeupException because consumer is shutdown, ignore and terminate.")
-          // Consumer will be closed
-          return
-        case e: Throwable =>
-          error("Error processing message, terminating consumer process: ", e)
-          // Consumer will be closed
-          return
-      }
-      messageCount += 1
-      try {
-        formatter.writeTo(new ConsumerRecord(msg.topic, msg.partition, msg.offset, msg.timestamp, msg.timestampType,
-          0, 0, msg.key, msg.value, msg.headers, Optional.empty[Integer]), output)
-      } catch {
-        case e: Throwable =>
-          if (skipMessageOnError) {
-            error("Error processing message, skipping this message: ", e)
-          } else {
-            // Consumer will be closed
-            throw e
-          }
-      }
-      if (checkErr(output, formatter)) {
-        // Consumer will be closed
-        return
-      }
-    }
-  }
-
-  def reportRecordCount(): Unit = {
-    System.err.println(s"Processed a total of $messageCount messages")
-  }
-
-  def checkErr(output: PrintStream, formatter: MessageFormatter): Boolean = {
-    val gotError = output.checkError()
-    if (gotError) {
-      // This means no one is listening to our output stream anymore, time to shutdown
-      System.err.println("Unable to write to standard out, closing consumer.")
-    }
-    gotError
-  }
-
-  private[tools] def consumerProps(config: ConsumerConfig): Properties = {
-    val props = new Properties
-    props ++= config.consumerProps
-    props ++= config.extraConsumerProps
-    setAutoOffsetResetValue(config, props)
-    props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServer)
-    if (props.getProperty(ConsumerConfig.CLIENT_ID_CONFIG) == null)
-      props.put(ConsumerConfig.CLIENT_ID_CONFIG, "console-consumer")
-    CommandLineUtils.maybeMergeOptions(
-      props, ConsumerConfig.ISOLATION_LEVEL_CONFIG, config.options, config.isolationLevelOpt)
-    props
-  }
-
-  /**
-    * Used by consumerProps to retrieve the correct value for the consumer parameter 'auto.offset.reset'.
-    *
-    * Order of priority is:
-    *   1. Explicitly set parameter via --consumer.property command line parameter
-    *   2. Explicit --from-beginning given -> 'earliest'
-    *   3. Default value of 'latest'
-    *
-    * In case both --from-beginning and an explicit value are specified an error is thrown if these
-    * are conflicting.
-    */
-  def setAutoOffsetResetValue(config: ConsumerConfig, props: Properties): Unit = {
-    val (earliestConfigValue, latestConfigValue) = ("earliest", "latest")
-
-    if (props.containsKey(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)) {
-      // auto.offset.reset parameter was specified on the command line
-      val autoResetOption = props.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)
-      if (config.options.has(config.resetBeginningOpt) && earliestConfigValue != autoResetOption) {
-        // conflicting options - latest und earliest, throw an error
-        System.err.println(s"Can't simultaneously specify --from-beginning and 'auto.offset.reset=$autoResetOption', " +
-          "please remove one option")
-        Exit.exit(1)
-      }
-      // nothing to do, checking for valid parameter values happens later and the specified
-      // value was already copied during .putall operation
-    } else {
-      // no explicit value for auto.offset.reset was specified
-      // if --from-beginning was specified use earliest, otherwise default to latest
-      val autoResetOption = if (config.options.has(config.resetBeginningOpt)) earliestConfigValue else latestConfigValue
-      props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, autoResetOption)
-    }
-  }
-
-  class ConsumerConfig(args: Array[String]) extends CommandDefaultOptions(args) {
-    val topicOpt = parser.accepts("topic", "The topic to consume on.")
-      .withRequiredArg
-      .describedAs("topic")
-      .ofType(classOf[String])
-    val whitelistOpt = parser.accepts("whitelist",
-      "DEPRECATED, use --include instead; ignored if --include specified. Regular expression specifying list of topics to include for consumption.")
-      .withRequiredArg
-      .describedAs("Java regex (String)")
-      .ofType(classOf[String])
-    val includeOpt = parser.accepts("include",
-      "Regular expression specifying list of topics to include for consumption.")
-      .withRequiredArg
-      .describedAs("Java regex (String)")
-      .ofType(classOf[String])
-    val partitionIdOpt = parser.accepts("partition", "The partition to consume from. Consumption " +
-      "starts from the end of the partition unless '--offset' is specified.")
-      .withRequiredArg
-      .describedAs("partition")
-      .ofType(classOf[java.lang.Integer])
-    val offsetOpt = parser.accepts("offset", "The offset to consume from (a non-negative number), or 'earliest' which means from beginning, or 'latest' which means from end")
-      .withRequiredArg
-      .describedAs("consume offset")
-      .ofType(classOf[String])
-      .defaultsTo("latest")
-    val consumerPropertyOpt = parser.accepts("consumer-property", "A mechanism to pass user-defined properties in the form key=value to the consumer.")
-      .withRequiredArg
-      .describedAs("consumer_prop")
-      .ofType(classOf[String])
-    val consumerConfigOpt = parser.accepts("consumer.config", s"Consumer config properties file. Note that $consumerPropertyOpt takes precedence over this config.")
-      .withRequiredArg
-      .describedAs("config file")
-      .ofType(classOf[String])
-    val messageFormatterOpt = parser.accepts("formatter", "The name of a class to use for formatting kafka messages for display.")
-      .withRequiredArg
-      .describedAs("class")
-      .ofType(classOf[String])
-      .defaultsTo(classOf[DefaultMessageFormatter].getName)
-    val messageFormatterArgOpt = parser.accepts("property",
-    """The properties to initialize the message formatter. Default properties include:
-      | print.timestamp=true|false
-      | print.key=true|false
-      | print.offset=true|false
-      | print.partition=true|false
-      | print.headers=true|false
-      | print.value=true|false
-      | key.separator=<key.separator>
-      | line.separator=<line.separator>
-      | headers.separator=<line.separator>
-      | null.literal=<null.literal>
-      | key.deserializer=<key.deserializer>
-      | value.deserializer=<value.deserializer>
-      | header.deserializer=<header.deserializer>
-      |
-      |Users can also pass in customized properties for their formatter; more specifically, users can pass in properties keyed with 'key.deserializer.', 'value.deserializer.' and 'headers.deserializer.' prefixes to configure their deserializers."""
-      .stripMargin)
-      .withRequiredArg
-      .describedAs("prop")
-      .ofType(classOf[String])
-    val messageFormatterConfigOpt = parser.accepts("formatter-config", s"Config properties file to initialize the message formatter. Note that $messageFormatterArgOpt takes precedence over this config.")
-      .withRequiredArg
-      .describedAs("config file")
-      .ofType(classOf[String])
-    val resetBeginningOpt = parser.accepts("from-beginning", "If the consumer does not already have an established offset to consume from, " +
-      "start with the earliest message present in the log rather than the latest message.")
-    val maxMessagesOpt = parser.accepts("max-messages", "The maximum number of messages to consume before exiting. If not set, consumption is continual.")
-      .withRequiredArg
-      .describedAs("num_messages")
-      .ofType(classOf[java.lang.Integer])
-    val timeoutMsOpt = parser.accepts("timeout-ms", "If specified, exit if no message is available for consumption for the specified interval.")
-      .withRequiredArg
-      .describedAs("timeout_ms")
-      .ofType(classOf[java.lang.Integer])
-    val skipMessageOnErrorOpt = parser.accepts("skip-message-on-error", "If there is an error when processing a message, " +
-      "skip it instead of halt.")
-    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: The server(s) to connect to.")
-      .withRequiredArg
-      .describedAs("server to connect to")
-      .ofType(classOf[String])
-    val keyDeserializerOpt = parser.accepts("key-deserializer")
-      .withRequiredArg
-      .describedAs("deserializer for key")
-      .ofType(classOf[String])
-    val valueDeserializerOpt = parser.accepts("value-deserializer")
-      .withRequiredArg
-      .describedAs("deserializer for values")
-      .ofType(classOf[String])
-    val enableSystestEventsLoggingOpt = parser.accepts("enable-systest-events",
-      "Log lifecycle events of the consumer in addition to logging consumed " +
-        "messages. (This is specific for system tests.)")
-    val isolationLevelOpt = parser.accepts("isolation-level",
-      "Set to read_committed in order to filter out transactional messages which are not committed. Set to read_uncommitted " +
-        "to read all messages.")
-      .withRequiredArg()
-      .ofType(classOf[String])
-      .defaultsTo("read_uncommitted")
-
-    val groupIdOpt = parser.accepts("group", "The consumer group id of the consumer.")
-      .withRequiredArg
-      .describedAs("consumer group id")
-      .ofType(classOf[String])
-
-    options = tryParse(parser, args)
-
-    CommandLineUtils.maybePrintHelpOrVersion(this, "This tool helps to read data from Kafka topics and outputs it to standard output.")
-
-    var groupIdPassed = true
-    val enableSystestEventsLogging = options.has(enableSystestEventsLoggingOpt)
-
-    // topic must be specified.
-    var topicArg: String = _
-    var includedTopicsArg: String = _
-    val extraConsumerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(consumerPropertyOpt))
-    val consumerProps = if (options.has(consumerConfigOpt))
-      Utils.loadProps(options.valueOf(consumerConfigOpt))
-    else
-      new Properties()
-    val fromBeginning = options.has(resetBeginningOpt)
-    val partitionArg = if (options.has(partitionIdOpt)) Some(options.valueOf(partitionIdOpt).intValue) else None
-    val skipMessageOnError = options.has(skipMessageOnErrorOpt)
-    val messageFormatterClass = Class.forName(options.valueOf(messageFormatterOpt))
-    val formatterArgs = if (options.has(messageFormatterConfigOpt))
-      Utils.loadProps(options.valueOf(messageFormatterConfigOpt))
-    else
-      new Properties()
-    formatterArgs ++= CommandLineUtils.parseKeyValueArgs(options.valuesOf(messageFormatterArgOpt))
-    val maxMessages = if (options.has(maxMessagesOpt)) options.valueOf(maxMessagesOpt).intValue else -1
-    val timeoutMs = if (options.has(timeoutMsOpt)) options.valueOf(timeoutMsOpt).intValue else -1
-    val bootstrapServer = options.valueOf(bootstrapServerOpt)
-    val keyDeserializer = options.valueOf(keyDeserializerOpt)
-    val valueDeserializer = options.valueOf(valueDeserializerOpt)
-    val formatter: MessageFormatter = messageFormatterClass.getDeclaredConstructor().newInstance().asInstanceOf[MessageFormatter]
-
-    if (keyDeserializer != null && keyDeserializer.nonEmpty) {
-      formatterArgs.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer)
-    }
-    if (valueDeserializer != null && valueDeserializer.nonEmpty) {
-      formatterArgs.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer)
-    }
-
-    formatter.configure(formatterArgs.asScala.asJava)
-
-    topicArg = options.valueOf(topicOpt)
-    includedTopicsArg = if (options.has(includeOpt))
-      options.valueOf(includeOpt)
-    else
-      options.valueOf(whitelistOpt)
-
-    val topicOrFilterArgs = List(topicArg, includedTopicsArg).filterNot(_ == null)
-    // user need to specify value for either --topic or one of the include filters options (--include or --whitelist)
-    if (topicOrFilterArgs.size != 1)
-      CommandLineUtils.printUsageAndExit(parser, s"Exactly one of --include/--topic is required. " +
-        s"${if (options.has(whitelistOpt)) "--whitelist is DEPRECATED use --include instead; ignored if --include specified."}")
-
-    if (partitionArg.isDefined) {
-      if (!options.has(topicOpt))
-        CommandLineUtils.printUsageAndExit(parser, "The topic is required when partition is specified.")
-      if (fromBeginning && options.has(offsetOpt))
-        CommandLineUtils.printUsageAndExit(parser, "Options from-beginning and offset cannot be specified together.")
-    } else if (options.has(offsetOpt))
-      CommandLineUtils.printUsageAndExit(parser, "The partition is required when offset is specified.")
-
-    def invalidOffset(offset: String): Nothing =
-      ToolsUtils.printUsageAndExit(parser, s"The provided offset value '$offset' is incorrect. Valid values are " +
-        "'earliest', 'latest', or a non-negative long.")
-
-    val offsetArg =
-      if (options.has(offsetOpt)) {
-        options.valueOf(offsetOpt).toLowerCase(Locale.ROOT) match {
-          case "earliest" => ListOffsetsRequest.EARLIEST_TIMESTAMP
-          case "latest" => ListOffsetsRequest.LATEST_TIMESTAMP
-          case offsetString =>
-            try {
-              val offset = offsetString.toLong
-              if (offset < 0)
-                invalidOffset(offsetString)
-              offset
-            } catch {
-              case _: NumberFormatException => invalidOffset(offsetString)
-            }
-        }
-      }
-      else if (fromBeginning) ListOffsetsRequest.EARLIEST_TIMESTAMP
-      else ListOffsetsRequest.LATEST_TIMESTAMP
-
-    CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
-
-    // if the group id is provided in more than place (through different means) all values must be the same
-    val groupIdsProvided = Set(
-      Option(options.valueOf(groupIdOpt)), // via --group
-      Option(consumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)), // via --consumer-property
-      Option(extraConsumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)) // via --consumer.config
-    ).flatten
-
-    if (groupIdsProvided.size > 1) {
-      CommandLineUtils.printUsageAndExit(parser, "The group ids provided in different places (directly using '--group', "
-        + "via '--consumer-property', or via '--consumer.config') do not match. "
-        + s"Detected group ids: ${groupIdsProvided.mkString("'", "', '", "'")}")
-    }
-
-    groupIdsProvided.headOption match {
-      case Some(group) =>
-        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, group)
-      case None =>
-        consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, s"console-consumer-${new Random().nextInt(100000)}")
-        // By default, avoid unnecessary expansion of the coordinator cache since
-        // the auto-generated group and its offsets is not intended to be used again
-        if (!consumerProps.containsKey(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG))
-          consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-        groupIdPassed = false
-    }
-
-    if (groupIdPassed && partitionArg.isDefined)
-      CommandLineUtils.printUsageAndExit(parser, "Options group and partition cannot be specified together.")
-
-    def tryParse(parser: OptionParser, args: Array[String]): OptionSet = {
-      try
-        parser.parse(args: _*)
-      catch {
-        case e: OptionException =>
-          ToolsUtils.printUsageAndExit(parser, e.getMessage)
-      }
-    }
-  }
-
-  private[tools] class ConsumerWrapper(
-    topic: Option[String],
-    partitionId: Option[Int],
-    offset: Option[Long],
-    includedTopics: Option[String],
-    consumer: Consumer[Array[Byte], Array[Byte]],
-    timeoutMs: Long = Long.MaxValue,
-    time: Time = Time.SYSTEM
-  ) {
-    consumerInit()
-    var recordIter = Collections.emptyList[ConsumerRecord[Array[Byte], Array[Byte]]]().iterator()
-
-    def consumerInit(): Unit = {
-      (topic, partitionId, offset, includedTopics) match {
-        case (Some(topic), Some(partitionId), Some(offset), None) =>
-          seek(topic, partitionId, offset)
-        case (Some(topic), Some(partitionId), None, None) =>
-          // default to latest if no offset is provided
-          seek(topic, partitionId, ListOffsetsRequest.LATEST_TIMESTAMP)
-        case (Some(topic), None, None, None) =>
-          consumer.subscribe(Collections.singletonList(topic))
-        case (None, None, None, Some(include)) =>
-          consumer.subscribe(Pattern.compile(include))
-        case _ =>
-          throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
-            "Exactly one of 'topic' or 'include' must be provided. " +
-            "If 'topic' is provided, an optional 'partition' may also be provided. " +
-            "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.")
-      }
-    }
-
-    def seek(topic: String, partitionId: Int, offset: Long): Unit = {
-      val topicPartition = new TopicPartition(topic, partitionId)
-      consumer.assign(Collections.singletonList(topicPartition))
-      offset match {
-        case ListOffsetsRequest.EARLIEST_TIMESTAMP => consumer.seekToBeginning(Collections.singletonList(topicPartition))
-        case ListOffsetsRequest.LATEST_TIMESTAMP => consumer.seekToEnd(Collections.singletonList(topicPartition))
-        case _ => consumer.seek(topicPartition, offset)
-      }
-    }
-
-    def resetUnconsumedOffsets(): Unit = {
-      val smallestUnconsumedOffsets = collection.mutable.Map[TopicPartition, Long]()
-      while (recordIter.hasNext) {
-        val record = recordIter.next()
-        val tp = new TopicPartition(record.topic, record.partition)
-        // avoid auto-committing offsets which haven't been consumed
-        smallestUnconsumedOffsets.getOrElseUpdate(tp, record.offset)
-      }
-      smallestUnconsumedOffsets.forKeyValue { (tp, offset) => consumer.seek(tp, offset) }
-    }
-
-    def receive(): ConsumerRecord[Array[Byte], Array[Byte]] = {
-      val startTimeMs = time.milliseconds
-      while (!recordIter.hasNext) {
-        recordIter = consumer.poll(Duration.ofMillis(timeoutMs)).iterator
-        if (!recordIter.hasNext && (time.milliseconds - startTimeMs > timeoutMs)) {
-          throw new TimeoutException()
-        }
-      }
-
-      recordIter.next
-    }
-
-    def wakeup(): Unit = {
-      this.consumer.wakeup()
-    }
-
-    def cleanup(): Unit = {
-      resetUnconsumedOffsets()
-      this.consumer.close()
-    }
-
-  }
-}
-
-class DefaultMessageFormatter extends MessageFormatter {
-  var printTimestamp = false
-  var printKey = false
-  var printValue = true
-  var printPartition = false
-  var printOffset = false
-  var printHeaders = false
-  var keySeparator = utfBytes("\t")
-  var lineSeparator = utfBytes("\n")
-  var headersSeparator = utfBytes(",")
-  var nullLiteral = utfBytes("null")
-
-  var keyDeserializer: Option[Deserializer[_]] = None
-  var valueDeserializer: Option[Deserializer[_]] = None
-  var headersDeserializer: Option[Deserializer[_]] = None
-
-  override def configure(configs: Map[String, _]): Unit = {
-    getPropertyIfExists(configs, "print.timestamp", getBoolProperty).foreach(printTimestamp = _)
-    getPropertyIfExists(configs, "print.key", getBoolProperty).foreach(printKey = _)
-    getPropertyIfExists(configs, "print.offset", getBoolProperty).foreach(printOffset = _)
-    getPropertyIfExists(configs, "print.partition", getBoolProperty).foreach(printPartition = _)
-    getPropertyIfExists(configs, "print.headers", getBoolProperty).foreach(printHeaders = _)
-    getPropertyIfExists(configs, "print.value", getBoolProperty).foreach(printValue = _)
-    getPropertyIfExists(configs, "key.separator", getByteProperty).foreach(keySeparator = _)
-    getPropertyIfExists(configs, "line.separator", getByteProperty).foreach(lineSeparator = _)
-    getPropertyIfExists(configs, "headers.separator", getByteProperty).foreach(headersSeparator = _)
-    getPropertyIfExists(configs, "null.literal", getByteProperty).foreach(nullLiteral = _)
-
-    keyDeserializer = getPropertyIfExists(configs, "key.deserializer", getDeserializerProperty(true))
-    valueDeserializer = getPropertyIfExists(configs, "value.deserializer", getDeserializerProperty(false))
-    headersDeserializer = getPropertyIfExists(configs, "headers.deserializer", getDeserializerProperty(false))
-  }
-
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
-
-    def writeSeparator(columnSeparator: Boolean): Unit = {
-      if (columnSeparator)
-        output.write(keySeparator)
-      else
-        output.write(lineSeparator)
-    }
-
-    def deserialize(deserializer: Option[Deserializer[_]], sourceBytes: Array[Byte], topic: String) = {
-      val nonNullBytes = Option(sourceBytes).getOrElse(nullLiteral)
-      val convertedBytes = deserializer
-        .map(d => utfBytes(d.deserialize(topic, consumerRecord.headers, nonNullBytes).toString))
-        .getOrElse(nonNullBytes)
-      convertedBytes
-    }
-
-    import consumerRecord._
-
-    if (printTimestamp) {
-      if (timestampType != TimestampType.NO_TIMESTAMP_TYPE)
-        output.write(utfBytes(s"$timestampType:$timestamp"))
-      else
-        output.write(utfBytes("NO_TIMESTAMP"))
-      writeSeparator(columnSeparator =  printOffset || printPartition || printHeaders || printKey || printValue)
-    }
-
-    if (printPartition) {
-      output.write(utfBytes("Partition:"))
-      output.write(utfBytes(partition().toString))
-      writeSeparator(columnSeparator = printOffset || printHeaders || printKey || printValue)
-    }
-
-    if (printOffset) {
-      output.write(utfBytes("Offset:"))
-      output.write(utfBytes(offset().toString))
-      writeSeparator(columnSeparator = printHeaders || printKey || printValue)
-    }
-
-    if (printHeaders) {
-      val headersIt = headers().iterator.asScala
-      if (headersIt.hasNext) {
-        headersIt.foreach { header =>
-          output.write(utfBytes(header.key() + ":"))
-          output.write(deserialize(headersDeserializer, header.value(), topic))
-          if (headersIt.hasNext) {
-            output.write(headersSeparator)
-          }
-        }
-      } else {
-        output.write(utfBytes("NO_HEADERS"))
-      }
-      writeSeparator(columnSeparator = printKey || printValue)
-    }
-
-    if (printKey) {
-      output.write(deserialize(keyDeserializer, key, topic))
-      writeSeparator(columnSeparator = printValue)
-    }
-
-    if (printValue) {
-      output.write(deserialize(valueDeserializer, value, topic))
-      output.write(lineSeparator)
-    }
-  }
-
-  private def propertiesWithKeyPrefixStripped(prefix: String, configs: Map[String, _]): Map[String, _] = {
-    val newConfigs = collection.mutable.Map[String, Any]()
-    configs.asScala.foreach { case (key, value) =>
-      if (key.startsWith(prefix) && key.length > prefix.length)
-        newConfigs.put(key.substring(prefix.length), value)
-    }
-    newConfigs.asJava
-  }
-
-  private def utfBytes(str: String) = str.getBytes(StandardCharsets.UTF_8)
-
-  private def getByteProperty(configs: Map[String, _], key: String): Array[Byte] = {
-    utfBytes(configs.get(key).asInstanceOf[String])
-  }
-
-  private def getBoolProperty(configs: Map[String, _], key: String): Boolean = {
-    configs.get(key).asInstanceOf[String].trim.equalsIgnoreCase("true")
-  }
-
-  private def getDeserializerProperty(isKey: Boolean)(configs: Map[String, _], propertyName: String): Deserializer[_] = {
-    val deserializer = Class.forName(configs.get(propertyName).asInstanceOf[String]).getDeclaredConstructor().newInstance().asInstanceOf[Deserializer[_]]
-    val deserializerConfig = propertiesWithKeyPrefixStripped(propertyName + ".", configs)
-      .asScala
-      .asJava
-    deserializer.configure(deserializerConfig, isKey)
-    deserializer
-  }
-
-  private def getPropertyIfExists[T](configs: Map[String, _], key: String, getter: (Map[String, _], String) => T): Option[T] = {
-    if (configs.containsKey(key))
-      Some(getter(configs, key))
-    else
-      None
-  }
-}
-
-class LoggingMessageFormatter extends MessageFormatter with LazyLogging {
-  private val defaultWriter: DefaultMessageFormatter = new DefaultMessageFormatter
-
-  override def configure(configs: Map[String, _]): Unit = defaultWriter.configure(configs)
-
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
-    import consumerRecord._
-    defaultWriter.writeTo(consumerRecord, output)
-    logger.info({if (timestampType != TimestampType.NO_TIMESTAMP_TYPE) s"$timestampType:$timestamp, " else ""} +
-                  s"key:${if (key == null) "null" else new String(key, StandardCharsets.UTF_8)}, " +
-                  s"value:${if (value == null) "null" else new String(value, StandardCharsets.UTF_8)}")
-  }
-}
-
-class NoOpMessageFormatter extends MessageFormatter {
-
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {}
-}
-
diff --git a/core/src/main/scala/kafka/utils/ToolsUtils.scala b/core/src/main/scala/kafka/utils/ToolsUtils.scala
index 7a2aa3111c557..edd79c00030b9 100644
--- a/core/src/main/scala/kafka/utils/ToolsUtils.scala
+++ b/core/src/main/scala/kafka/utils/ToolsUtils.scala
@@ -69,8 +69,7 @@ object ToolsUtils {
   /**
    * This is a simple wrapper around `CommandLineUtils.printUsageAndExit`.
    * It is needed for tools migration (KAFKA-14525), as there is no Java equivalent for return type `Nothing`.
-   * Can be removed once [[kafka.admin.ConsumerGroupCommand]], [[kafka.tools.ConsoleConsumer]]
-   * and [[kafka.tools.ConsoleProducer]] are migrated.
+   * Can be removed once [[kafka.admin.ConsumerGroupCommand]] and [[kafka.tools.ConsoleProducer]] are migrated.
    *
    * @param parser Command line options parser.
    * @param message Error message.
diff --git a/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala b/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
deleted file mode 100644
index 244a9cfb148d0..0000000000000
--- a/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
+++ /dev/null
@@ -1,65 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.io.PrintStream
-
-import kafka.utils.TestUtils
-import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.header.Headers
-import org.apache.kafka.common.serialization.Deserializer
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.Test
-import org.mockito.Mockito._
-
-class CustomDeserializer extends Deserializer[String] {
-
-  override def deserialize(topic: String, data: Array[Byte]): String = {
-    assertNotNull(topic, "topic must not be null")
-    new String(data)
-  }
-
-  override def deserialize(topic: String, headers: Headers, data: Array[Byte]): String = {
-    println("WITH HEADERS")
-    new String(data)
-  }
-}
-
-class CustomDeserializerTest {
-
-  @Test
-  def checkFormatterCallDeserializerWithHeaders(): Unit = {
-    val formatter = new DefaultMessageFormatter()
-    formatter.valueDeserializer = Some(new CustomDeserializer)
-    val output = TestUtils.grabConsoleOutput(formatter.writeTo(
-      new ConsumerRecord("topic_test", 1, 1L, "key".getBytes, "value".getBytes), mock(classOf[PrintStream])))
-    assertTrue(output.contains("WITH HEADERS"), "DefaultMessageFormatter should call `deserialize` method with headers.")
-    formatter.close()
-  }
-
-  @Test
-  def checkDeserializerTopicIsNotNull(): Unit = {
-    val formatter = new DefaultMessageFormatter()
-    formatter.keyDeserializer = Some(new CustomDeserializer)
-
-    formatter.writeTo(new ConsumerRecord("topic_test", 1, 1L, "key".getBytes, "value".getBytes),
-      mock(classOf[PrintStream]))
-
-    formatter.close()
-  }
-}
diff --git a/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala b/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala
deleted file mode 100644
index 12bbc9481107b..0000000000000
--- a/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala
+++ /dev/null
@@ -1,234 +0,0 @@
-/**
-  * Licensed to the Apache Software Foundation (ASF) under one or more
-  * contributor license agreements.  See the NOTICE file distributed with
-  * this work for additional information regarding copyright ownership.
-  * The ASF licenses this file to You under the Apache License, Version 2.0
-  * (the "License"); you may not use this file except in compliance with
-  * the License.  You may obtain a copy of the License at
-  *
-  * http://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software
-  * distributed under the License is distributed on an "AS IS" BASIS,
-  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  * See the License for the specific language governing permissions and
-  * limitations under the License.
-  */
-
-package kafka.tools
-
-import java.io.{ByteArrayOutputStream, Closeable, PrintStream}
-import java.nio.charset.StandardCharsets
-import java.util
-
-import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.header.Header
-import org.apache.kafka.common.header.internals.{RecordHeader, RecordHeaders}
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.serialization.Deserializer
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.{Arguments, MethodSource}
-
-import java.util.Optional
-import scala.jdk.CollectionConverters._
-
-class DefaultMessageFormatterTest {
-  import DefaultMessageFormatterTest._
-
-  @ParameterizedTest
-  @MethodSource(Array("parameters"))
-  def testWriteRecord(name: String, record: ConsumerRecord[Array[Byte], Array[Byte]], properties: Map[String, String], expected: String): Unit = {
-    withResource(new ByteArrayOutputStream()) { baos =>
-      withResource(new PrintStream(baos)) { ps =>
-        val formatter = buildFormatter(properties)
-        formatter.writeTo(record, ps)
-        val actual = new String(baos.toByteArray(), StandardCharsets.UTF_8)
-        assertEquals(expected, actual)
-
-      }
-    }
-  }
-}
-
-object DefaultMessageFormatterTest {
-  def parameters: java.util.stream.Stream[Arguments] = {
-    Seq(
-      Arguments.of(
-        "print nothing",
-        consumerRecord(),
-        Map("print.value" -> "false"),
-        ""),
-      Arguments.of(
-        "print key",
-        consumerRecord(),
-        Map("print.key" -> "true",
-          "print.value" -> "false"),
-        "someKey\n"),
-      Arguments.of(
-        "print value",
-        consumerRecord(),
-        Map(),
-        "someValue\n"),
-      Arguments.of(
-        "print empty timestamp",
-        consumerRecord(timestampType = TimestampType.NO_TIMESTAMP_TYPE),
-        Map("print.timestamp" -> "true",
-            "print.value" -> "false"),
-        "NO_TIMESTAMP\n"),
-      Arguments.of(
-        "print log append time timestamp",
-        consumerRecord(timestampType = TimestampType.LOG_APPEND_TIME),
-        Map("print.timestamp" -> "true",
-            "print.value" -> "false"),
-        "LogAppendTime:1234\n"),
-      Arguments.of(
-        "print create time timestamp",
-        consumerRecord(timestampType = TimestampType.CREATE_TIME),
-        Map("print.timestamp" -> "true",
-            "print.value" -> "false"),
-        "CreateTime:1234\n"),
-      Arguments.of(
-        "print partition",
-        consumerRecord(),
-        Map("print.partition" -> "true",
-            "print.value" -> "false"),
-        "Partition:9\n"),
-      Arguments.of(
-        "print offset",
-        consumerRecord(),
-        Map("print.offset" -> "true",
-            "print.value" -> "false"),
-        "Offset:9876\n"),
-      Arguments.of(
-        "print headers",
-        consumerRecord(),
-        Map("print.headers" -> "true",
-            "print.value" -> "false"),
-        "h1:v1,h2:v2\n"),
-      Arguments.of(
-        "print empty headers",
-        consumerRecord(headers = Nil),
-        Map("print.headers" -> "true",
-            "print.value" -> "false"),
-        "NO_HEADERS\n"),
-      Arguments.of(
-        "print all possible fields with default delimiters",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.timestamp" -> "true",
-            "print.partition" -> "true",
-            "print.offset" -> "true",
-            "print.headers" -> "true",
-            "print.value" -> "true"),
-        "CreateTime:1234\tPartition:9\tOffset:9876\th1:v1,h2:v2\tsomeKey\tsomeValue\n"),
-      Arguments.of(
-        "print all possible fields with custom delimiters",
-        consumerRecord(),
-        Map("key.separator" -> "|",
-            "line.separator" -> "^",
-            "headers.separator" -> "#",
-            "print.key" -> "true",
-            "print.timestamp" -> "true",
-            "print.partition" -> "true",
-            "print.offset" -> "true",
-            "print.headers" -> "true",
-            "print.value" -> "true"),
-        "CreateTime:1234|Partition:9|Offset:9876|h1:v1#h2:v2|someKey|someValue^"),
-      Arguments.of(
-        "print key with custom deserializer",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.headers" -> "true",
-            "print.value" -> "true",
-            "key.deserializer" -> "kafka.tools.UpperCaseDeserializer"),
-        "h1:v1,h2:v2\tSOMEKEY\tsomeValue\n"),
-      Arguments.of(
-        "print value with custom deserializer",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.headers" -> "true",
-            "print.value" -> "true",
-            "value.deserializer" -> "kafka.tools.UpperCaseDeserializer"),
-        "h1:v1,h2:v2\tsomeKey\tSOMEVALUE\n"),
-      Arguments.of(
-        "print headers with custom deserializer",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.headers" -> "true",
-            "print.value" -> "true",
-            "headers.deserializer" -> "kafka.tools.UpperCaseDeserializer"),
-        "h1:V1,h2:V2\tsomeKey\tsomeValue\n"),
-      Arguments.of(
-        "print key and value",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.value" -> "true"),
-        "someKey\tsomeValue\n"),
-      Arguments.of(
-        "print fields in the beginning, middle and the end",
-        consumerRecord(),
-        Map("print.key" -> "true",
-            "print.value" -> "true",
-            "print.partition" -> "true"),
-        "Partition:9\tsomeKey\tsomeValue\n"),
-      Arguments.of(
-        "null value without custom null literal",
-        consumerRecord(value = null),
-        Map("print.key" -> "true"),
-        "someKey\tnull\n"),
-      Arguments.of(
-        "null value with custom null literal",
-        consumerRecord(value = null),
-        Map("print.key" -> "true",
-            "null.literal" -> "NULL"),
-        "someKey\tNULL\n"),
-    ).asJava.stream()
-  }
-
-  private def buildFormatter(propsToSet: Map[String, String]): DefaultMessageFormatter = {
-    val formatter = new DefaultMessageFormatter()
-    formatter.configure(propsToSet.asJava)
-    formatter
-  }
-
-
-  private def header(key: String, value: String) = {
-    new RecordHeader(key, value.getBytes(StandardCharsets.UTF_8))
-  }
-
-  private def consumerRecord(key: String = "someKey",
-                             value: String = "someValue",
-                             headers: Iterable[Header] = Seq(header("h1", "v1"), header("h2", "v2")),
-                             partition: Int = 9,
-                             offset: Long = 9876,
-                             timestamp: Long = 1234,
-                             timestampType: TimestampType = TimestampType.CREATE_TIME) = {
-    new ConsumerRecord[Array[Byte], Array[Byte]](
-      "someTopic",
-      partition,
-      offset,
-      timestamp,
-      timestampType,
-      0,
-      0,
-      if (key == null) null else key.getBytes(StandardCharsets.UTF_8),
-      if (value == null) null else value.getBytes(StandardCharsets.UTF_8),
-      new RecordHeaders(headers.asJava),
-      Optional.empty[Integer])
-  }
-
-  private def withResource[Resource <: Closeable, Result](resource: Resource)(handler: Resource => Result): Result = {
-    try {
-      handler(resource)
-    } finally {
-      resource.close()
-    }
-  }
-}
-
-class UpperCaseDeserializer extends Deserializer[String] {
-  override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = {}
-  override def deserialize(topic: String, data: Array[Byte]): String = new String(data, StandardCharsets.UTF_8).toUpperCase
-  override def close(): Unit = {}
-}
diff --git a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala b/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
deleted file mode 100644
index 054cede8e62e7..0000000000000
--- a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
+++ /dev/null
@@ -1,676 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.io.{ByteArrayOutputStream, PrintStream}
-import java.util.{HashMap, Optional, Map => JMap}
-import java.time.Duration
-import kafka.tools.ConsoleConsumer.ConsumerWrapper
-import kafka.utils.{Exit, TestUtils}
-import org.apache.kafka.clients.consumer.{ConsumerRecord, MockConsumer, OffsetResetStrategy}
-import org.apache.kafka.common.{MessageFormatter, TopicPartition}
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.clients.consumer.ConsumerConfig
-import org.apache.kafka.test.MockDeserializer
-import org.mockito.Mockito._
-import org.mockito.ArgumentMatchers
-import ArgumentMatchers._
-import org.apache.kafka.clients.consumer.Consumer
-import org.apache.kafka.clients.consumer.ConsumerRecords
-import org.apache.kafka.common.errors.TimeoutException
-import org.apache.kafka.common.header.internals.RecordHeaders
-import org.apache.kafka.server.util.MockTime
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.{BeforeEach, Test}
-
-import scala.jdk.CollectionConverters._
-
-class ConsoleConsumerTest {
-
-  @BeforeEach
-  def setup(): Unit = {
-    ConsoleConsumer.messageCount = 0
-  }
-
-  @Test
-  def shouldThrowTimeoutExceptionWhenTimeoutIsReached(): Unit = {
-    val topic = "test"
-    val time = new MockTime
-    val timeoutMs = 1000
-
-    val mockConsumer = mock(classOf[Consumer[Array[Byte], Array[Byte]]])
-
-    when(mockConsumer.poll(Duration.ofMillis(timeoutMs))).thenAnswer { _ =>
-      time.sleep(timeoutMs / 2 + 1)
-      ConsumerRecords.EMPTY
-    }
-
-    val consumer = new ConsumerWrapper(
-      topic = Some(topic),
-      partitionId = None,
-      offset = None,
-      includedTopics = None,
-      consumer = mockConsumer,
-      timeoutMs = timeoutMs,
-      time = time
-    )
-
-    assertThrows(classOf[TimeoutException], () => consumer.receive())
-  }
-
-  @Test
-  def shouldResetUnConsumedOffsetsBeforeExit(): Unit = {
-    val topic = "test"
-    val maxMessages: Int = 123
-    val totalMessages: Int = 700
-    val startOffset: java.lang.Long = 0L
-
-    val mockConsumer = new MockConsumer[Array[Byte], Array[Byte]](OffsetResetStrategy.EARLIEST)
-    val tp1 = new TopicPartition(topic, 0)
-    val tp2 = new TopicPartition(topic, 1)
-
-    val consumer = new ConsumerWrapper(Some(topic), None, None, None, mockConsumer)
-
-    mockConsumer.rebalance(List(tp1, tp2).asJava)
-    mockConsumer.updateBeginningOffsets(Map(tp1 -> startOffset, tp2 -> startOffset).asJava)
-
-    0 until totalMessages foreach { i =>
-      // add all records, each partition should have half of `totalMessages`
-      mockConsumer.addRecord(new ConsumerRecord[Array[Byte], Array[Byte]](topic, i % 2, i / 2, "key".getBytes, "value".getBytes))
-    }
-
-    val formatter = mock(classOf[MessageFormatter])
-
-    ConsoleConsumer.process(maxMessages, formatter, consumer, System.out, skipMessageOnError = false)
-    assertEquals(totalMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2))
-
-    consumer.resetUnconsumedOffsets()
-    assertEquals(maxMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2))
-
-    verify(formatter, times(maxMessages)).writeTo(any(), any())
-  }
-
-  @Test
-  def shouldLimitReadsToMaxMessageLimit(): Unit = {
-    val consumer = mock(classOf[ConsumerWrapper])
-    val formatter = mock(classOf[MessageFormatter])
-    val record = new ConsumerRecord("foo", 1, 1, Array[Byte](), Array[Byte]())
-
-    val messageLimit: Int = 10
-    when(consumer.receive()).thenReturn(record)
-
-    ConsoleConsumer.process(messageLimit, formatter, consumer, System.out, true)
-
-    verify(consumer, times(messageLimit)).receive()
-    verify(formatter, times(messageLimit)).writeTo(any(), any())
-
-    consumer.cleanup()
-  }
-
-  @Test
-  def shouldStopWhenOutputCheckErrorFails(): Unit = {
-    val consumer = mock(classOf[ConsumerWrapper])
-    val formatter = mock(classOf[MessageFormatter])
-    val printStream = mock(classOf[PrintStream])
-
-    val record = new ConsumerRecord("foo", 1, 1, Array[Byte](), Array[Byte]())
-
-    when(consumer.receive()).thenReturn(record)
-    //Simulate an error on System.out after the first record has been printed
-    when(printStream.checkError()).thenReturn(true)
-
-    ConsoleConsumer.process(-1, formatter, consumer, printStream, true)
-
-    verify(formatter).writeTo(any(), ArgumentMatchers.eq(printStream))
-    verify(consumer).receive()
-    verify(printStream).checkError()
-
-    consumer.cleanup()
-  }
-
-  @Test
-  def shouldParseValidConsumerValidConfig(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(true, config.fromBeginning)
-  }
-
-  @Test
-  def shouldParseIncludeArgument(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--include", "includeTest*",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("includeTest*", config.includedTopicsArg)
-    assertEquals(true, config.fromBeginning)
-  }
-
-  @Test
-  def shouldParseWhitelistArgument(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--whitelist", "whitelistTest*",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("whitelistTest*", config.includedTopicsArg)
-    assertEquals(true, config.fromBeginning)
-  }
-
-  @Test
-  def shouldIgnoreWhitelistArgumentIfIncludeSpecified(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--include", "includeTest*",
-      "--whitelist", "whitelistTest*",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("includeTest*", config.includedTopicsArg)
-    assertEquals(true, config.fromBeginning)
-  }
-
-  @Test
-  def shouldParseValidSimpleConsumerValidConfigWithNumericOffset(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--partition", "0",
-      "--offset", "3")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(0, config.partitionArg.get)
-    assertEquals(3, config.offsetArg)
-    assertEquals(false, config.fromBeginning)
-
-  }
-
-  @Test
-  def shouldExitOnUnrecognizedNewConsumerOption(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-
-    //Given
-    val args: Array[String] = Array(
-      "--new-consumer",
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--from-beginning")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def shouldParseValidSimpleConsumerValidConfigWithStringOffset(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--partition", "0",
-      "--offset", "LatEst",
-      "--property", "print.value=false")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(0, config.partitionArg.get)
-    assertEquals(-1, config.offsetArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals(false, config.formatter.asInstanceOf[DefaultMessageFormatter].printValue)
-  }
-
-  @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetLatest(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=latest")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetEarliest(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=earliest")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=earliest",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(true, config.fromBeginning)
-    assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldParseValidConsumerConfigWithNoOffsetReset(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginning(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=latest",
-      "--from-beginning")
-    try {
-      val config = new ConsoleConsumer.ConsumerConfig(args)
-      assertThrows(classOf[IllegalArgumentException], () => ConsoleConsumer.consumerProps(config))
-    }
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def shouldParseConfigsFromFile(): Unit = {
-    val propsFile = TestUtils.tempPropertiesFile(Map("request.timeout.ms" -> "1000", "group.id" -> "group1"))
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer.config", propsFile.getAbsolutePath
-    )
-
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    assertEquals("1000", config.consumerProps.getProperty("request.timeout.ms"))
-    assertEquals("group1", config.consumerProps.getProperty("group.id"))
-  }
-
-  @Test
-  def groupIdsProvidedInDifferentPlacesMustMatch(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-
-    // different in all three places
-    var propsFile = TestUtils.tempPropertiesFile(Map("group.id" -> "group-from-file"))
-    var args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "group-from-arguments",
-      "--consumer-property", "group.id=group-from-properties",
-      "--consumer.config", propsFile.getAbsolutePath
-    )
-
-    assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-
-    // the same in all three places
-    propsFile = TestUtils.tempPropertiesFile(Map("group.id" -> "test-group"))
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "test-group",
-      "--consumer-property", "group.id=test-group",
-      "--consumer.config", propsFile.getAbsolutePath
-    )
-
-    var config = new ConsoleConsumer.ConsumerConfig(args)
-    var props = ConsoleConsumer.consumerProps(config)
-    assertEquals("test-group", props.getProperty("group.id"))
-
-    // different via --consumer-property and --consumer.config
-    propsFile = TestUtils.tempPropertiesFile(Map("group.id" -> "group-from-file"))
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--consumer-property", "group.id=group-from-properties",
-      "--consumer.config", propsFile.getAbsolutePath
-    )
-
-    assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-
-    // different via --consumer-property and --group
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "group-from-arguments",
-      "--consumer-property", "group.id=group-from-properties"
-    )
-
-    assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-
-    // different via --group and --consumer.config
-    propsFile = TestUtils.tempPropertiesFile(Map("group.id" -> "group-from-file"))
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "group-from-arguments",
-      "--consumer.config", propsFile.getAbsolutePath
-    )
-    assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-
-    // via --group only
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "group-from-arguments"
-    )
-
-    config = new ConsoleConsumer.ConsumerConfig(args)
-    props = ConsoleConsumer.consumerProps(config)
-    assertEquals("group-from-arguments", props.getProperty("group.id"))
-
-    Exit.resetExitProcedure()
-  }
-
-  @Test
-  def testCustomPropertyShouldBePassedToConfigureMethod(): Unit = {
-    val args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--property", "print.key=true",
-      "--property", "key.deserializer=org.apache.kafka.test.MockDeserializer",
-      "--property", "key.deserializer.my-props=abc"
-    )
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    assertTrue(config.formatter.isInstanceOf[DefaultMessageFormatter])
-    assertTrue(config.formatterArgs.containsKey("key.deserializer.my-props"))
-    val formatter = config.formatter.asInstanceOf[DefaultMessageFormatter]
-    assertTrue(formatter.keyDeserializer.get.isInstanceOf[MockDeserializer])
-    val keyDeserializer = formatter.keyDeserializer.get.asInstanceOf[MockDeserializer]
-    assertEquals(1, keyDeserializer.configs.size)
-    assertEquals("abc", keyDeserializer.configs.get("my-props"))
-    assertTrue(keyDeserializer.isKey)
-  }
-
-  @Test
-  def testCustomConfigShouldBePassedToConfigureMethod(): Unit = {
-    val propsFile = TestUtils.tempPropertiesFile(Map("key.deserializer.my-props" -> "abc", "print.key" -> "false"))
-    val args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--property", "print.key=true",
-      "--property", "key.deserializer=org.apache.kafka.test.MockDeserializer",
-      "--formatter-config", propsFile.getAbsolutePath
-    )
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    assertTrue(config.formatter.isInstanceOf[DefaultMessageFormatter])
-    assertTrue(config.formatterArgs.containsKey("key.deserializer.my-props"))
-    val formatter = config.formatter.asInstanceOf[DefaultMessageFormatter]
-    assertTrue(formatter.keyDeserializer.get.isInstanceOf[MockDeserializer])
-    val keyDeserializer = formatter.keyDeserializer.get.asInstanceOf[MockDeserializer]
-    assertEquals(1, keyDeserializer.configs.size)
-    assertEquals("abc", keyDeserializer.configs.get("my-props"))
-    assertTrue(keyDeserializer.isKey)
-  }
-
-  @Test
-  def shouldParseGroupIdFromBeginningGivenTogether(): Unit = {
-    // Start from earliest
-    var args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "test-group",
-      "--from-beginning")
-
-    var config = new ConsoleConsumer.ConsumerConfig(args)
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(-2, config.offsetArg)
-    assertEquals(true, config.fromBeginning)
-
-    // Start from latest
-    args = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "test-group"
-    )
-
-    config = new ConsoleConsumer.ConsumerConfig(args)
-    assertEquals("localhost:9092", config.bootstrapServer)
-    assertEquals("test", config.topicArg)
-    assertEquals(-1, config.offsetArg)
-    assertEquals(false, config.fromBeginning)
-  }
-
-  @Test
-  def shouldExitOnGroupIdAndPartitionGivenTogether(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--group", "test-group",
-      "--partition", "0")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def shouldExitOnOffsetWithoutPartition(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--offset", "10")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def testDefaultMessageFormatter(): Unit = {
-    val record = new ConsumerRecord("topic", 0, 123, "key".getBytes, "value".getBytes)
-    val formatter = new DefaultMessageFormatter()
-    val configs: JMap[String, String] = new HashMap()
-
-    formatter.configure(configs)
-    var out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("value\n", out.toString)
-
-    configs.put("print.key", "true")
-    formatter.configure(configs)
-    out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("key\tvalue\n", out.toString)
-
-    configs.put("print.partition", "true")
-    formatter.configure(configs)
-    out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("Partition:0\tkey\tvalue\n", out.toString)
-
-    configs.put("print.timestamp", "true")
-    formatter.configure(configs)
-    out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("NO_TIMESTAMP\tPartition:0\tkey\tvalue\n", out.toString)
-
-    configs.put("print.offset", "true")
-    formatter.configure(configs)
-    out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("NO_TIMESTAMP\tPartition:0\tOffset:123\tkey\tvalue\n", out.toString)
-
-    out = new ByteArrayOutputStream()
-    val record2 = new ConsumerRecord("topic", 0, 123, 123L, TimestampType.CREATE_TIME, -1, -1, "key".getBytes, "value".getBytes,
-      new RecordHeaders(), Optional.empty[Integer])
-    formatter.writeTo(record2, new PrintStream(out))
-    assertEquals("CreateTime:123\tPartition:0\tOffset:123\tkey\tvalue\n", out.toString)
-    formatter.close()
-  }
-
-  @Test
-  def testNoOpMessageFormatter(): Unit = {
-    val record = new ConsumerRecord("topic", 0, 123, "key".getBytes, "value".getBytes)
-    val formatter = new NoOpMessageFormatter()
-
-    formatter.configure(new HashMap())
-    val out = new ByteArrayOutputStream()
-    formatter.writeTo(record, new PrintStream(out))
-    assertEquals("", out.toString)
-  }
-
-  @Test
-  def shouldExitIfNoTopicOrFilterSpecified(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def shouldExitIfTopicAndIncludeSpecified(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--include", "includeTest*")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def shouldExitIfTopicAndWhitelistSpecified(): Unit = {
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--whitelist", "whitelistTest*")
-
-    try assertThrows(classOf[IllegalArgumentException], () => new ConsoleConsumer.ConsumerConfig(args))
-    finally Exit.resetExitProcedure()
-  }
-
-  @Test
-  def testClientIdOverride(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--from-beginning",
-      "--consumer-property", "client.id=consumer-1")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("consumer-1", consumerProperties.getProperty(ConsumerConfig.CLIENT_ID_CONFIG))
-  }
-
-  @Test
-  def testDefaultClientId(): Unit = {
-    //Given
-    val args: Array[String] = Array(
-      "--bootstrap-server", "localhost:9092",
-      "--topic", "test",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.consumerProps(config)
-
-    //Then
-    assertEquals("console-consumer", consumerProperties.getProperty(ConsumerConfig.CLIENT_ID_CONFIG))
-  }
-}
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java
index 03877c6859563..1f3fc98b77828 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.streams.integration;
 
-import kafka.tools.ConsoleConsumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.IntegerDeserializer;
@@ -62,6 +61,8 @@
 import org.apache.kafka.streams.state.ReadOnlySessionStore;
 import org.apache.kafka.test.MockMapper;
 import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.consumer.ConsoleConsumer;
+import org.apache.kafka.tools.consumer.ConsoleConsumerOptions;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
@@ -1117,7 +1118,7 @@ private <K, V> String readWindowedKeyedMessagesViaConsoleConsumer(final Deserial
                                                                       final Deserializer<V> valueDeserializer,
                                                                       final Class innerClass,
                                                                       final int numMessages,
-                                                                      final boolean printTimestamp) {
+                                                                      final boolean printTimestamp) throws Exception {
         final ByteArrayOutputStream newConsole = new ByteArrayOutputStream();
         final PrintStream originalStream = System.out;
         try (final PrintStream newStream = new PrintStream(newConsole)) {
@@ -1139,8 +1140,7 @@ private <K, V> String readWindowedKeyedMessagesViaConsoleConsumer(final Deserial
                 "--property", "key.deserializer.window.size.ms=500",
             };
 
-            ConsoleConsumer.messageCount_$eq(0); //reset the message count
-            ConsoleConsumer.run(new ConsoleConsumer.ConsumerConfig(args));
+            ConsoleConsumer.run(new ConsoleConsumerOptions(args));
             newStream.flush();
             System.setOut(originalStream);
             return newConsole.toString();
diff --git a/tests/kafkatest/services/console_consumer.py b/tests/kafkatest/services/console_consumer.py
index eb5d3d9f85dc2..fb87f20df1993 100644
--- a/tests/kafkatest/services/console_consumer.py
+++ b/tests/kafkatest/services/console_consumer.py
@@ -21,7 +21,7 @@
 
 from kafkatest.directory_layout.kafka_path import KafkaPathResolverMixin
 from kafkatest.services.monitor.jmx import JmxMixin, JmxTool
-from kafkatest.version import DEV_BRANCH, LATEST_0_8_2, LATEST_0_9, LATEST_0_10_0, V_0_10_0_0, V_0_11_0_0, V_2_0_0
+from kafkatest.version import DEV_BRANCH, LATEST_0_8_2, LATEST_0_9, LATEST_0_10_0, V_0_10_0_0, V_0_11_0_0, V_2_0_0, LATEST_3_7
 from kafkatest.services.kafka.util import fix_opts_for_new_jvm
 
 """
@@ -210,7 +210,10 @@ def start_cmd(self, node):
 
         # LoggingMessageFormatter was introduced after 0.9
         if node.version > LATEST_0_9:
-            cmd += " --formatter kafka.tools.LoggingMessageFormatter"
+            if node.version > LATEST_3_7:
+                cmd += " --formatter org.apache.kafka.tools.consumer.LoggingMessageFormatter"
+            else:
+                cmd += " --formatter kafka.tools.LoggingMessageFormatter"
 
         if self.enable_systest_events:
             # enable systest events is only available in 0.10.0 and later
diff --git a/tests/kafkatest/version.py b/tests/kafkatest/version.py
index 6d0b1ae329f75..3e3da017bc4a0 100644
--- a/tests/kafkatest/version.py
+++ b/tests/kafkatest/version.py
@@ -263,3 +263,7 @@ def get_version(node=None):
 V_3_6_0 = KafkaVersion("3.6.0")
 V_3_6_1 = KafkaVersion("3.6.1")
 LATEST_3_6 = V_3_6_1
+
+# 3.7.x version
+V_3_7_0 = KafkaVersion("3.7.0")
+LATEST_3_7 = V_3_7_0
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
new file mode 100644
index 0000000000000..f84fb88c23f5f
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
@@ -0,0 +1,234 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import java.io.PrintStream;
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
+import java.util.concurrent.CountDownLatch;
+import java.util.regex.Pattern;
+import java.util.Collections;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.requests.ListOffsetsRequest;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Consumer that dumps messages to standard out.
+ */
+public class ConsoleConsumer {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ConsoleConsumer.class);
+    private static final CountDownLatch SHUTDOWN_LATCH = new CountDownLatch(1);
+
+    static int messageCount = 0;
+
+    public static void main(String[] args) throws Exception {
+        ConsoleConsumerOptions opts = new ConsoleConsumerOptions(args);
+        try {
+            run(opts);
+        } catch (AuthenticationException ae) {
+            LOG.error("Authentication failed: terminating consumer process", ae);
+            Exit.exit(1);
+        } catch (Throwable t) {
+            LOG.error("Unknown error when running consumer: ", t);
+            Exit.exit(1);
+        }
+    }
+
+    public static void run(ConsoleConsumerOptions opts) {
+        messageCount = 0;
+        long timeoutMs = opts.timeoutMs() >= 0 ? opts.timeoutMs() : Long.MAX_VALUE;
+        Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(opts.consumerProps(), new ByteArrayDeserializer(), new ByteArrayDeserializer());
+        ConsumerWrapper consumerWrapper = opts.partitionArg().isPresent()
+            ? new ConsumerWrapper(Optional.of(opts.topicArg()), opts.partitionArg(), OptionalLong.of(opts.offsetArg()), Optional.empty(), consumer, timeoutMs)
+            : new ConsumerWrapper(Optional.of(opts.topicArg()), OptionalInt.empty(), OptionalLong.empty(), Optional.ofNullable(opts.includedTopicsArg()), consumer, timeoutMs);
+
+        addShutdownHook(consumerWrapper, opts);
+
+        try {
+            process(opts.maxMessages(), opts.formatter(), consumerWrapper, System.out, opts.skipMessageOnError());
+        } finally {
+            consumerWrapper.cleanup();
+            opts.formatter().close();
+            reportRecordCount();
+
+            SHUTDOWN_LATCH.countDown();
+        }
+    }
+
+    static void addShutdownHook(ConsumerWrapper consumer, ConsoleConsumerOptions conf) {
+        Exit.addShutdownHook("consumer-shutdown-hook", () -> {
+            try {
+                consumer.wakeup();
+                SHUTDOWN_LATCH.await();
+            } catch (Throwable t) {
+                LOG.error("Exception while running shutdown hook: ", t);
+            }
+            if (conf.enableSystestEventsLogging()) {
+                System.out.println("shutdown_complete");
+            }
+        });
+    }
+
+    static void process(int maxMessages, MessageFormatter formatter, ConsumerWrapper consumer, PrintStream output, boolean skipMessageOnError) {
+        while (messageCount < maxMessages || maxMessages == -1) {
+            ConsumerRecord<byte[], byte[]> msg;
+            try {
+                msg = consumer.receive();
+            } catch (WakeupException we) {
+                LOG.trace("Caught WakeupException because consumer is shutdown, ignore and terminate.");
+                // Consumer will be closed
+                return;
+            } catch (Throwable t) {
+                LOG.error("Error processing message, terminating consumer process: ", t);
+                // Consumer will be closed
+                return;
+            }
+            messageCount += 1;
+            try {
+                formatter.writeTo(new ConsumerRecord<>(msg.topic(), msg.partition(), msg.offset(), msg.timestamp(), msg.timestampType(),
+                    0, 0, msg.key(), msg.value(), msg.headers(), Optional.empty()), output);
+            } catch (Throwable t) {
+                if (skipMessageOnError) {
+                    LOG.error("Error processing message, skipping this message: ", t);
+                } else {
+                    // Consumer will be closed
+                    throw t;
+                }
+            }
+            if (checkErr(output)) {
+                // Consumer will be closed
+                return;
+            }
+        }
+    }
+
+    static void reportRecordCount() {
+        System.err.println("Processed a total of " + messageCount + " messages");
+    }
+
+    static boolean checkErr(PrintStream output) {
+        boolean gotError = output.checkError();
+        if (gotError) {
+            // This means no one is listening to our output stream anymore, time to shut down
+            System.err.println("Unable to write to standard out, closing consumer.");
+        }
+        return gotError;
+    }
+
+    public static class ConsumerWrapper {
+        final Optional<String> topic;
+        final OptionalInt partitionId;
+        final OptionalLong offset;
+        final Optional<String> includedTopics;
+        final Consumer<byte[], byte[]> consumer;
+        final long timeoutMs;
+        final Time time = Time.SYSTEM;
+
+        Iterator<ConsumerRecord<byte[], byte[]>> recordIter = Collections.emptyIterator();
+
+        public ConsumerWrapper(Optional<String> topic,
+                               OptionalInt partitionId,
+                               OptionalLong offset,
+                               Optional<String> includedTopics,
+                               Consumer<byte[], byte[]> consumer,
+                               long timeoutMs) {
+            this.topic = topic;
+            this.partitionId = partitionId;
+            this.offset = offset;
+            this.includedTopics = includedTopics;
+            this.consumer = consumer;
+            this.timeoutMs = timeoutMs;
+
+            if (topic.isPresent() && partitionId.isPresent() && offset.isPresent() && !includedTopics.isPresent()) {
+                seek(topic.get(), partitionId.getAsInt(), offset.getAsLong());
+            } else if (topic.isPresent() && partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
+                // default to latest if no offset is provided
+                seek(topic.get(), partitionId.getAsInt(), ListOffsetsRequest.LATEST_TIMESTAMP);
+            } else if (topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
+                consumer.subscribe(Collections.singletonList(topic.get()));
+            } else if (!topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && includedTopics.isPresent()) {
+                consumer.subscribe(Pattern.compile(includedTopics.get()));
+            } else {
+                throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
+                        "Exactly one of 'topic' or 'include' must be provided. " +
+                        "If 'topic' is provided, an optional 'partition' may also be provided. " +
+                        "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.");
+            }
+        }
+
+        private void seek(String topic, int partitionId, long offset) {
+            TopicPartition topicPartition = new TopicPartition(topic, partitionId);
+            consumer.assign(Collections.singletonList(topicPartition));
+            if (offset == ListOffsetsRequest.EARLIEST_TIMESTAMP) {
+                consumer.seekToBeginning(Collections.singletonList(topicPartition));
+            } else if (offset == ListOffsetsRequest.LATEST_TIMESTAMP) {
+                consumer.seekToEnd(Collections.singletonList(topicPartition));
+            } else {
+                consumer.seek(topicPartition, offset);
+            }
+        }
+
+        void resetUnconsumedOffsets() {
+            Map<TopicPartition, Long> smallestUnconsumedOffsets = new HashMap<>();
+            while (recordIter.hasNext()) {
+                ConsumerRecord<byte[], byte[]> record = recordIter.next();
+                TopicPartition tp = new TopicPartition(record.topic(), record.partition());
+                // avoid auto-committing offsets which haven't been consumed
+                smallestUnconsumedOffsets.putIfAbsent(tp, record.offset());
+            }
+            smallestUnconsumedOffsets.forEach(consumer::seek);
+        }
+
+        ConsumerRecord<byte[], byte[]> receive() {
+            long startTimeMs = time.milliseconds();
+            while (!recordIter.hasNext()) {
+                recordIter = consumer.poll(Duration.ofMillis(timeoutMs)).iterator();
+                if (!recordIter.hasNext() && (time.milliseconds() - startTimeMs > timeoutMs)) {
+                    throw new TimeoutException();
+                }
+            }
+            return recordIter.next();
+        }
+
+        void wakeup() {
+            this.consumer.wakeup();
+        }
+
+        void cleanup() {
+            resetUnconsumedOffsets();
+            this.consumer.close();
+        }
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
new file mode 100644
index 0000000000000..a713afb2bf22c
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
@@ -0,0 +1,414 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import joptsimple.OptionException;
+import joptsimple.OptionSpec;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.requests.ListOffsetsRequest;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.util.CommandDefaultOptions;
+import org.apache.kafka.server.util.CommandLineUtils;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.OptionalInt;
+import java.util.Properties;
+import java.util.Random;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public final class ConsoleConsumerOptions extends CommandDefaultOptions {
+
+    private static final Random RANDOM = new Random();
+
+    private final OptionSpec<String> topicOpt;
+    private final OptionSpec<String> whitelistOpt;
+    private final OptionSpec<String> includeOpt;
+    private final OptionSpec<Integer> partitionIdOpt;
+    private final OptionSpec<String> offsetOpt;
+    private final OptionSpec<String> messageFormatterOpt;
+    private final OptionSpec<String> messageFormatterArgOpt;
+    private final OptionSpec<String> messageFormatterConfigOpt;
+    private final OptionSpec<?> resetBeginningOpt;
+    private final OptionSpec<Integer> maxMessagesOpt;
+    private final OptionSpec<Integer> timeoutMsOpt;
+    private final OptionSpec<?> skipMessageOnErrorOpt;
+    private final OptionSpec<String> bootstrapServerOpt;
+    private final OptionSpec<String> keyDeserializerOpt;
+    private final OptionSpec<String> valueDeserializerOpt;
+    private final OptionSpec<?> enableSystestEventsLoggingOpt;
+    private final OptionSpec<String> isolationLevelOpt;
+    private final OptionSpec<String> groupIdOpt;
+
+    private final Properties consumerProps;
+    private final long offset;
+    private final MessageFormatter formatter;
+
+    public ConsoleConsumerOptions(String[] args) throws IOException {
+        super(args);
+        topicOpt = parser.accepts("topic", "The topic to consume on.")
+                .withRequiredArg()
+                .describedAs("topic")
+                .ofType(String.class);
+        whitelistOpt = parser.accepts("whitelist",
+                        "DEPRECATED, use --include instead; ignored if --include specified. Regular expression specifying list of topics to include for consumption.")
+                .withRequiredArg()
+                .describedAs("Java regex (String)")
+                .ofType(String.class);
+        includeOpt = parser.accepts("include",
+                        "Regular expression specifying list of topics to include for consumption.")
+                .withRequiredArg()
+                .describedAs("Java regex (String)")
+                .ofType(String.class);
+        partitionIdOpt = parser.accepts("partition",
+                        "The partition to consume from. Consumption starts from the end of the partition unless '--offset' is specified.")
+                .withRequiredArg()
+                .describedAs("partition")
+                .ofType(Integer.class);
+        offsetOpt = parser.accepts("offset", "The offset to consume from (a non-negative number), or 'earliest' which means from beginning, or 'latest' which means from end")
+                .withRequiredArg()
+                .describedAs("consume offset")
+                .ofType(String.class)
+                .defaultsTo("latest");
+        OptionSpec<String> consumerPropertyOpt = parser.accepts("consumer-property", "A mechanism to pass user-defined properties in the form key=value to the consumer.")
+                .withRequiredArg()
+                .describedAs("consumer_prop")
+                .ofType(String.class);
+        OptionSpec<String> consumerConfigOpt = parser.accepts("consumer.config", "Consumer config properties file. Note that " + consumerPropertyOpt + " takes precedence over this config.")
+                .withRequiredArg()
+                .describedAs("config file")
+                .ofType(String.class);
+        messageFormatterOpt = parser.accepts("formatter", "The name of a class to use for formatting kafka messages for display.")
+                .withRequiredArg()
+                .describedAs("class")
+                .ofType(String.class)
+                .defaultsTo(DefaultMessageFormatter.class.getName());
+        messageFormatterArgOpt = parser.accepts("property",
+                        "The properties to initialize the message formatter. Default properties include: \n" +
+                            " print.timestamp=true|false\n" +
+                            " print.key=true|false\n" +
+                            " print.offset=true|false\n" +
+                            " print.partition=true|false\n" +
+                            " print.headers=true|false\n" +
+                            " print.value=true|false\n" +
+                            " key.separator=<key.separator>\n" +
+                            " line.separator=<line.separator>\n" +
+                            " headers.separator=<line.separator>\n" +
+                            " null.literal=<null.literal>\n" +
+                            " key.deserializer=<key.deserializer>\n" +
+                            " value.deserializer=<value.deserializer>\n" +
+                            " header.deserializer=<header.deserializer>\n" +
+                            "\nUsers can also pass in customized properties for their formatter; more specifically, users can pass in properties keyed with 'key.deserializer.', 'value.deserializer.' and 'headers.deserializer.' prefixes to configure their deserializers.")
+                .withRequiredArg()
+                .describedAs("prop")
+                .ofType(String.class);
+        messageFormatterConfigOpt = parser.accepts("formatter-config", "Config properties file to initialize the message formatter. Note that " + messageFormatterArgOpt + " takes precedence over this config.")
+                .withRequiredArg()
+                .describedAs("config file")
+                .ofType(String.class);
+        resetBeginningOpt = parser.accepts("from-beginning", "If the consumer does not already have an established offset to consume from, " +
+                "start with the earliest message present in the log rather than the latest message.");
+        maxMessagesOpt = parser.accepts("max-messages", "The maximum number of messages to consume before exiting. If not set, consumption is continual.")
+                .withRequiredArg()
+                .describedAs("num_messages")
+                .ofType(Integer.class);
+        timeoutMsOpt = parser.accepts("timeout-ms", "If specified, exit if no message is available for consumption for the specified interval.")
+                .withRequiredArg()
+                .describedAs("timeout_ms")
+                .ofType(Integer.class);
+        skipMessageOnErrorOpt = parser.accepts("skip-message-on-error", "If there is an error when processing a message, " +
+                "skip it instead of halt.");
+        bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: The server(s) to connect to.")
+                .withRequiredArg()
+                .describedAs("server to connect to")
+                .ofType(String.class);
+        keyDeserializerOpt = parser.accepts("key-deserializer")
+                .withRequiredArg()
+                .describedAs("deserializer for key")
+                .ofType(String.class);
+        valueDeserializerOpt = parser.accepts("value-deserializer")
+                .withRequiredArg()
+                .describedAs("deserializer for values")
+                .ofType(String.class);
+        enableSystestEventsLoggingOpt = parser.accepts("enable-systest-events",
+                "Log lifecycle events of the consumer in addition to logging consumed messages. (This is specific for system tests.)");
+        isolationLevelOpt = parser.accepts("isolation-level",
+                        "Set to read_committed in order to filter out transactional messages which are not committed. Set to read_uncommitted " +
+                                "to read all messages.")
+                .withRequiredArg()
+                .ofType(String.class)
+                .defaultsTo("read_uncommitted");
+        groupIdOpt = parser.accepts("group", "The consumer group id of the consumer.")
+                .withRequiredArg()
+                .describedAs("consumer group id")
+                .ofType(String.class);
+
+        try {
+            options = parser.parse(args);
+        } catch (OptionException oe) {
+            CommandLineUtils.printUsageAndExit(parser, oe.getMessage());
+        }
+
+        CommandLineUtils.maybePrintHelpOrVersion(this, "This tool helps to read data from Kafka topics and outputs it to standard output.");
+
+        checkRequiredArgs();
+
+        Properties consumerPropsFromFile = options.has(consumerConfigOpt)
+                ? Utils.loadProps(options.valueOf(consumerConfigOpt))
+                : new Properties();
+        Properties extraConsumerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(consumerPropertyOpt));
+        Set<String> groupIdsProvided = checkConsumerGroup(consumerPropsFromFile, extraConsumerProps);
+        consumerProps = buildConsumerProps(consumerPropsFromFile, extraConsumerProps, groupIdsProvided);
+        offset = parseOffset();
+        formatter = buildFormatter();
+    }
+
+    private void checkRequiredArgs() {
+        List<String> topicOrFilterArgs = new ArrayList<>(Arrays.asList(topicArg(), includedTopicsArg()));
+        topicOrFilterArgs.removeIf(Objects::isNull);
+        // user need to specify value for either --topic or one of the include filters options (--include or --whitelist)
+        if (topicOrFilterArgs.size() != 1) {
+            CommandLineUtils.printUsageAndExit(parser, "Exactly one of --include/--topic is required. " +
+                    (options.has(whitelistOpt) ? "--whitelist is DEPRECATED use --include instead; ignored if --include specified." : ""));
+        }
+
+        if (partitionArg().isPresent()) {
+            if (!options.has(topicOpt)) {
+                CommandLineUtils.printUsageAndExit(parser, "The topic is required when partition is specified.");
+            }
+            if (fromBeginning() && options.has(offsetOpt)) {
+                CommandLineUtils.printUsageAndExit(parser, "Options from-beginning and offset cannot be specified together.");
+            }
+        } else if (options.has(offsetOpt)) {
+            CommandLineUtils.printUsageAndExit(parser, "The partition is required when offset is specified.");
+        }
+
+        CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt);
+    }
+
+    private Set<String> checkConsumerGroup(Properties consumerPropsFromFile, Properties extraConsumerProps) {
+        // if the group id is provided in more than place (through different means) all values must be the same
+        Set<String> groupIdsProvided = new HashSet<>();
+        if (options.has(groupIdOpt)) {
+            groupIdsProvided.add(options.valueOf(groupIdOpt));
+        }
+
+        if (consumerPropsFromFile.containsKey(ConsumerConfig.GROUP_ID_CONFIG)) {
+            groupIdsProvided.add((String) consumerPropsFromFile.get(ConsumerConfig.GROUP_ID_CONFIG));
+        }
+
+        if (extraConsumerProps.containsKey(ConsumerConfig.GROUP_ID_CONFIG)) {
+            groupIdsProvided.add(extraConsumerProps.getProperty(ConsumerConfig.GROUP_ID_CONFIG));
+        }
+        if (groupIdsProvided.size() > 1) {
+            CommandLineUtils.printUsageAndExit(parser, "The group ids provided in different places (directly using '--group', "
+                    + "via '--consumer-property', or via '--consumer.config') do not match. "
+                    + "Detected group ids: "
+                    + groupIdsProvided.stream().map(group -> "'" + group + "'").collect(Collectors.joining(", ")));
+        }
+        if (!groupIdsProvided.isEmpty() && partitionArg().isPresent()) {
+            CommandLineUtils.printUsageAndExit(parser, "Options group and partition cannot be specified together.");
+        }
+        return groupIdsProvided;
+    }
+
+    private Properties buildConsumerProps(Properties consumerPropsFromFile, Properties extraConsumerProps, Set<String> groupIdsProvided) {
+        Properties consumerProps = new Properties();
+        consumerProps.putAll(consumerPropsFromFile);
+        consumerProps.putAll(extraConsumerProps);
+        setAutoOffsetResetValue(consumerProps);
+        consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer());
+        if (consumerProps.getProperty(ConsumerConfig.CLIENT_ID_CONFIG) == null) {
+            consumerProps.put(ConsumerConfig.CLIENT_ID_CONFIG, "console-consumer");
+        }
+        CommandLineUtils.maybeMergeOptions(consumerProps, ConsumerConfig.ISOLATION_LEVEL_CONFIG, options, isolationLevelOpt);
+
+        if (groupIdsProvided.isEmpty()) {
+            consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "console-consumer-" + RANDOM.nextInt(100000));
+            // By default, avoid unnecessary expansion of the coordinator cache since
+            // the auto-generated group and its offsets is not intended to be used again
+            if (!consumerProps.containsKey(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG)) {
+                consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
+            }
+        } else {
+            consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, groupIdsProvided.iterator().next());
+        }
+        return consumerProps;
+    }
+
+    /**
+     * Used to retrieve the correct value for the consumer parameter 'auto.offset.reset'.
+     * Order of priority is:
+     *   1. Explicitly set parameter via --consumer.property command line parameter
+     *   2. Explicit --from-beginning given -> 'earliest'
+     *   3. Default value of 'latest'
+     * In case both --from-beginning and an explicit value are specified an error is thrown if these
+     * are conflicting.
+     */
+    private void setAutoOffsetResetValue(Properties props) {
+        String earliestConfigValue = "earliest";
+        String latestConfigValue = "latest";
+
+        if (props.containsKey(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)) {
+            // auto.offset.reset parameter was specified on the command line
+            String autoResetOption = props.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG);
+            if (fromBeginning() && !earliestConfigValue.equals(autoResetOption)) {
+                // conflicting options - latest und earliest, throw an error
+                System.err.println("Can't simultaneously specify --from-beginning and 'auto.offset.reset=" + autoResetOption + "', " +
+                        "please remove one option");
+                Exit.exit(1);
+            }
+            // nothing to do, checking for valid parameter values happens later and the specified
+            // value was already copied during .putall operation
+        } else {
+            // no explicit value for auto.offset.reset was specified
+            // if --from-beginning was specified use earliest, otherwise default to latest
+            String autoResetOption = fromBeginning() ? earliestConfigValue : latestConfigValue;
+            props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, autoResetOption);
+        }
+    }
+
+    private long parseOffset() {
+        if (options.has(offsetOpt)) {
+            switch (options.valueOf(offsetOpt).toLowerCase(Locale.ROOT)) {
+                case "earliest":
+                    return ListOffsetsRequest.EARLIEST_TIMESTAMP;
+                case "latest":
+                    return ListOffsetsRequest.LATEST_TIMESTAMP;
+                default:
+                    String offsetString = options.valueOf(offsetOpt);
+                    try {
+                        long offset = Long.parseLong(offsetString);
+                        if (offset < 0) {
+                            invalidOffset(offsetString);
+                        }
+                        return offset;
+                    } catch (NumberFormatException nfe) {
+                        invalidOffset(offsetString);
+                    }
+            }
+        } else if (fromBeginning()) {
+            return ListOffsetsRequest.EARLIEST_TIMESTAMP;
+        }
+        return ListOffsetsRequest.LATEST_TIMESTAMP;
+    }
+
+    private void invalidOffset(String offset) {
+        CommandLineUtils.printUsageAndExit(parser, "The provided offset value '" + offset + "' is incorrect. Valid values are " +
+                "'earliest', 'latest', or a non-negative long.");
+    }
+
+    private MessageFormatter buildFormatter() {
+        MessageFormatter formatter = null;
+        try {
+            Class<?> messageFormatterClass = Class.forName(options.valueOf(messageFormatterOpt));
+            formatter = (MessageFormatter) messageFormatterClass.getDeclaredConstructor().newInstance();
+
+            Properties formatterArgs = formatterArgs();
+            Map<String, String> formatterConfigs = new HashMap<>();
+            for (final String name : formatterArgs.stringPropertyNames()) {
+                formatterConfigs.put(name, formatterArgs.getProperty(name));
+            }
+
+            formatter.configure(formatterConfigs);
+
+        } catch (Exception e) {
+            CommandLineUtils.printUsageAndExit(parser, e.getMessage());
+        }
+        return formatter;
+    }
+
+    Properties consumerProps() {
+        return consumerProps;
+    }
+
+    boolean fromBeginning() {
+        return options.has(resetBeginningOpt);
+    }
+
+    long offsetArg() {
+        return offset;
+    }
+
+    boolean skipMessageOnError() {
+        return options.has(skipMessageOnErrorOpt);
+    }
+
+    OptionalInt partitionArg() {
+        if (options.has(partitionIdOpt)) {
+            return OptionalInt.of(options.valueOf(partitionIdOpt));
+        }
+        return OptionalInt.empty();
+    }
+
+    String topicArg() {
+        return options.valueOf(topicOpt);
+    }
+
+    int maxMessages() {
+        return options.has(maxMessagesOpt) ? options.valueOf(maxMessagesOpt) : -1;
+    }
+
+    int timeoutMs() {
+        return options.has(timeoutMsOpt) ? options.valueOf(timeoutMsOpt) : -1;
+    }
+
+    boolean enableSystestEventsLogging() {
+        return options.has(enableSystestEventsLoggingOpt);
+    }
+
+    String bootstrapServer() {
+        return options.valueOf(bootstrapServerOpt);
+    }
+
+    String includedTopicsArg() {
+        return options.has(includeOpt)
+                ? options.valueOf(includeOpt)
+                : options.valueOf(whitelistOpt);
+    }
+
+    Properties formatterArgs() throws IOException {
+        Properties formatterArgs = options.has(messageFormatterConfigOpt)
+                ? Utils.loadProps(options.valueOf(messageFormatterConfigOpt))
+                : new Properties();
+        String keyDeserializer = options.valueOf(keyDeserializerOpt);
+        if (keyDeserializer != null && !keyDeserializer.isEmpty()) {
+            formatterArgs.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer);
+        }
+        String valueDeserializer = options.valueOf(valueDeserializerOpt);
+        if (valueDeserializer != null && !valueDeserializer.isEmpty()) {
+            formatterArgs.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer);
+        }
+        formatterArgs.putAll(CommandLineUtils.parseKeyValueArgs(options.valuesOf(messageFormatterArgOpt)));
+
+        return formatterArgs;
+    }
+
+    MessageFormatter formatter() {
+        return formatter;
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/DefaultMessageFormatter.java b/tools/src/main/java/org/apache/kafka/tools/consumer/DefaultMessageFormatter.java
new file mode 100644
index 0000000000000..402deb8358ea2
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/DefaultMessageFormatter.java
@@ -0,0 +1,210 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Optional;
+
+class DefaultMessageFormatter implements MessageFormatter {
+
+    private static final Logger LOG = LoggerFactory.getLogger(DefaultMessageFormatter.class);
+
+    private boolean printTimestamp = false;
+    private boolean printKey = false;
+    private boolean printValue = true;
+    private boolean printPartition = false;
+    private boolean printOffset = false;
+    private boolean printHeaders = false;
+    private byte[] keySeparator = utfBytes("\t");
+    private byte[] lineSeparator = utfBytes("\n");
+    private byte[] headersSeparator = utfBytes(",");
+    private byte[] nullLiteral = utfBytes("null");
+
+    private Optional<Deserializer<?>> keyDeserializer = Optional.empty();
+    private Optional<Deserializer<?>> valueDeserializer = Optional.empty();
+    private Optional<Deserializer<?>> headersDeserializer = Optional.empty();
+
+    @Override
+    public void configure(Map<String, ?> configs) {
+        if (configs.containsKey("print.timestamp")) {
+            printTimestamp = getBoolProperty(configs, "print.timestamp");
+        }
+        if (configs.containsKey("print.key")) {
+            printKey = getBoolProperty(configs, "print.key");
+        }
+        if (configs.containsKey("print.offset")) {
+            printOffset = getBoolProperty(configs, "print.offset");
+        }
+        if (configs.containsKey("print.partition")) {
+            printPartition = getBoolProperty(configs, "print.partition");
+        }
+        if (configs.containsKey("print.headers")) {
+            printHeaders = getBoolProperty(configs, "print.headers");
+        }
+        if (configs.containsKey("print.value")) {
+            printValue = getBoolProperty(configs, "print.value");
+        }
+        if (configs.containsKey("key.separator")) {
+            keySeparator = getByteProperty(configs, "key.separator");
+        }
+        if (configs.containsKey("line.separator")) {
+            lineSeparator = getByteProperty(configs, "line.separator");
+        }
+        if (configs.containsKey("headers.separator")) {
+            headersSeparator = getByteProperty(configs, "headers.separator");
+        }
+        if (configs.containsKey("null.literal")) {
+            nullLiteral = getByteProperty(configs, "null.literal");
+        }
+
+        keyDeserializer = getDeserializerProperty(configs, true, "key.deserializer");
+        valueDeserializer = getDeserializerProperty(configs, false, "value.deserializer");
+        headersDeserializer = getDeserializerProperty(configs, false, "headers.deserializer");
+    }
+
+    // for testing
+    public boolean printValue() {
+        return printValue;
+    }
+
+    // for testing
+    public Optional<Deserializer<?>> keyDeserializer() {
+        return keyDeserializer;
+    }
+
+    private void writeSeparator(PrintStream output, boolean columnSeparator) {
+        try {
+            if (columnSeparator) {
+                output.write(keySeparator);
+            } else {
+                output.write(lineSeparator);
+            }
+        } catch (IOException ioe) {
+            LOG.error("Unable to write the separator to the output", ioe);
+        }
+    }
+
+    private byte[] deserialize(ConsumerRecord<byte[], byte[]> consumerRecord, Optional<Deserializer<?>> deserializer, byte[] sourceBytes, String topic) {
+        byte[] nonNullBytes = sourceBytes != null ? sourceBytes : nullLiteral;
+        return deserializer.map(value -> utfBytes(value.deserialize(topic, consumerRecord.headers(), nonNullBytes).toString())).orElse(nonNullBytes);
+    }
+
+    @Override
+    public void writeTo(ConsumerRecord<byte[], byte[]> consumerRecord, PrintStream output) {
+        try {
+            if (printTimestamp) {
+                if (consumerRecord.timestampType() != TimestampType.NO_TIMESTAMP_TYPE) {
+                    output.print(consumerRecord.timestampType() + ":" + consumerRecord.timestamp());
+                } else {
+                    output.print("NO_TIMESTAMP");
+                }
+                writeSeparator(output, printOffset || printPartition || printHeaders || printKey || printValue);
+            }
+
+            if (printPartition) {
+                output.print("Partition:");
+                output.print(consumerRecord.partition());
+                writeSeparator(output, printOffset || printHeaders || printKey || printValue);
+            }
+
+            if (printOffset) {
+                output.print("Offset:");
+                output.print(consumerRecord.offset());
+                writeSeparator(output, printHeaders || printKey || printValue);
+            }
+
+            if (printHeaders) {
+                Iterator<Header> headersIt = consumerRecord.headers().iterator();
+                if (!headersIt.hasNext()) {
+                    output.print("NO_HEADERS");
+                } else {
+                    while (headersIt.hasNext()) {
+                        Header header = headersIt.next();
+                        output.print(header.key() + ":");
+                        output.write(deserialize(consumerRecord, headersDeserializer, header.value(), consumerRecord.topic()));
+                        if (headersIt.hasNext()) {
+                            output.write(headersSeparator);
+                        }
+                    }
+                }
+                writeSeparator(output, printKey || printValue);
+            }
+
+            if (printKey) {
+                output.write(deserialize(consumerRecord, keyDeserializer, consumerRecord.key(), consumerRecord.topic()));
+                writeSeparator(output, printValue);
+            }
+
+            if (printValue) {
+                output.write(deserialize(consumerRecord, valueDeserializer, consumerRecord.value(), consumerRecord.topic()));
+                output.write(lineSeparator);
+            }
+        } catch (IOException ioe) {
+            LOG.error("Unable to write the consumer record to the output", ioe);
+        }
+    }
+
+    private Map<String, ?> propertiesWithKeyPrefixStripped(String prefix, Map<String, ?> configs) {
+        final Map<String, Object> newConfigs = new HashMap<>();
+        for (Map.Entry<String, ?> entry : configs.entrySet()) {
+            if (entry.getKey().startsWith(prefix) && entry.getKey().length() > prefix.length()) {
+                newConfigs.put(entry.getKey().substring(prefix.length()), entry.getValue());
+            }
+        }
+        return newConfigs;
+    }
+
+    private byte[] utfBytes(String str) {
+        return str.getBytes(StandardCharsets.UTF_8);
+    }
+
+    private byte[] getByteProperty(Map<String, ?> configs, String key) {
+        return utfBytes((String) configs.get(key));
+    }
+
+    private boolean getBoolProperty(Map<String, ?> configs, String key) {
+        return ((String) configs.get(key)).trim().equalsIgnoreCase("true");
+    }
+
+    private Optional<Deserializer<?>> getDeserializerProperty(Map<String, ?> configs, boolean isKey, String key) {
+        if (configs.containsKey(key)) {
+            String name = (String) configs.get(key);
+            try {
+                Deserializer<?> deserializer = (Deserializer<?>) Class.forName(name).getDeclaredConstructor().newInstance();
+                Map<String, ?> deserializerConfig = propertiesWithKeyPrefixStripped(key + ".", configs);
+                deserializer.configure(deserializerConfig, isKey);
+                return Optional.of(deserializer);
+            } catch (Exception e) {
+                LOG.error("Unable to instantiate a deserializer from " + name, e);
+            }
+        }
+        return Optional.empty();
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/LoggingMessageFormatter.java b/tools/src/main/java/org/apache/kafka/tools/consumer/LoggingMessageFormatter.java
new file mode 100644
index 0000000000000..f010bbaea9b28
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/LoggingMessageFormatter.java
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.record.TimestampType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+
+class LoggingMessageFormatter implements MessageFormatter {
+
+    private static final Logger LOG = LoggerFactory.getLogger(LoggingMessageFormatter.class);
+    private final DefaultMessageFormatter defaultWriter = new DefaultMessageFormatter();
+
+    @Override
+    public void configure(Map<String, ?> configs) {
+        defaultWriter.configure(configs);
+    }
+
+    @Override
+    public void writeTo(ConsumerRecord<byte[], byte[]> consumerRecord, PrintStream output) {
+        defaultWriter.writeTo(consumerRecord, output);
+        String timestamp = consumerRecord.timestampType() != TimestampType.NO_TIMESTAMP_TYPE
+                ? consumerRecord.timestampType() + ":" + consumerRecord.timestamp() + ", "
+                : "";
+        String key = "key:" + (consumerRecord.key() == null ? "null " : new String(consumerRecord.key(), StandardCharsets.UTF_8) + ", ");
+        String value = "value:" + (consumerRecord.value() == null ? "null" : new String(consumerRecord.value(), StandardCharsets.UTF_8));
+        LOG.info(timestamp + key + value);
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/NoOpMessageFormatter.java b/tools/src/main/java/org/apache/kafka/tools/consumer/NoOpMessageFormatter.java
new file mode 100644
index 0000000000000..cba9e9b67289e
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/NoOpMessageFormatter.java
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+
+import java.io.PrintStream;
+
+class NoOpMessageFormatter implements MessageFormatter {
+
+    @Override
+    public void writeTo(ConsumerRecord<byte[], byte[]> consumerRecord, PrintStream output) {
+
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
index 9b698b0cc5f17..fdc732ea29ab1 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
@@ -28,6 +28,8 @@
 import org.apache.kafka.storage.internals.log.LogConfig;
 
 import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
 import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -198,6 +200,14 @@ public static String formatReplicaThrottles(Map<TopicPartition, List<Integer>> m
             .collect(Collectors.joining(","));
     }
 
+    public static File tempPropertiesFile(Map<String, String> properties) throws IOException {
+        StringBuilder sb = new StringBuilder();
+        for (Map.Entry<String, String> entry : properties.entrySet()) {
+            sb.append(entry.getKey() + "=" + entry.getValue() + System.lineSeparator());
+        }
+        return org.apache.kafka.test.TestUtils.tempFile(sb.toString());
+    }
+
     public static class MockExitProcedure implements Exit.Procedure {
         private boolean hasExited = false;
         private int statusCode;
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
new file mode 100644
index 0000000000000..523122c4cdd98
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
@@ -0,0 +1,621 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.common.requests.ListOffsetsRequest;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.test.MockDeserializer;
+import org.apache.kafka.tools.ToolsTestUtils;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class ConsoleConsumerOptionsTest {
+
+    @Test
+    public void shouldParseValidConsumerValidConfig() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertTrue(config.fromBeginning());
+        assertFalse(config.enableSystestEventsLogging());
+        assertFalse(config.skipMessageOnError());
+        assertEquals(-1, config.maxMessages());
+        assertEquals(-1, config.timeoutMs());
+    }
+
+    @Test
+    public void shouldParseIncludeArgument() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--include", "includeTest*",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("includeTest*", config.includedTopicsArg());
+        assertTrue(config.fromBeginning());
+    }
+
+    @Test
+    public void shouldParseWhitelistArgument() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--whitelist", "whitelistTest*",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("whitelistTest*", config.includedTopicsArg());
+        assertTrue(config.fromBeginning());
+    }
+
+    @Test
+    public void shouldIgnoreWhitelistArgumentIfIncludeSpecified() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--include", "includeTest*",
+            "--whitelist", "whitelistTest*",
+            "--from-beginning"
+        };
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("includeTest*", config.includedTopicsArg());
+        assertTrue(config.fromBeginning());
+    }
+
+    @Test
+    public void shouldParseValidSimpleConsumerValidConfigWithNumericOffset() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0",
+            "--offset", "3"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertTrue(config.partitionArg().isPresent());
+        assertEquals(0, config.partitionArg().getAsInt());
+        assertEquals(3, config.offsetArg());
+        assertFalse(config.fromBeginning());
+    }
+
+    @Test
+    public void shouldExitOnUnrecognizedNewConsumerOption() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--new-consumer",
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--from-beginning"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitIfPartitionButNoTopic() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--include", "test.*",
+            "--partition", "0"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitIfFromBeginningAndOffset() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0",
+            "--from-beginning",
+            "--offset", "123"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldParseValidSimpleConsumerValidConfigWithStringOffset() throws Exception {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0",
+            "--offset", "LatEst",
+            "--property", "print.value=false"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertTrue(config.partitionArg().isPresent());
+        assertEquals(0, config.partitionArg().getAsInt());
+        assertEquals(-1, config.offsetArg());
+        assertFalse(config.fromBeginning());
+        assertFalse(((DefaultMessageFormatter) config.formatter()).printValue());
+    }
+
+    @Test
+    public void shouldParseValidConsumerConfigWithAutoOffsetResetLatest() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer-property", "auto.offset.reset=latest"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertFalse(config.fromBeginning());
+        assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
+    }
+
+    @Test
+    public void shouldParseValidConsumerConfigWithAutoOffsetResetEarliest() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer-property", "auto.offset.reset=earliest"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertFalse(config.fromBeginning());
+        assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
+    }
+
+    @Test
+    public void shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer-property", "auto.offset.reset=earliest",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertTrue(config.fromBeginning());
+        assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
+    }
+
+    @Test
+    public void shouldParseValidConsumerConfigWithNoOffsetReset() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertFalse(config.fromBeginning());
+        assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
+    }
+
+    @Test
+    public void shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginning() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer-property", "auto.offset.reset=latest",
+            "--from-beginning"
+        };
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldParseConfigsFromFile() throws IOException {
+        Map<String, String> configs = new HashMap<>();
+        configs.put("request.timeout.ms", "1000");
+        configs.put("group.id", "group1");
+        File propsFile = ToolsTestUtils.tempPropertiesFile(configs);
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer.config", propsFile.getAbsolutePath()
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        assertEquals("1000", config.consumerProps().get("request.timeout.ms"));
+        assertEquals("group1", config.consumerProps().get("group.id"));
+    }
+
+    @Test
+    public void groupIdsProvidedInDifferentPlacesMustMatch() throws IOException {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        // different in all three places
+        File propsFile = ToolsTestUtils.tempPropertiesFile(Collections.singletonMap("group.id", "group-from-file"));
+        final String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "group-from-arguments",
+            "--consumer-property", "group.id=group-from-properties",
+            "--consumer.config", propsFile.getAbsolutePath()
+        };
+
+        assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+
+        // the same in all three places
+        propsFile = ToolsTestUtils.tempPropertiesFile(Collections.singletonMap("group.id", "test-group"));
+        final String[] args1 = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "test-group",
+            "--consumer-property", "group.id=test-group",
+            "--consumer.config", propsFile.getAbsolutePath()
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args1);
+        Properties props = config.consumerProps();
+        assertEquals("test-group", props.getProperty("group.id"));
+
+        // different via --consumer-property and --consumer.config
+        propsFile = ToolsTestUtils.tempPropertiesFile(Collections.singletonMap("group.id", "group-from-file"));
+        final String[] args2 = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--consumer-property", "group.id=group-from-properties",
+            "--consumer.config", propsFile.getAbsolutePath()
+        };
+
+        assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args2));
+
+        // different via --consumer-property and --group
+        final String[] args3 = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "group-from-arguments",
+            "--consumer-property", "group.id=group-from-properties"
+        };
+
+        assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args3));
+
+        // different via --group and --consumer.config
+        propsFile = ToolsTestUtils.tempPropertiesFile(Collections.singletonMap("group.id", "group-from-file"));
+        final String[] args4 = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "group-from-arguments",
+            "--consumer.config", propsFile.getAbsolutePath()
+        };
+        assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args4));
+
+        // via --group only
+        final String[] args5 = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "group-from-arguments"
+        };
+
+        config = new ConsoleConsumerOptions(args5);
+        props = config.consumerProps();
+        assertEquals("group-from-arguments", props.getProperty("group.id"));
+
+        Exit.resetExitProcedure();
+    }
+
+    @Test
+    public void testCustomPropertyShouldBePassedToConfigureMethod() throws Exception {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--property", "print.key=true",
+            "--property", "key.deserializer=org.apache.kafka.test.MockDeserializer",
+            "--property", "key.deserializer.my-props=abc"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertTrue(config.formatter() instanceof DefaultMessageFormatter);
+        assertTrue(config.formatterArgs().containsKey("key.deserializer.my-props"));
+        DefaultMessageFormatter formatter = (DefaultMessageFormatter) config.formatter();
+        assertTrue(formatter.keyDeserializer().isPresent());
+        assertTrue(formatter.keyDeserializer().get() instanceof MockDeserializer);
+        MockDeserializer keyDeserializer = (MockDeserializer) formatter.keyDeserializer().get();
+        assertEquals(1, keyDeserializer.configs.size());
+        assertEquals("abc", keyDeserializer.configs.get("my-props"));
+        assertTrue(keyDeserializer.isKey);
+    }
+
+    @Test
+    public void testCustomConfigShouldBePassedToConfigureMethod() throws Exception {
+        Map<String, String> configs = new HashMap<>();
+        configs.put("key.deserializer.my-props", "abc");
+        configs.put("print.key", "false");
+        File propsFile = ToolsTestUtils.tempPropertiesFile(configs);
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--property", "print.key=true",
+            "--property", "key.deserializer=org.apache.kafka.test.MockDeserializer",
+            "--formatter-config", propsFile.getAbsolutePath()
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+
+        assertTrue(config.formatter() instanceof DefaultMessageFormatter);
+        assertTrue(config.formatterArgs().containsKey("key.deserializer.my-props"));
+        DefaultMessageFormatter formatter = (DefaultMessageFormatter) config.formatter();
+        assertTrue(formatter.keyDeserializer().isPresent());
+        assertTrue(formatter.keyDeserializer().get() instanceof MockDeserializer);
+        MockDeserializer keyDeserializer = (MockDeserializer) formatter.keyDeserializer().get();
+        assertEquals(1, keyDeserializer.configs.size());
+        assertEquals("abc", keyDeserializer.configs.get("my-props"));
+        assertTrue(keyDeserializer.isKey);
+    }
+
+    @Test
+    public void shouldParseGroupIdFromBeginningGivenTogether() throws IOException {
+        // Start from earliest
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "test-group",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertEquals(-2, config.offsetArg());
+        assertTrue(config.fromBeginning());
+
+        // Start from latest
+        args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "test-group"
+        };
+
+        config = new ConsoleConsumerOptions(args);
+        assertEquals("localhost:9092", config.bootstrapServer());
+        assertEquals("test", config.topicArg());
+        assertEquals(-1, config.offsetArg());
+        assertFalse(config.fromBeginning());
+    }
+
+    @Test
+    public void shouldExitOnGroupIdAndPartitionGivenTogether() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--group", "test-group",
+            "--partition", "0"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitOnOffsetWithoutPartition() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--offset", "10"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitIfNoTopicOrFilterSpecified() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitIfTopicAndIncludeSpecified() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--include", "includeTest*"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitIfTopicAndWhitelistSpecified() {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--whitelist", "whitelistTest*"
+        };
+
+        try {
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(args));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void testClientIdOverride() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--from-beginning",
+            "--consumer-property", "client.id=consumer-1"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("consumer-1", consumerProperties.getProperty(ConsumerConfig.CLIENT_ID_CONFIG));
+    }
+
+    @Test
+    public void testDefaultClientId() throws IOException {
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--from-beginning"
+        };
+
+        ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
+        Properties consumerProperties = config.consumerProps();
+
+        assertEquals("console-consumer", consumerProperties.getProperty(ConsumerConfig.CLIENT_ID_CONFIG));
+    }
+
+    @Test
+    public void testParseOffset() throws Exception {
+        Exit.setExitProcedure((code, message) -> {
+            throw new IllegalArgumentException(message);
+        });
+
+        try {
+            final String[] badOffset = new String[]{
+                "--bootstrap-server", "localhost:9092",
+                "--topic", "test",
+                "--partition", "0",
+                "--offset", "bad"
+            };
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(badOffset));
+
+            final String[] negativeOffset = new String[]{
+                "--bootstrap-server", "localhost:9092",
+                "--topic", "test",
+                "--partition", "0",
+                "--offset", "-100"
+            };
+            assertThrows(IllegalArgumentException.class, () -> new ConsoleConsumerOptions(negativeOffset));
+
+            final String[] earliestOffset = new String[]{
+                "--bootstrap-server", "localhost:9092",
+                "--topic", "test",
+                "--partition", "0",
+                "--offset", "earliest"
+            };
+            ConsoleConsumerOptions config = new ConsoleConsumerOptions(earliestOffset);
+            assertEquals(ListOffsetsRequest.EARLIEST_TIMESTAMP, config.offsetArg());
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
new file mode 100644
index 0000000000000..008893f9c505a
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
@@ -0,0 +1,211 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.MockConsumer;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.requests.ListOffsetsRequest;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.server.util.MockTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.PrintStream;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ConsoleConsumerTest {
+
+    @BeforeEach
+    public void setup() {
+        ConsoleConsumer.messageCount = 0;
+    }
+
+    @Test
+    public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() {
+        String topic = "test";
+        final Time time = new MockTime();
+        final int timeoutMs = 1000;
+
+        @SuppressWarnings("unchecked")
+        Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
+
+        when(mockConsumer.poll(Duration.ofMillis(timeoutMs))).thenAnswer(invocation -> {
+            time.sleep(timeoutMs / 2 + 1);
+            return ConsumerRecords.EMPTY;
+        });
+
+        ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
+                Optional.of(topic),
+                OptionalInt.empty(),
+                OptionalLong.empty(),
+                Optional.empty(),
+                mockConsumer,
+                timeoutMs
+        );
+
+        assertThrows(TimeoutException.class, consumer::receive);
+    }
+
+    @Test
+    public void shouldResetUnConsumedOffsetsBeforeExit() {
+        String topic = "test";
+        int maxMessages = 123;
+        int totalMessages = 700;
+        long startOffset = 0L;
+
+        MockConsumer<byte[], byte[]> mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
+        TopicPartition tp1 = new TopicPartition(topic, 0);
+        TopicPartition tp2 = new TopicPartition(topic, 1);
+
+        ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
+                Optional.of(topic),
+                OptionalInt.empty(),
+                OptionalLong.empty(),
+                Optional.empty(),
+                mockConsumer,
+                1000L);
+
+        mockConsumer.rebalance(Arrays.asList(tp1, tp2));
+        Map<TopicPartition, Long> offsets = new HashMap<>();
+        offsets.put(tp1, startOffset);
+        offsets.put(tp2, startOffset);
+        mockConsumer.updateBeginningOffsets(offsets);
+
+        for (int i = 0; i < totalMessages; i++) {
+            // add all records, each partition should have half of `totalMessages`
+            mockConsumer.addRecord(new ConsumerRecord<>(topic, i % 2, i / 2, "key".getBytes(), "value".getBytes()));
+        }
+
+        MessageFormatter formatter = mock(MessageFormatter.class);
+
+        ConsoleConsumer.process(maxMessages, formatter, consumer, System.out, false);
+        assertEquals(totalMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2));
+
+        consumer.resetUnconsumedOffsets();
+        assertEquals(maxMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2));
+
+        verify(formatter, times(maxMessages)).writeTo(any(), any());
+        consumer.cleanup();
+    }
+
+    @Test
+    public void shouldLimitReadsToMaxMessageLimit() {
+        ConsoleConsumer.ConsumerWrapper consumer = mock(ConsoleConsumer.ConsumerWrapper.class);
+        MessageFormatter formatter = mock(MessageFormatter.class);
+        ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>("foo", 1, 1, new byte[0], new byte[0]);
+
+        int messageLimit = 10;
+        when(consumer.receive()).thenReturn(record);
+
+        ConsoleConsumer.process(messageLimit, formatter, consumer, System.out, true);
+
+        verify(consumer, times(messageLimit)).receive();
+        verify(formatter, times(messageLimit)).writeTo(any(), any());
+
+        consumer.cleanup();
+    }
+
+    @Test
+    public void shouldStopWhenOutputCheckErrorFails() {
+        ConsoleConsumer.ConsumerWrapper consumer = mock(ConsoleConsumer.ConsumerWrapper.class);
+        MessageFormatter formatter = mock(MessageFormatter.class);
+        PrintStream printStream = mock(PrintStream.class);
+
+        ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>("foo", 1, 1, new byte[0], new byte[0]);
+
+        when(consumer.receive()).thenReturn(record);
+        //Simulate an error on System.out after the first record has been printed
+        when(printStream.checkError()).thenReturn(true);
+
+        ConsoleConsumer.process(-1, formatter, consumer, printStream, true);
+
+        verify(formatter).writeTo(any(), eq(printStream));
+        verify(consumer).receive();
+        verify(printStream).checkError();
+
+        consumer.cleanup();
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void shouldSeekWhenOffsetIsSet() {
+        Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
+        TopicPartition tp0 = new TopicPartition("test", 0);
+
+        ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
+                Optional.of(tp0.topic()),
+                OptionalInt.of(tp0.partition()),
+                OptionalLong.empty(),
+                Optional.empty(),
+                mockConsumer,
+                1000L);
+
+        verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
+        verify(mockConsumer).seekToEnd(eq(Collections.singletonList(tp0)));
+        consumer.cleanup();
+        reset(mockConsumer);
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(
+                Optional.of(tp0.topic()),
+                OptionalInt.of(tp0.partition()),
+                OptionalLong.of(123L),
+                Optional.empty(),
+                mockConsumer,
+                1000L);
+
+        verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
+        verify(mockConsumer).seek(eq(tp0), eq(123L));
+        consumer.cleanup();
+        reset(mockConsumer);
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(
+                Optional.of(tp0.topic()),
+                OptionalInt.of(tp0.partition()),
+                OptionalLong.of(ListOffsetsRequest.EARLIEST_TIMESTAMP),
+                Optional.empty(),
+                mockConsumer,
+                1000L);
+
+        verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
+        verify(mockConsumer).seekToBeginning(eq(Collections.singletonList(tp0)));
+        consumer.cleanup();
+        reset(mockConsumer);
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/DefaultMessageFormatterTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/DefaultMessageFormatterTest.java
new file mode 100644
index 0000000000000..917d44d3a4f61
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/DefaultMessageFormatterTest.java
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class DefaultMessageFormatterTest {
+
+    @Test
+    public void testDefaultMessageFormatter() {
+        ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>("topic", 0, 123, "key".getBytes(), "value".getBytes());
+        MessageFormatter formatter = new DefaultMessageFormatter();
+        Map<String, String> configs = new HashMap<>();
+
+        formatter.configure(configs);
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("value\n", out.toString());
+
+        configs.put("print.key", "true");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("key\tvalue\n", out.toString());
+
+        configs.put("print.partition", "true");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("Partition:0\tkey\tvalue\n", out.toString());
+
+        configs.put("print.timestamp", "true");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("NO_TIMESTAMP\tPartition:0\tkey\tvalue\n", out.toString());
+
+        configs.put("print.offset", "true");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("NO_TIMESTAMP\tPartition:0\tOffset:123\tkey\tvalue\n", out.toString());
+
+        configs.put("print.headers", "true");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("NO_TIMESTAMP\tPartition:0\tOffset:123\tNO_HEADERS\tkey\tvalue\n", out.toString());
+
+        RecordHeaders headers = new RecordHeaders();
+        headers.add("h1", "v1".getBytes());
+        headers.add("h2", "v2".getBytes());
+        record = new ConsumerRecord<>("topic", 0, 123, 123L, TimestampType.CREATE_TIME, -1, -1, "key".getBytes(), "value".getBytes(),
+                headers, Optional.empty());
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123\tPartition:0\tOffset:123\th1:v1,h2:v2\tkey\tvalue\n", out.toString());
+
+        configs.put("print.value", "false");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123\tPartition:0\tOffset:123\th1:v1,h2:v2\tkey\n", out.toString());
+
+        configs.put("key.separator", "<sep>");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:v1,h2:v2<sep>key\n", out.toString());
+
+        configs.put("line.separator", "<end>");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:v1,h2:v2<sep>key<end>", out.toString());
+
+        configs.put("headers.separator", "|");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:v1|h2:v2<sep>key<end>", out.toString());
+
+        record = new ConsumerRecord<>("topic", 0, 123, 123L, TimestampType.CREATE_TIME, -1, -1, "key".getBytes(), "value".getBytes(),
+                headers, Optional.empty());
+
+        configs.put("key.deserializer", UpperCaseDeserializer.class.getName());
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:v1|h2:v2<sep>KEY<end>", out.toString());
+
+        configs.put("headers.deserializer", UpperCaseDeserializer.class.getName());
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:V1|h2:V2<sep>KEY<end>", out.toString());
+
+        record = new ConsumerRecord<>("topic", 0, 123, 123L, TimestampType.CREATE_TIME, -1, -1, "key".getBytes(), null,
+                headers, Optional.empty());
+
+        configs.put("print.value", "true");
+        configs.put("null.literal", "<null>");
+        formatter.configure(configs);
+        out = new ByteArrayOutputStream();
+        formatter.writeTo(record, new PrintStream(out));
+        assertEquals("CreateTime:123<sep>Partition:0<sep>Offset:123<sep>h1:V1|h2:V2<sep>KEY<sep><null><end>", out.toString());
+        formatter.close();
+    }
+
+    static class UpperCaseDeserializer implements Deserializer<String> {
+
+        @Override
+        public String deserialize(String topic, byte[] data) {
+            return new String(data, StandardCharsets.UTF_8).toUpperCase(Locale.ROOT);
+        }
+    }
+}
+
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/NoOpMessageFormatterTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/NoOpMessageFormatterTest.java
new file mode 100644
index 0000000000000..1652a484d10ff
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/NoOpMessageFormatterTest.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.util.HashMap;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class NoOpMessageFormatterTest {
+
+    @Test
+    public void testNoOpMessageFormatter() {
+        ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>("topic", 0, 123, "key".getBytes(), "value".getBytes());
+        try (MessageFormatter formatter = new NoOpMessageFormatter()) {
+            formatter.configure(new HashMap<>());
+            ByteArrayOutputStream out = new ByteArrayOutputStream();
+            formatter.writeTo(record, new PrintStream(out));
+            assertEquals("", out.toString());
+        }
+    }
+}

From be6653c8bc25717e25a7db164527635a6579b4cc Mon Sep 17 00:00:00 2001
From: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
Date: Tue, 13 Feb 2024 22:13:53 +0000
Subject: [PATCH 035/258] KAFKA-16225 [1/N]: Set metadata.log.dir to broker in
 KRAFT mode in integration test

Fix the flakiness of LogDirFailureTest by setting a separate metadata.log.dir for brokers in KRAFT mode.

The test was flaky because as we call causeLogDirFailure some times we impact the first log.dir which also is KafkaConfig.metadataLogDir as we don't have metadata.log.dir. So to fix the flakiness we need to explicitly set metadata.log.dir to diff log dir than the ones we could potentially fail for the tests.

This is part 1 of the fixes. Delivering them separately as the other issues were not as clear cut.

Reviewers: Gaurav Narula <gaurav_narula2@apple.com>, Justine Olshan <jolshan@confluent.io>, Greg Harris <greg.harris@aiven.io>
---
 .../integration/kafka/api/IntegrationTestHarness.scala     | 5 +++++
 .../test/scala/unit/kafka/server/LogDirFailureTest.scala   | 7 ++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
index 2381976f378f9..5d3115c98ec63 100644
--- a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
+++ b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
@@ -70,6 +70,11 @@ abstract class IntegrationTestHarness extends KafkaServerTestHarness {
     if (isNewGroupCoordinatorEnabled()) {
       cfgs.foreach(_.setProperty(KafkaConfig.NewGroupCoordinatorEnableProp, "true"))
     }
+
+    if(isKRaftTest()) {
+      cfgs.foreach(_.setProperty(KafkaConfig.MetadataLogDirProp, TestUtils.tempDir().getAbsolutePath))
+    }
+
     insertControllerListenersIfNeeded(cfgs)
     cfgs.map(KafkaConfig.fromProps)
   }
diff --git a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
index c63f4664096ac..57be102bea042 100644
--- a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
@@ -21,7 +21,7 @@ import java.util.Collections
 import java.util.concurrent.{ExecutionException, TimeUnit}
 import kafka.api.IntegrationTestHarness
 import kafka.controller.{OfflineReplica, PartitionAndReplica}
-import kafka.utils.TestUtils.{waitUntilTrue, Checkpoint, LogDirFailureType, Roll}
+import kafka.utils.TestUtils.{Checkpoint, LogDirFailureType, Roll, waitUntilTrue}
 import kafka.utils.{CoreUtils, Exit, TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.consumer.Consumer
 import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
@@ -56,6 +56,7 @@ class LogDirFailureTest extends IntegrationTestHarness {
   override def setUp(testInfo: TestInfo): Unit = {
     super.setUp(testInfo)
     createTopic(topic, partitionNum, brokerCount)
+    ensureConsistentKRaftMetadata()
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
@@ -168,6 +169,10 @@ class LogDirFailureTest extends IntegrationTestHarness {
   }
 
   def testProduceAfterLogDirFailureOnLeader(failureType: LogDirFailureType, quorum: String): Unit = {
+    if (isKRaftTest()) {
+      val value = configs.map(c => c.brokerId -> c.logDirs.contains(c.metadataLogDir))
+      logger.warn(s">>>>>> ${value.mkString(",")}")
+    }
     val consumer = createConsumer()
     subscribeAndWaitForAssignment(topic, consumer)
 

From d24abe0edebad37e554adea47408c3063037f744 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Tue, 13 Feb 2024 23:29:29 -0800
Subject: [PATCH 036/258] MINOR: Refactor GroupMetadataManagerTest (#15348)

`GroupMetadataManagerTest` class got a little under control. We have too many things defined in it. As a first steps, this patch extracts all the inner classes. It also extracts all the helper methods. However, the logic is not changed at all.

Reviewers: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Justine Olshan <jolshan@confluent.io>
---
 checkstyle/suppressions.xml                   |    4 +-
 .../kafka/coordinator/group/Assertions.java   |  221 ++
 .../group/GroupMetadataManagerTest.java       | 2466 ++++-------------
 .../GroupMetadataManagerTestContext.java      | 1277 +++++++++
 .../group/MetadataImageBuilder.java           |   66 +
 .../group/MockPartitionAssignor.java          |   46 +
 .../group/OffsetMetadataManagerTest.java      |    2 +-
 .../coordinator/group/RecordHelpersTest.java  |  144 +-
 .../group/consumer/ConsumerGroupBuilder.java  |  114 +
 .../consumer/ConsumerGroupMemberTest.java     |    8 +-
 .../group/consumer/ConsumerGroupTest.java     |    8 +-
 .../consumer/TargetAssignmentBuilderTest.java |   20 +-
 12 files changed, 2226 insertions(+), 2150 deletions(-)
 create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
 create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
 create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MetadataImageBuilder.java
 create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MockPartitionAssignor.java
 create mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupBuilder.java

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 6a8bc0b1852eb..7486ef9a80d3e 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -336,11 +336,11 @@
     <suppress checks="CyclomaticComplexity"
               files="(ConsumerGroupMember|GroupMetadataManager|GeneralUniformAssignmentBuilder).java"/>
     <suppress checks="(NPathComplexity|MethodLength)"
-              files="(GroupMetadataManager|ConsumerGroupTest|GroupMetadataManagerTest|GeneralUniformAssignmentBuilder).java"/>
+              files="(GroupMetadataManager|ConsumerGroupTest|GroupMetadataManagerTest|GroupMetadataManagerTestContext|GeneralUniformAssignmentBuilder).java"/>
     <suppress checks="NPathComplexity"
               files="CoordinatorRuntime.java"/>
     <suppress checks="ClassFanOutComplexity"
-              files="(GroupMetadataManager|GroupMetadataManagerTest|GroupCoordinatorService|GroupCoordinatorServiceTest).java"/>
+              files="(GroupMetadataManager|GroupMetadataManagerTest|GroupMetadataManagerTestContext|GroupCoordinatorService|GroupCoordinatorServiceTest).java"/>
     <suppress checks="ParameterNumber"
               files="(ConsumerGroupMember|GroupMetadataManager|GroupCoordinatorConfig).java"/>
     <suppress checks="ClassDataAbstractionCouplingCheck"
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
new file mode 100644
index 0000000000000..77d4c7d03c8d8
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataValue;
+import org.apache.kafka.server.common.ApiMessageAndVersion;
+import org.opentest4j.AssertionFailedError;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class Assertions {
+    public static <T> void assertUnorderedListEquals(
+        List<T> expected,
+        List<T> actual
+    ) {
+        assertEquals(new HashSet<>(expected), new HashSet<>(actual));
+    }
+
+    public static void assertResponseEquals(
+        ConsumerGroupHeartbeatResponseData expected,
+        ConsumerGroupHeartbeatResponseData actual
+    ) {
+        if (!responseEquals(expected, actual)) {
+            assertionFailure()
+                .expected(expected)
+                .actual(actual)
+                .buildAndThrow();
+        }
+    }
+
+    private static boolean responseEquals(
+        ConsumerGroupHeartbeatResponseData expected,
+        ConsumerGroupHeartbeatResponseData actual
+    ) {
+        if (expected.throttleTimeMs() != actual.throttleTimeMs()) return false;
+        if (expected.errorCode() != actual.errorCode()) return false;
+        if (!Objects.equals(expected.errorMessage(), actual.errorMessage())) return false;
+        if (!Objects.equals(expected.memberId(), actual.memberId())) return false;
+        if (expected.memberEpoch() != actual.memberEpoch()) return false;
+        if (expected.heartbeatIntervalMs() != actual.heartbeatIntervalMs()) return false;
+        // Unordered comparison of the assignments.
+        return responseAssignmentEquals(expected.assignment(), actual.assignment());
+    }
+
+    private static boolean responseAssignmentEquals(
+        ConsumerGroupHeartbeatResponseData.Assignment expected,
+        ConsumerGroupHeartbeatResponseData.Assignment actual
+    ) {
+        if (expected == actual) return true;
+        if (expected == null) return false;
+        if (actual == null) return false;
+
+        return Objects.equals(fromAssignment(expected.topicPartitions()), fromAssignment(actual.topicPartitions()));
+    }
+
+    private static Map<Uuid, Set<Integer>> fromAssignment(
+        List<ConsumerGroupHeartbeatResponseData.TopicPartitions> assignment
+    ) {
+        if (assignment == null) return null;
+
+        Map<Uuid, Set<Integer>> assignmentMap = new HashMap<>();
+        assignment.forEach(topicPartitions ->
+            assignmentMap.put(topicPartitions.topicId(), new HashSet<>(topicPartitions.partitions()))
+        );
+        return assignmentMap;
+    }
+
+    public static void assertRecordsEquals(
+        List<Record> expectedRecords,
+        List<Record> actualRecords
+    ) {
+        try {
+            assertEquals(expectedRecords.size(), actualRecords.size());
+
+            for (int i = 0; i < expectedRecords.size(); i++) {
+                Record expectedRecord = expectedRecords.get(i);
+                Record actualRecord = actualRecords.get(i);
+                assertRecordEquals(expectedRecord, actualRecord);
+            }
+        } catch (AssertionFailedError e) {
+            assertionFailure()
+                .expected(expectedRecords)
+                .actual(actualRecords)
+                .buildAndThrow();
+        }
+    }
+
+    public static void assertRecordEquals(
+        Record expected,
+        Record actual
+    ) {
+        try {
+            assertApiMessageAndVersionEquals(expected.key(), actual.key());
+            assertApiMessageAndVersionEquals(expected.value(), actual.value());
+        } catch (AssertionFailedError e) {
+            assertionFailure()
+                .expected(expected)
+                .actual(actual)
+                .buildAndThrow();
+        }
+    }
+
+    private static void assertApiMessageAndVersionEquals(
+        ApiMessageAndVersion expected,
+        ApiMessageAndVersion actual
+    ) {
+        if (expected == actual) return;
+
+        assertEquals(expected.version(), actual.version());
+
+        if (actual.message() instanceof ConsumerGroupCurrentMemberAssignmentValue) {
+            // The order of the topics stored in ConsumerGroupCurrentMemberAssignmentValue is not
+            // always guaranteed. Therefore, we need a special comparator.
+            ConsumerGroupCurrentMemberAssignmentValue expectedValue =
+                (ConsumerGroupCurrentMemberAssignmentValue) expected.message();
+            ConsumerGroupCurrentMemberAssignmentValue actualValue =
+                (ConsumerGroupCurrentMemberAssignmentValue) actual.message();
+
+            assertEquals(expectedValue.memberEpoch(), actualValue.memberEpoch());
+            assertEquals(expectedValue.previousMemberEpoch(), actualValue.previousMemberEpoch());
+            assertEquals(expectedValue.targetMemberEpoch(), actualValue.targetMemberEpoch());
+            assertEquals(expectedValue.error(), actualValue.error());
+            assertEquals(expectedValue.metadataVersion(), actualValue.metadataVersion());
+            assertEquals(expectedValue.metadataBytes(), actualValue.metadataBytes());
+
+            // We transform those to Maps before comparing them.
+            assertEquals(fromTopicPartitions(expectedValue.assignedPartitions()),
+                fromTopicPartitions(actualValue.assignedPartitions()));
+            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingRevocation()),
+                fromTopicPartitions(actualValue.partitionsPendingRevocation()));
+            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingAssignment()),
+                fromTopicPartitions(actualValue.partitionsPendingAssignment()));
+        } else if (actual.message() instanceof ConsumerGroupPartitionMetadataValue) {
+            // The order of the racks stored in the PartitionMetadata of the ConsumerGroupPartitionMetadataValue
+            // is not always guaranteed. Therefore, we need a special comparator.
+            ConsumerGroupPartitionMetadataValue expectedValue =
+                (ConsumerGroupPartitionMetadataValue) expected.message();
+            ConsumerGroupPartitionMetadataValue actualValue =
+                (ConsumerGroupPartitionMetadataValue) actual.message();
+
+            List<ConsumerGroupPartitionMetadataValue.TopicMetadata> expectedTopicMetadataList =
+                expectedValue.topics();
+            List<ConsumerGroupPartitionMetadataValue.TopicMetadata> actualTopicMetadataList =
+                actualValue.topics();
+
+            if (expectedTopicMetadataList.size() != actualTopicMetadataList.size()) {
+                fail("Topic metadata lists have different sizes");
+            }
+
+            for (int i = 0; i < expectedTopicMetadataList.size(); i++) {
+                ConsumerGroupPartitionMetadataValue.TopicMetadata expectedTopicMetadata =
+                    expectedTopicMetadataList.get(i);
+                ConsumerGroupPartitionMetadataValue.TopicMetadata actualTopicMetadata =
+                    actualTopicMetadataList.get(i);
+
+                assertEquals(expectedTopicMetadata.topicId(), actualTopicMetadata.topicId());
+                assertEquals(expectedTopicMetadata.topicName(), actualTopicMetadata.topicName());
+                assertEquals(expectedTopicMetadata.numPartitions(), actualTopicMetadata.numPartitions());
+
+                List<ConsumerGroupPartitionMetadataValue.PartitionMetadata> expectedPartitionMetadataList =
+                    expectedTopicMetadata.partitionMetadata();
+                List<ConsumerGroupPartitionMetadataValue.PartitionMetadata> actualPartitionMetadataList =
+                    actualTopicMetadata.partitionMetadata();
+
+                // If the list is empty, rack information wasn't available for any replica of
+                // the partition and hence, the entry wasn't added to the record.
+                if (expectedPartitionMetadataList.size() != actualPartitionMetadataList.size()) {
+                    fail("Partition metadata lists have different sizes");
+                } else if (!expectedPartitionMetadataList.isEmpty() && !actualPartitionMetadataList.isEmpty()) {
+                    for (int j = 0; j < expectedTopicMetadataList.size(); j++) {
+                        ConsumerGroupPartitionMetadataValue.PartitionMetadata expectedPartitionMetadata =
+                            expectedPartitionMetadataList.get(j);
+                        ConsumerGroupPartitionMetadataValue.PartitionMetadata actualPartitionMetadata =
+                            actualPartitionMetadataList.get(j);
+
+                        assertEquals(expectedPartitionMetadata.partition(), actualPartitionMetadata.partition());
+                        assertUnorderedListEquals(expectedPartitionMetadata.racks(), actualPartitionMetadata.racks());
+                    }
+                }
+            }
+        } else {
+            assertEquals(expected.message(), actual.message());
+        }
+    }
+
+    private static Map<Uuid, Set<Integer>> fromTopicPartitions(
+        List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> assignment
+    ) {
+        Map<Uuid, Set<Integer>> assignmentMap = new HashMap<>();
+        assignment.forEach(topicPartitions ->
+            assignmentMap.put(topicPartitions.topicId(), new HashSet<>(topicPartitions.partitions()))
+        );
+        return assignmentMap;
+    }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index f529daa80e574..4a47fd3dc32a5 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -50,76 +50,43 @@
 import org.apache.kafka.common.message.SyncGroupRequestData.SyncGroupRequestAssignment;
 import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.metadata.PartitionRecord;
-import org.apache.kafka.common.metadata.RegisterBrokerRecord;
 import org.apache.kafka.common.metadata.RemoveTopicRecord;
 import org.apache.kafka.common.metadata.TopicRecord;
-import org.apache.kafka.common.network.ClientInformation;
-import org.apache.kafka.common.network.ListenerName;
-import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.ApiMessage;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.requests.RequestContext;
-import org.apache.kafka.common.requests.RequestHeader;
-import org.apache.kafka.common.security.auth.KafkaPrincipal;
-import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.LogContext;
-import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.coordinator.group.MockCoordinatorTimer.ExpiredTimeout;
 import org.apache.kafka.coordinator.group.MockCoordinatorTimer.ScheduledTimeout;
-import org.apache.kafka.coordinator.group.assignor.AssignmentSpec;
 import org.apache.kafka.coordinator.group.assignor.GroupAssignment;
 import org.apache.kafka.coordinator.group.assignor.MemberAssignment;
 import org.apache.kafka.coordinator.group.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.assignor.PartitionAssignorException;
-import org.apache.kafka.coordinator.group.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.coordinator.group.classic.ClassicGroupState;
 import org.apache.kafka.coordinator.group.consumer.Assignment;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataValue;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataValue;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberValue;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataKey;
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataValue;
-import org.apache.kafka.coordinator.group.generated.GroupMetadataKey;
 import org.apache.kafka.coordinator.group.generated.GroupMetadataValue;
 import org.apache.kafka.coordinator.group.classic.ClassicGroup;
 import org.apache.kafka.coordinator.group.classic.ClassicGroupMember;
-import org.apache.kafka.coordinator.group.classic.ClassicGroupState;
 import org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetricsShard;
 import org.apache.kafka.coordinator.group.runtime.CoordinatorResult;
 import org.apache.kafka.image.MetadataDelta;
 import org.apache.kafka.image.MetadataImage;
 import org.apache.kafka.image.MetadataProvenance;
-import org.apache.kafka.image.TopicImage;
-import org.apache.kafka.image.TopicsImage;
-import org.apache.kafka.server.common.ApiMessageAndVersion;
 import org.apache.kafka.server.common.MetadataVersion;
-import org.apache.kafka.timeline.SnapshotRegistry;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
-import org.opentest4j.AssertionFailedError;
 
-import java.net.InetAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -132,6 +99,10 @@
 import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
 import static org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
 import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID;
+import static org.apache.kafka.coordinator.group.Assertions.assertRecordEquals;
+import static org.apache.kafka.coordinator.group.Assertions.assertRecordsEquals;
+import static org.apache.kafka.coordinator.group.Assertions.assertResponseEquals;
+import static org.apache.kafka.coordinator.group.Assertions.assertUnorderedListEquals;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.appendGroupMetadataErrorToResponseError;
@@ -148,7 +119,6 @@
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.STABLE;
 import static org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetrics.CLASSIC_GROUP_COMPLETED_REBALANCES_SENSOR_NAME;
 import static org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetrics.CONSUMER_GROUP_REBALANCES_SENSOR_NAME;
-import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
@@ -164,1093 +134,6 @@
 import static org.mockito.Mockito.when;
 
 public class GroupMetadataManagerTest {
-    static class MockPartitionAssignor implements PartitionAssignor {
-        private final String name;
-        private GroupAssignment prepareGroupAssignment = null;
-
-        MockPartitionAssignor(String name) {
-            this.name = name;
-        }
-
-        public void prepareGroupAssignment(GroupAssignment prepareGroupAssignment) {
-            this.prepareGroupAssignment = prepareGroupAssignment;
-        }
-
-        @Override
-        public String name() {
-            return name;
-        }
-
-        @Override
-        public GroupAssignment assign(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) throws PartitionAssignorException {
-            return prepareGroupAssignment;
-        }
-    }
-
-    public static class MetadataImageBuilder {
-        private MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY);
-
-        public MetadataImageBuilder addTopic(
-            Uuid topicId,
-            String topicName,
-            int numPartitions
-        ) {
-            // For testing purposes, the following criteria are used:
-            // - Number of replicas for each partition: 2
-            // - Number of brokers available in the cluster: 4
-            delta.replay(new TopicRecord().setTopicId(topicId).setName(topicName));
-            for (int i = 0; i < numPartitions; i++) {
-                delta.replay(new PartitionRecord()
-                    .setTopicId(topicId)
-                    .setPartitionId(i)
-                    .setReplicas(Arrays.asList(i % 4, (i + 1) % 4)));
-            }
-            return this;
-        }
-
-        /**
-         * Add rack Ids for 4 broker Ids.
-         *
-         * For testing purposes, each broker is mapped
-         * to a rack Id with the same broker Id as a suffix.
-         */
-        public MetadataImageBuilder addRacks() {
-            for (int i = 0; i < 4; i++) {
-                delta.replay(new RegisterBrokerRecord().setBrokerId(i).setRack("rack" + i));
-            }
-            return this;
-        }
-
-        public MetadataImage build() {
-            return delta.apply(MetadataProvenance.EMPTY);
-        }
-    }
-
-    static class ConsumerGroupBuilder {
-        private final String groupId;
-        private final int groupEpoch;
-        private int assignmentEpoch;
-        private final Map<String, ConsumerGroupMember> members = new HashMap<>();
-        private final Map<String, Assignment> assignments = new HashMap<>();
-        private Map<String, TopicMetadata> subscriptionMetadata;
-
-        public ConsumerGroupBuilder(String groupId, int groupEpoch) {
-            this.groupId = groupId;
-            this.groupEpoch = groupEpoch;
-            this.assignmentEpoch = 0;
-        }
-
-        public ConsumerGroupBuilder withMember(ConsumerGroupMember member) {
-            this.members.put(member.memberId(), member);
-            return this;
-        }
-
-        public ConsumerGroupBuilder withSubscriptionMetadata(Map<String, TopicMetadata> subscriptionMetadata) {
-            this.subscriptionMetadata = subscriptionMetadata;
-            return this;
-        }
-
-        public ConsumerGroupBuilder withAssignment(String memberId, Map<Uuid, Set<Integer>> assignment) {
-            this.assignments.put(memberId, new Assignment(assignment));
-            return this;
-        }
-
-        public ConsumerGroupBuilder withAssignmentEpoch(int assignmentEpoch) {
-            this.assignmentEpoch = assignmentEpoch;
-            return this;
-        }
-
-        public List<Record> build(TopicsImage topicsImage) {
-            List<Record> records = new ArrayList<>();
-
-            // Add subscription records for members.
-            members.forEach((memberId, member) -> {
-                records.add(RecordHelpers.newMemberSubscriptionRecord(groupId, member));
-            });
-
-            // Add subscription metadata.
-            if (subscriptionMetadata == null) {
-                subscriptionMetadata = new HashMap<>();
-                members.forEach((memberId, member) -> {
-                    member.subscribedTopicNames().forEach(topicName -> {
-                        TopicImage topicImage = topicsImage.getTopic(topicName);
-                        if (topicImage != null) {
-                            subscriptionMetadata.put(topicName, new TopicMetadata(
-                                topicImage.id(),
-                                topicImage.name(),
-                                topicImage.partitions().size(),
-                                Collections.emptyMap()));
-                        }
-                    });
-                });
-            }
-
-            if (!subscriptionMetadata.isEmpty()) {
-                records.add(RecordHelpers.newGroupSubscriptionMetadataRecord(groupId, subscriptionMetadata));
-            }
-
-            // Add group epoch record.
-            records.add(RecordHelpers.newGroupEpochRecord(groupId, groupEpoch));
-
-            // Add target assignment records.
-            assignments.forEach((memberId, assignment) -> {
-                records.add(RecordHelpers.newTargetAssignmentRecord(groupId, memberId, assignment.partitions()));
-            });
-
-            // Add target assignment epoch.
-            records.add(RecordHelpers.newTargetAssignmentEpochRecord(groupId, assignmentEpoch));
-
-            // Add current assignment records for members.
-            members.forEach((memberId, member) -> {
-                records.add(RecordHelpers.newCurrentAssignmentRecord(groupId, member));
-            });
-
-            return records;
-        }
-    }
-
-    static class GroupMetadataManagerTestContext {
-        static class Builder {
-            final private MockTime time = new MockTime();
-            final private MockCoordinatorTimer<Void, Record> timer = new MockCoordinatorTimer<>(time);
-            final private LogContext logContext = new LogContext();
-            final private SnapshotRegistry snapshotRegistry = new SnapshotRegistry(logContext);
-            final private TopicPartition groupMetadataTopicPartition = new TopicPartition("topic", 0);
-            private MetadataImage metadataImage;
-            private List<PartitionAssignor> consumerGroupAssignors = Collections.singletonList(new MockPartitionAssignor("range"));
-            private List<ConsumerGroupBuilder> consumerGroupBuilders = new ArrayList<>();
-            private int consumerGroupMaxSize = Integer.MAX_VALUE;
-            private int consumerGroupMetadataRefreshIntervalMs = Integer.MAX_VALUE;
-            private int classicGroupMaxSize = Integer.MAX_VALUE;
-            private int classicGroupInitialRebalanceDelayMs = 3000;
-            final private int classicGroupNewMemberJoinTimeoutMs = 5 * 60 * 1000;
-            private int classicGroupMinSessionTimeoutMs = 10;
-            private int classicGroupMaxSessionTimeoutMs = 10 * 60 * 1000;
-            private GroupCoordinatorMetricsShard metrics = mock(GroupCoordinatorMetricsShard.class);
-
-            public Builder withMetadataImage(MetadataImage metadataImage) {
-                this.metadataImage = metadataImage;
-                return this;
-            }
-
-            public Builder withAssignors(List<PartitionAssignor> assignors) {
-                this.consumerGroupAssignors = assignors;
-                return this;
-            }
-
-            public Builder withConsumerGroup(ConsumerGroupBuilder builder) {
-                this.consumerGroupBuilders.add(builder);
-                return this;
-            }
-
-            public Builder withConsumerGroupMaxSize(int consumerGroupMaxSize) {
-                this.consumerGroupMaxSize = consumerGroupMaxSize;
-                return this;
-            }
-
-            public Builder withConsumerGroupMetadataRefreshIntervalMs(int consumerGroupMetadataRefreshIntervalMs) {
-                this.consumerGroupMetadataRefreshIntervalMs = consumerGroupMetadataRefreshIntervalMs;
-                return this;
-            }
-
-            public Builder withClassicGroupMaxSize(int classicGroupMaxSize) {
-                this.classicGroupMaxSize = classicGroupMaxSize;
-                return this;
-            }
-
-            public Builder withClassicGroupInitialRebalanceDelayMs(int classicGroupInitialRebalanceDelayMs) {
-                this.classicGroupInitialRebalanceDelayMs = classicGroupInitialRebalanceDelayMs;
-                return this;
-            }
-
-            public Builder withClassicGroupMinSessionTimeoutMs(int classicGroupMinSessionTimeoutMs) {
-                this.classicGroupMinSessionTimeoutMs = classicGroupMinSessionTimeoutMs;
-                return this;
-            }
-
-            public Builder withClassicGroupMaxSessionTimeoutMs(int classicGroupMaxSessionTimeoutMs) {
-                this.classicGroupMaxSessionTimeoutMs = classicGroupMaxSessionTimeoutMs;
-                return this;
-            }
-
-            public GroupMetadataManagerTestContext build() {
-                if (metadataImage == null) metadataImage = MetadataImage.EMPTY;
-                if (consumerGroupAssignors == null) consumerGroupAssignors = Collections.emptyList();
-
-                GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext(
-                    time,
-                    timer,
-                    snapshotRegistry,
-                    metrics,
-                    new GroupMetadataManager.Builder()
-                        .withSnapshotRegistry(snapshotRegistry)
-                        .withLogContext(logContext)
-                        .withTime(time)
-                        .withTimer(timer)
-                        .withMetadataImage(metadataImage)
-                        .withConsumerGroupHeartbeatInterval(5000)
-                        .withConsumerGroupSessionTimeout(45000)
-                        .withConsumerGroupMaxSize(consumerGroupMaxSize)
-                        .withConsumerGroupAssignors(consumerGroupAssignors)
-                        .withConsumerGroupMetadataRefreshIntervalMs(consumerGroupMetadataRefreshIntervalMs)
-                        .withClassicGroupMaxSize(classicGroupMaxSize)
-                        .withClassicGroupMinSessionTimeoutMs(classicGroupMinSessionTimeoutMs)
-                        .withClassicGroupMaxSessionTimeoutMs(classicGroupMaxSessionTimeoutMs)
-                        .withClassicGroupInitialRebalanceDelayMs(classicGroupInitialRebalanceDelayMs)
-                        .withClassicGroupNewMemberJoinTimeoutMs(classicGroupNewMemberJoinTimeoutMs)
-                        .withGroupCoordinatorMetricsShard(metrics)
-                        .build(),
-                    classicGroupInitialRebalanceDelayMs,
-                    classicGroupNewMemberJoinTimeoutMs
-                );
-
-                consumerGroupBuilders.forEach(builder -> {
-                    builder.build(metadataImage.topics()).forEach(context::replay);
-                });
-
-                context.commit();
-
-                return context;
-            }
-        }
-
-        final MockTime time;
-        final MockCoordinatorTimer<Void, Record> timer;
-        final SnapshotRegistry snapshotRegistry;
-        final GroupCoordinatorMetricsShard metrics;
-        final GroupMetadataManager groupMetadataManager;
-        final int classicGroupInitialRebalanceDelayMs;
-        final int classicGroupNewMemberJoinTimeoutMs;
-
-        long lastCommittedOffset = 0L;
-        long lastWrittenOffset = 0L;
-
-        public GroupMetadataManagerTestContext(
-            MockTime time,
-            MockCoordinatorTimer<Void, Record> timer,
-            SnapshotRegistry snapshotRegistry,
-            GroupCoordinatorMetricsShard metrics,
-            GroupMetadataManager groupMetadataManager,
-            int classicGroupInitialRebalanceDelayMs,
-            int classicGroupNewMemberJoinTimeoutMs
-        ) {
-            this.time = time;
-            this.timer = timer;
-            this.snapshotRegistry = snapshotRegistry;
-            this.metrics = metrics;
-            this.groupMetadataManager = groupMetadataManager;
-            this.classicGroupInitialRebalanceDelayMs = classicGroupInitialRebalanceDelayMs;
-            this.classicGroupNewMemberJoinTimeoutMs = classicGroupNewMemberJoinTimeoutMs;
-            snapshotRegistry.getOrCreateSnapshot(lastWrittenOffset);
-        }
-
-        public void commit() {
-            long lastCommittedOffset = this.lastCommittedOffset;
-            this.lastCommittedOffset = lastWrittenOffset;
-            snapshotRegistry.deleteSnapshotsUpTo(lastCommittedOffset);
-        }
-
-        public void rollback() {
-            lastWrittenOffset = lastCommittedOffset;
-            snapshotRegistry.revertToSnapshot(lastCommittedOffset);
-        }
-
-        public ConsumerGroup.ConsumerGroupState consumerGroupState(
-            String groupId
-        ) {
-            return groupMetadataManager
-                .getOrMaybeCreateConsumerGroup(groupId, false)
-                .state();
-        }
-
-        public ConsumerGroupMember.MemberState consumerGroupMemberState(
-            String groupId,
-            String memberId
-        ) {
-            return groupMetadataManager
-                .getOrMaybeCreateConsumerGroup(groupId, false)
-                .getOrMaybeCreateMember(memberId, false)
-                .state();
-        }
-
-        public CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGroupHeartbeat(
-            ConsumerGroupHeartbeatRequestData request
-        ) {
-            RequestContext context = new RequestContext(
-                new RequestHeader(
-                    ApiKeys.CONSUMER_GROUP_HEARTBEAT,
-                    ApiKeys.CONSUMER_GROUP_HEARTBEAT.latestVersion(),
-                    "client",
-                    0
-                ),
-                "1",
-                InetAddress.getLoopbackAddress(),
-                KafkaPrincipal.ANONYMOUS,
-                ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
-                SecurityProtocol.PLAINTEXT,
-                ClientInformation.EMPTY,
-                false
-            );
-
-            CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result = groupMetadataManager.consumerGroupHeartbeat(
-                context,
-                request
-            );
-
-            result.records().forEach(this::replay);
-            return result;
-        }
-
-        public List<ExpiredTimeout<Void, Record>> sleep(long ms) {
-            time.sleep(ms);
-            List<ExpiredTimeout<Void, Record>> timeouts = timer.poll();
-            timeouts.forEach(timeout -> {
-                if (timeout.result.replayRecords()) {
-                    timeout.result.records().forEach(this::replay);
-                }
-            });
-            return timeouts;
-        }
-
-        public ScheduledTimeout<Void, Record> assertSessionTimeout(
-            String groupId,
-            String memberId,
-            long delayMs
-        ) {
-            ScheduledTimeout<Void, Record> timeout =
-                timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
-            assertNotNull(timeout);
-            assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
-            return timeout;
-        }
-
-        public void assertNoSessionTimeout(
-            String groupId,
-            String memberId
-        ) {
-            ScheduledTimeout<Void, Record> timeout =
-                timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
-            assertNull(timeout);
-        }
-
-        public ScheduledTimeout<Void, Record> assertRevocationTimeout(
-            String groupId,
-            String memberId,
-            long delayMs
-        ) {
-            ScheduledTimeout<Void, Record> timeout =
-                timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
-            assertNotNull(timeout);
-            assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
-            return timeout;
-        }
-
-        public void assertNoRevocationTimeout(
-            String groupId,
-            String memberId
-        ) {
-            ScheduledTimeout<Void, Record> timeout =
-                timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
-            assertNull(timeout);
-        }
-
-        ClassicGroup createClassicGroup(String groupId) {
-            return groupMetadataManager.getOrMaybeCreateClassicGroup(groupId, true);
-        }
-
-        public JoinResult sendClassicGroupJoin(
-            JoinGroupRequestData request
-        ) {
-            return sendClassicGroupJoin(request, false);
-        }
-
-        public JoinResult sendClassicGroupJoin(
-            JoinGroupRequestData request,
-            boolean requireKnownMemberId
-        ) {
-            return sendClassicGroupJoin(request, requireKnownMemberId, false);
-        }
-
-        public JoinResult sendClassicGroupJoin(
-            JoinGroupRequestData request,
-            boolean requireKnownMemberId,
-            boolean supportSkippingAssignment
-        ) {
-            // requireKnownMemberId is true: version >= 4 (See JoinGroupRequest#requiresKnownMemberId())
-            // supportSkippingAssignment is true: version >= 9 (See JoinGroupRequest#supportsSkippingAssignment())
-            short joinGroupVersion = 3;
-
-            if (requireKnownMemberId) {
-                joinGroupVersion = 4;
-                if (supportSkippingAssignment) {
-                    joinGroupVersion = ApiKeys.JOIN_GROUP.latestVersion();
-                }
-            }
-
-            RequestContext context = new RequestContext(
-                new RequestHeader(
-                    ApiKeys.JOIN_GROUP,
-                    joinGroupVersion,
-                    "client",
-                    0
-                ),
-                "1",
-                InetAddress.getLoopbackAddress(),
-                KafkaPrincipal.ANONYMOUS,
-                ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
-                SecurityProtocol.PLAINTEXT,
-                ClientInformation.EMPTY,
-                false
-            );
-
-            CompletableFuture<JoinGroupResponseData> responseFuture = new CompletableFuture<>();
-            CoordinatorResult<Void, Record> coordinatorResult = groupMetadataManager.classicGroupJoin(
-                context,
-                request,
-                responseFuture
-            );
-
-            return new JoinResult(responseFuture, coordinatorResult);
-        }
-
-        public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteRebalance(
-            String groupId
-        ) throws Exception {
-            ClassicGroup group = createClassicGroup(groupId);
-
-            JoinGroupResponseData leaderJoinResponse =
-                joinClassicGroupAsDynamicMemberAndCompleteJoin(new JoinGroupRequestBuilder()
-                    .withGroupId("group-id")
-                    .withMemberId(UNKNOWN_MEMBER_ID)
-                    .withDefaultProtocolTypeAndProtocols()
-                    .withRebalanceTimeoutMs(10000)
-                    .withSessionTimeoutMs(5000)
-                    .build());
-
-            assertEquals(1, leaderJoinResponse.generationId());
-            assertTrue(group.isInState(COMPLETING_REBALANCE));
-
-            SyncResult syncResult = sendClassicGroupSync(new SyncGroupRequestBuilder()
-                .withGroupId("group-id")
-                .withMemberId(leaderJoinResponse.memberId())
-                .withGenerationId(leaderJoinResponse.generationId())
-                .build());
-
-            assertEquals(
-                Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
-                syncResult.records
-            );
-            // Simulate a successful write to the log.
-            syncResult.appendFuture.complete(null);
-
-            assertTrue(syncResult.syncFuture.isDone());
-            assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
-            assertTrue(group.isInState(STABLE));
-
-            return leaderJoinResponse;
-        }
-
-        public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteJoin(
-            JoinGroupRequestData request
-        ) throws ExecutionException, InterruptedException {
-            boolean requireKnownMemberId = true;
-            String newMemberId = request.memberId();
-
-            if (request.memberId().equals(UNKNOWN_MEMBER_ID)) {
-                // Since member id is required, we need another round to get the successful join group result.
-                JoinResult firstJoinResult = sendClassicGroupJoin(
-                    request,
-                    requireKnownMemberId
-                );
-                assertTrue(firstJoinResult.records.isEmpty());
-                assertTrue(firstJoinResult.joinFuture.isDone());
-                assertEquals(Errors.MEMBER_ID_REQUIRED.code(), firstJoinResult.joinFuture.get().errorCode());
-                newMemberId = firstJoinResult.joinFuture.get().memberId();
-            }
-
-            // Second round
-            JoinGroupRequestData secondRequest = new JoinGroupRequestData()
-                .setGroupId(request.groupId())
-                .setMemberId(newMemberId)
-                .setProtocolType(request.protocolType())
-                .setProtocols(request.protocols())
-                .setSessionTimeoutMs(request.sessionTimeoutMs())
-                .setRebalanceTimeoutMs(request.rebalanceTimeoutMs())
-                .setReason(request.reason());
-
-            JoinResult secondJoinResult = sendClassicGroupJoin(
-                secondRequest,
-                requireKnownMemberId
-            );
-
-            assertTrue(secondJoinResult.records.isEmpty());
-            List<ExpiredTimeout<Void, Record>> timeouts = sleep(classicGroupInitialRebalanceDelayMs);
-            assertEquals(1, timeouts.size());
-            assertTrue(secondJoinResult.joinFuture.isDone());
-            assertEquals(Errors.NONE.code(), secondJoinResult.joinFuture.get().errorCode());
-
-            return secondJoinResult.joinFuture.get();
-        }
-
-        public JoinGroupResponseData joinClassicGroupAndCompleteJoin(
-            JoinGroupRequestData request,
-            boolean requireKnownMemberId,
-            boolean supportSkippingAssignment
-        ) throws ExecutionException, InterruptedException {
-            return joinClassicGroupAndCompleteJoin(
-                request,
-                requireKnownMemberId,
-                supportSkippingAssignment,
-                classicGroupInitialRebalanceDelayMs
-            );
-        }
-
-        public JoinGroupResponseData joinClassicGroupAndCompleteJoin(
-            JoinGroupRequestData request,
-            boolean requireKnownMemberId,
-            boolean supportSkippingAssignment,
-            int advanceClockMs
-        ) throws ExecutionException, InterruptedException {
-            if (requireKnownMemberId && request.groupInstanceId().isEmpty()) {
-                return joinClassicGroupAsDynamicMemberAndCompleteJoin(request);
-            }
-
-            try {
-                JoinResult joinResult = sendClassicGroupJoin(
-                    request,
-                    requireKnownMemberId,
-                    supportSkippingAssignment
-                );
-
-                sleep(advanceClockMs);
-                assertTrue(joinResult.joinFuture.isDone());
-                assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
-                return joinResult.joinFuture.get();
-            } catch (Exception e) {
-                fail("Failed to due: " + e.getMessage());
-            }
-            return null;
-        }
-
-        public SyncResult sendClassicGroupSync(SyncGroupRequestData request) {
-            RequestContext context = new RequestContext(
-                new RequestHeader(
-                    ApiKeys.SYNC_GROUP,
-                    ApiKeys.SYNC_GROUP.latestVersion(),
-                    "client",
-                    0
-                ),
-                "1",
-                InetAddress.getLoopbackAddress(),
-                KafkaPrincipal.ANONYMOUS,
-                ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
-                SecurityProtocol.PLAINTEXT,
-                ClientInformation.EMPTY,
-                false
-            );
-
-            CompletableFuture<SyncGroupResponseData> responseFuture = new CompletableFuture<>();
-
-            CoordinatorResult<Void, Record> coordinatorResult = groupMetadataManager.classicGroupSync(
-                context,
-                request,
-                responseFuture
-            );
-
-            return new SyncResult(responseFuture, coordinatorResult);
-        }
-
-        public RebalanceResult staticMembersJoinAndRebalance(
-            String groupId,
-            String leaderInstanceId,
-            String followerInstanceId
-        ) throws Exception {
-            return staticMembersJoinAndRebalance(
-                groupId,
-                leaderInstanceId,
-                followerInstanceId,
-                10000,
-                5000
-            );
-        }
-
-        public RebalanceResult staticMembersJoinAndRebalance(
-            String groupId,
-            String leaderInstanceId,
-            String followerInstanceId,
-            int rebalanceTimeoutMs,
-            int sessionTimeoutMs
-        ) throws Exception {
-            ClassicGroup group = createClassicGroup("group-id");
-
-            JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
-                .withGroupId(groupId)
-                .withGroupInstanceId(leaderInstanceId)
-                .withMemberId(UNKNOWN_MEMBER_ID)
-                .withProtocolType("consumer")
-                .withProtocolSuperset()
-                .withRebalanceTimeoutMs(rebalanceTimeoutMs)
-                .withSessionTimeoutMs(sessionTimeoutMs)
-                .build();
-
-            JoinResult leaderJoinResult = sendClassicGroupJoin(joinRequest);
-            JoinResult followerJoinResult = sendClassicGroupJoin(joinRequest.setGroupInstanceId(followerInstanceId));
-
-            assertTrue(leaderJoinResult.records.isEmpty());
-            assertTrue(followerJoinResult.records.isEmpty());
-            assertFalse(leaderJoinResult.joinFuture.isDone());
-            assertFalse(followerJoinResult.joinFuture.isDone());
-
-            // The goal for two timer advance is to let first group initial join complete and set newMemberAdded flag to false. Next advance is
-            // to trigger the rebalance as needed for follower delayed join. One large time advance won't help because we could only populate one
-            // delayed join from purgatory and the new delayed op is created at that time and never be triggered.
-            assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
-            assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
-
-            assertTrue(leaderJoinResult.joinFuture.isDone());
-            assertTrue(followerJoinResult.joinFuture.isDone());
-            assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
-            assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
-            assertEquals(1, leaderJoinResult.joinFuture.get().generationId());
-            assertEquals(1, followerJoinResult.joinFuture.get().generationId());
-            assertEquals(2, group.size());
-            assertEquals(1, group.generationId());
-            assertTrue(group.isInState(COMPLETING_REBALANCE));
-
-            String leaderId = leaderJoinResult.joinFuture.get().memberId();
-            String followerId = followerJoinResult.joinFuture.get().memberId();
-            List<SyncGroupRequestAssignment> assignment = new ArrayList<>();
-            assignment.add(new SyncGroupRequestAssignment().setMemberId(leaderId)
-                .setAssignment(new byte[]{1}));
-            assignment.add(new SyncGroupRequestAssignment().setMemberId(followerId)
-                .setAssignment(new byte[]{2}));
-
-            SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
-                .withGroupId(groupId)
-                .withGroupInstanceId(leaderInstanceId)
-                .withMemberId(leaderId)
-                .withGenerationId(1)
-                .withAssignment(assignment)
-                .build();
-
-            SyncResult leaderSyncResult = sendClassicGroupSync(syncRequest);
-
-            // The generated record should contain the new assignment.
-            Map<String, byte[]> groupAssignment = assignment.stream().collect(Collectors.toMap(
-                SyncGroupRequestAssignment::memberId, SyncGroupRequestAssignment::assignment
-            ));
-            assertEquals(
-                Collections.singletonList(
-                    RecordHelpers.newGroupMetadataRecord(group, groupAssignment, MetadataVersion.latestTesting())),
-                leaderSyncResult.records
-            );
-
-            // Simulate a successful write to the log.
-            leaderSyncResult.appendFuture.complete(null);
-
-            assertTrue(leaderSyncResult.syncFuture.isDone());
-            assertEquals(Errors.NONE.code(), leaderSyncResult.syncFuture.get().errorCode());
-            assertTrue(group.isInState(STABLE));
-
-            SyncResult followerSyncResult = sendClassicGroupSync(
-                syncRequest.setGroupInstanceId(followerInstanceId)
-                    .setMemberId(followerId)
-                    .setAssignments(Collections.emptyList())
-            );
-
-            assertTrue(followerSyncResult.records.isEmpty());
-            assertTrue(followerSyncResult.syncFuture.isDone());
-            assertEquals(Errors.NONE.code(), followerSyncResult.syncFuture.get().errorCode());
-            assertTrue(group.isInState(STABLE));
-
-            assertEquals(2, group.size());
-            assertEquals(1, group.generationId());
-
-            return new RebalanceResult(
-                1,
-                leaderId,
-                leaderSyncResult.syncFuture.get().assignment(),
-                followerId,
-                followerSyncResult.syncFuture.get().assignment()
-            );
-        }
-
-        public PendingMemberGroupResult setupGroupWithPendingMember(ClassicGroup group) throws Exception {
-            // Add the first member
-            JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
-                .withGroupId("group-id")
-                .withMemberId(UNKNOWN_MEMBER_ID)
-                .withDefaultProtocolTypeAndProtocols()
-                .withRebalanceTimeoutMs(10000)
-                .withSessionTimeoutMs(5000)
-                .build();
-
-            JoinGroupResponseData leaderJoinResponse =
-                joinClassicGroupAsDynamicMemberAndCompleteJoin(joinRequest);
-
-            List<SyncGroupRequestAssignment> assignment = new ArrayList<>();
-            assignment.add(new SyncGroupRequestAssignment().setMemberId(leaderJoinResponse.memberId()));
-            SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
-                .withGroupId("group-id")
-                .withMemberId(leaderJoinResponse.memberId())
-                .withGenerationId(leaderJoinResponse.generationId())
-                .withAssignment(assignment)
-                .build();
-
-            SyncResult syncResult = sendClassicGroupSync(syncRequest);
-
-            // Now the group is stable, with the one member that joined above
-            assertEquals(
-                Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
-                syncResult.records
-            );
-            // Simulate a successful write to log.
-            syncResult.appendFuture.complete(null);
-
-            assertTrue(syncResult.syncFuture.isDone());
-            assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
-
-            // Start the join for the second member
-            JoinResult followerJoinResult = sendClassicGroupJoin(
-                joinRequest.setMemberId(UNKNOWN_MEMBER_ID)
-            );
-
-            assertTrue(followerJoinResult.records.isEmpty());
-            assertFalse(followerJoinResult.joinFuture.isDone());
-
-            JoinResult leaderJoinResult = sendClassicGroupJoin(
-                joinRequest.setMemberId(leaderJoinResponse.memberId())
-            );
-
-            assertTrue(leaderJoinResult.records.isEmpty());
-            assertTrue(group.isInState(COMPLETING_REBALANCE));
-            assertTrue(leaderJoinResult.joinFuture.isDone());
-            assertTrue(followerJoinResult.joinFuture.isDone());
-            assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
-            assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
-            assertEquals(leaderJoinResult.joinFuture.get().generationId(), followerJoinResult.joinFuture.get().generationId());
-            assertEquals(leaderJoinResponse.memberId(), leaderJoinResult.joinFuture.get().leader());
-            assertEquals(leaderJoinResponse.memberId(), followerJoinResult.joinFuture.get().leader());
-
-            int nextGenerationId = leaderJoinResult.joinFuture.get().generationId();
-            String followerId = followerJoinResult.joinFuture.get().memberId();
-
-            // Stabilize the group
-            syncResult = sendClassicGroupSync(syncRequest.setGenerationId(nextGenerationId));
-
-            assertEquals(
-                Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
-                syncResult.records
-            );
-            // Simulate a successful write to log.
-            syncResult.appendFuture.complete(null);
-
-            assertTrue(syncResult.syncFuture.isDone());
-            assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
-            assertTrue(group.isInState(STABLE));
-
-            // Re-join an existing member, to transition the group to PreparingRebalance state.
-            leaderJoinResult = sendClassicGroupJoin(
-                joinRequest.setMemberId(leaderJoinResponse.memberId()));
-
-            assertTrue(leaderJoinResult.records.isEmpty());
-            assertFalse(leaderJoinResult.joinFuture.isDone());
-            assertTrue(group.isInState(PREPARING_REBALANCE));
-
-            // Create a pending member in the group
-            JoinResult pendingMemberJoinResult = sendClassicGroupJoin(
-                joinRequest
-                    .setMemberId(UNKNOWN_MEMBER_ID)
-                    .setSessionTimeoutMs(2500),
-                true
-            );
-
-            assertTrue(pendingMemberJoinResult.records.isEmpty());
-            assertTrue(pendingMemberJoinResult.joinFuture.isDone());
-            assertEquals(Errors.MEMBER_ID_REQUIRED.code(), pendingMemberJoinResult.joinFuture.get().errorCode());
-            assertEquals(1, group.numPendingJoinMembers());
-
-            // Re-join the second existing member
-            followerJoinResult = sendClassicGroupJoin(
-                joinRequest.setMemberId(followerId).setSessionTimeoutMs(5000)
-            );
-
-            assertTrue(followerJoinResult.records.isEmpty());
-            assertFalse(followerJoinResult.joinFuture.isDone());
-            assertTrue(group.isInState(PREPARING_REBALANCE));
-            assertEquals(2, group.size());
-            assertEquals(1, group.numPendingJoinMembers());
-
-            return new PendingMemberGroupResult(
-                leaderJoinResponse.memberId(),
-                followerId,
-                pendingMemberJoinResult.joinFuture.get()
-            );
-        }
-
-        public void verifySessionExpiration(ClassicGroup group, int timeoutMs) {
-            Set<String> expectedHeartbeatKeys = group.allMembers().stream()
-                                                     .map(member -> classicGroupHeartbeatKey(group.groupId(), member.memberId())).collect(Collectors.toSet());
-
-            // Member should be removed as session expires.
-            List<ExpiredTimeout<Void, Record>> timeouts = sleep(timeoutMs);
-            List<Record> expectedRecords = Collections.singletonList(newGroupMetadataRecord(
-                group.groupId(),
-                new GroupMetadataValue()
-                    .setMembers(Collections.emptyList())
-                    .setGeneration(group.generationId())
-                    .setLeader(null)
-                    .setProtocolType("consumer")
-                    .setProtocol(null)
-                    .setCurrentStateTimestamp(time.milliseconds()),
-                MetadataVersion.latestTesting()));
-
-
-            Set<String> heartbeatKeys = timeouts.stream().map(timeout -> timeout.key).collect(Collectors.toSet());
-            assertEquals(expectedHeartbeatKeys, heartbeatKeys);
-
-            // Only the last member leaving the group should result in the empty group metadata record.
-            int timeoutsSize = timeouts.size();
-            assertEquals(expectedRecords, timeouts.get(timeoutsSize - 1).result.records());
-            assertNoOrEmptyResult(timeouts.subList(0, timeoutsSize - 1));
-            assertTrue(group.isInState(EMPTY));
-            assertEquals(0, group.size());
-        }
-
-        public HeartbeatResponseData sendClassicGroupHeartbeat(
-            HeartbeatRequestData request
-        ) {
-            RequestContext context = new RequestContext(
-                new RequestHeader(
-                    ApiKeys.HEARTBEAT,
-                    ApiKeys.HEARTBEAT.latestVersion(),
-                    "client",
-                    0
-                ),
-                "1",
-                InetAddress.getLoopbackAddress(),
-                KafkaPrincipal.ANONYMOUS,
-                ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
-                SecurityProtocol.PLAINTEXT,
-                ClientInformation.EMPTY,
-                false
-            );
-
-            return groupMetadataManager.classicGroupHeartbeat(
-                context,
-                request
-            );
-        }
-
-        public List<ListGroupsResponseData.ListedGroup> sendListGroups(List<String> statesFilter, List<String> typesFilter) {
-            Set<String> statesFilterSet = statesFilter.stream().collect(Collectors.toSet());
-            Set<String> typesFilterSet = typesFilter.stream().collect(Collectors.toSet());
-            return groupMetadataManager.listGroups(statesFilterSet, typesFilterSet, lastCommittedOffset);
-        }
-
-        public List<ConsumerGroupDescribeResponseData.DescribedGroup> sendConsumerGroupDescribe(List<String> groupIds) {
-            return groupMetadataManager.consumerGroupDescribe(groupIds, lastCommittedOffset);
-        }
-
-        public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(List<String> groupIds) {
-            return groupMetadataManager.describeGroups(groupIds, lastCommittedOffset);
-        }
-
-        public void verifyHeartbeat(
-            String groupId,
-            JoinGroupResponseData joinResponse,
-            Errors expectedError
-        ) {
-            HeartbeatRequestData request = new HeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(joinResponse.memberId())
-                .setGenerationId(joinResponse.generationId());
-
-            if (expectedError == Errors.UNKNOWN_MEMBER_ID) {
-                assertThrows(UnknownMemberIdException.class, () -> sendClassicGroupHeartbeat(request));
-            } else {
-                HeartbeatResponseData response = sendClassicGroupHeartbeat(request);
-                assertEquals(expectedError.code(), response.errorCode());
-            }
-        }
-
-        public List<JoinGroupResponseData> joinWithNMembers(
-            String groupId,
-            int numMembers,
-            int rebalanceTimeoutMs,
-            int sessionTimeoutMs
-        ) {
-            ClassicGroup group = createClassicGroup(groupId);
-            boolean requireKnownMemberId = true;
-
-            // First join requests
-            JoinGroupRequestData request = new JoinGroupRequestBuilder()
-                .withGroupId(groupId)
-                .withMemberId(UNKNOWN_MEMBER_ID)
-                .withDefaultProtocolTypeAndProtocols()
-                .withRebalanceTimeoutMs(rebalanceTimeoutMs)
-                .withSessionTimeoutMs(sessionTimeoutMs)
-                .build();
-
-            List<String> memberIds = IntStream.range(0, numMembers).mapToObj(i -> {
-                JoinResult joinResult = sendClassicGroupJoin(request, requireKnownMemberId);
-
-                assertTrue(joinResult.records.isEmpty());
-                assertTrue(joinResult.joinFuture.isDone());
-
-                try {
-                    return joinResult.joinFuture.get().memberId();
-                } catch (Exception e) {
-                    fail("Unexpected exception: " + e.getMessage());
-                }
-                return null;
-            }).collect(Collectors.toList());
-
-            // Second join requests
-            List<CompletableFuture<JoinGroupResponseData>> secondJoinFutures = IntStream.range(0, numMembers).mapToObj(i -> {
-                JoinResult joinResult = sendClassicGroupJoin(request.setMemberId(memberIds.get(i)), requireKnownMemberId);
-
-                assertTrue(joinResult.records.isEmpty());
-                assertFalse(joinResult.joinFuture.isDone());
-
-                return joinResult.joinFuture;
-            }).collect(Collectors.toList());
-
-            // Advance clock by initial rebalance delay.
-            assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
-            secondJoinFutures.forEach(future -> assertFalse(future.isDone()));
-            // Advance clock by rebalance timeout to complete join phase.
-            assertNoOrEmptyResult(sleep(rebalanceTimeoutMs));
-
-            List<JoinGroupResponseData> joinResponses = secondJoinFutures.stream().map(future -> {
-                assertTrue(future.isDone());
-                try {
-                    assertEquals(Errors.NONE.code(), future.get().errorCode());
-                    return future.get();
-                } catch (Exception e) {
-                    fail("Unexpected exception: " + e.getMessage());
-                }
-                return null;
-            }).collect(Collectors.toList());
-
-            assertEquals(numMembers, group.size());
-            assertTrue(group.isInState(COMPLETING_REBALANCE));
-
-            return joinResponses;
-        }
-
-        public CoordinatorResult<LeaveGroupResponseData, Record> sendClassicGroupLeave(
-            LeaveGroupRequestData request
-        ) {
-            RequestContext context = new RequestContext(
-                new RequestHeader(
-                    ApiKeys.LEAVE_GROUP,
-                    ApiKeys.LEAVE_GROUP.latestVersion(),
-                    "client",
-                    0
-                ),
-                "1",
-                InetAddress.getLoopbackAddress(),
-                KafkaPrincipal.ANONYMOUS,
-                ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
-                SecurityProtocol.PLAINTEXT,
-                ClientInformation.EMPTY,
-                false
-            );
-
-            return groupMetadataManager.classicGroupLeave(context, request);
-        }
-
-        private void verifyDescribeGroupsReturnsDeadGroup(String groupId) {
-            List<DescribeGroupsResponseData.DescribedGroup> describedGroups =
-                describeGroups(Collections.singletonList(groupId));
-
-            assertEquals(
-                Collections.singletonList(new DescribeGroupsResponseData.DescribedGroup()
-                    .setGroupId("group-id")
-                    .setGroupState(DEAD.toString())
-                ),
-                describedGroups
-            );
-        }
-
-        private ApiMessage messageOrNull(ApiMessageAndVersion apiMessageAndVersion) {
-            if (apiMessageAndVersion == null) {
-                return null;
-            } else {
-                return apiMessageAndVersion.message();
-            }
-        }
-
-        private void replay(
-            Record record
-        ) {
-            ApiMessageAndVersion key = record.key();
-            ApiMessageAndVersion value = record.value();
-
-            if (key == null) {
-                throw new IllegalStateException("Received a null key in " + record);
-            }
-
-            switch (key.version()) {
-                case GroupMetadataKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (GroupMetadataKey) key.message(),
-                        (GroupMetadataValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupMemberMetadataKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupMemberMetadataKey) key.message(),
-                        (ConsumerGroupMemberMetadataValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupMetadataKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupMetadataKey) key.message(),
-                        (ConsumerGroupMetadataValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupPartitionMetadataKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupPartitionMetadataKey) key.message(),
-                        (ConsumerGroupPartitionMetadataValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupTargetAssignmentMemberKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupTargetAssignmentMemberKey) key.message(),
-                        (ConsumerGroupTargetAssignmentMemberValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupTargetAssignmentMetadataKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupTargetAssignmentMetadataKey) key.message(),
-                        (ConsumerGroupTargetAssignmentMetadataValue) messageOrNull(value)
-                    );
-                    break;
-
-                case ConsumerGroupCurrentMemberAssignmentKey.HIGHEST_SUPPORTED_VERSION:
-                    groupMetadataManager.replay(
-                        (ConsumerGroupCurrentMemberAssignmentKey) key.message(),
-                        (ConsumerGroupCurrentMemberAssignmentValue) messageOrNull(value)
-                    );
-                    break;
-
-                default:
-                    throw new IllegalStateException("Received an unknown record type " + key.version()
-                        + " in " + record);
-            }
-
-            lastWrittenOffset++;
-            snapshotRegistry.getOrCreateSnapshot(lastWrittenOffset);
-        }
-    }
 
     @Test
     public void testConsumerHeartbeatRequestValidation() {
@@ -2917,7 +1800,7 @@ public void testReconciliationProcess() {
                             .setPartitions(Arrays.asList(0, 1)),
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(0))
+                            .setPartitions(Collections.singletonList(0))
                     ))),
             result.response()
         );
@@ -2957,10 +1840,10 @@ public void testReconciliationProcess() {
                     .setTopicPartitions(Arrays.asList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(3)),
+                            .setPartitions(Collections.singletonList(3)),
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(2))
+                            .setPartitions(Collections.singletonList(2))
                     ))),
             result.response()
         );
@@ -3016,7 +1899,7 @@ public void testReconciliationProcess() {
                     .setPartitions(Arrays.asList(0, 1)),
                 new ConsumerGroupHeartbeatRequestData.TopicPartitions()
                     .setTopicId(barTopicId)
-                    .setPartitions(Arrays.asList(0))
+                    .setPartitions(Collections.singletonList(0))
             )));
 
         assertResponseEquals(
@@ -3031,7 +1914,7 @@ public void testReconciliationProcess() {
                             .setPartitions(Arrays.asList(0, 1)),
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(0))
+                            .setPartitions(Collections.singletonList(0))
                     ))),
             result.response()
         );
@@ -3082,10 +1965,10 @@ public void testReconciliationProcess() {
                 .setMemberEpoch(11)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(1))))),
+                            .setPartitions(Collections.singletonList(1))))),
             result.response()
         );
 
@@ -3134,10 +2017,10 @@ public void testReconciliationProcess() {
             .setTopicPartitions(Arrays.asList(
                 new ConsumerGroupHeartbeatRequestData.TopicPartitions()
                     .setTopicId(fooTopicId)
-                    .setPartitions(Arrays.asList(3)),
+                    .setPartitions(Collections.singletonList(3)),
                 new ConsumerGroupHeartbeatRequestData.TopicPartitions()
                     .setTopicId(barTopicId)
-                    .setPartitions(Arrays.asList(2))
+                    .setPartitions(Collections.singletonList(2))
             )));
 
         assertResponseEquals(
@@ -3152,7 +2035,7 @@ public void testReconciliationProcess() {
                             .setPartitions(Arrays.asList(2, 3)),
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(2))
+                            .setPartitions(Collections.singletonList(2))
                     ))),
             result.response()
         );
@@ -3190,7 +2073,7 @@ public void testReconciliationProcess() {
                             .setPartitions(Arrays.asList(4, 5)),
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(barTopicId)
-                            .setPartitions(Arrays.asList(1))))),
+                            .setPartitions(Collections.singletonList(1))))),
             result.response()
         );
 
@@ -3308,7 +2191,7 @@ public void testReconciliationRestartsWhenNewTargetAssignmentIsInstalled() {
                 .setMemberEpoch(10)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1))))),
@@ -3392,10 +2275,10 @@ public void testReconciliationRestartsWhenNewTargetAssignmentIsInstalled() {
                 .setMemberEpoch(10)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0))))),
+                            .setPartitions(Collections.singletonList(0))))),
             result.response()
         );
 
@@ -3668,7 +2551,7 @@ public void testSubscriptionMetadataRefreshedAfterGroupIsLoaded() {
                 .setMemberEpoch(11)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1, 2, 3, 4, 5))
@@ -3797,7 +2680,7 @@ public void testSubscriptionMetadataRefreshedAgainAfterWriteFailure() {
                 .setMemberEpoch(11)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1, 2, 3, 4, 5))
@@ -3914,7 +2797,7 @@ public void testGroupIdsByTopics() {
         // M2 in group 2 subscribes to foo.
         context.replay(RecordHelpers.newMemberSubscriptionRecord(groupId2,
             new ConsumerGroupMember.Builder("group2-m2")
-                .setSubscribedTopicNames(Arrays.asList("foo"))
+                .setSubscribedTopicNames(Collections.singletonList("foo"))
                 .build()));
 
         assertEquals(mkSet(groupId2), context.groupMetadataManager.groupsSubscribedToTopic("foo"));
@@ -3976,19 +2859,19 @@ public void testOnNewMetadataImage() {
         // M1 in group 3 subscribes to d.
         context.replay(RecordHelpers.newMemberSubscriptionRecord("group3",
             new ConsumerGroupMember.Builder("group3-m1")
-                .setSubscribedTopicNames(Arrays.asList("d"))
+                .setSubscribedTopicNames(Collections.singletonList("d"))
                 .build()));
 
         // M1 in group 4 subscribes to e.
         context.replay(RecordHelpers.newMemberSubscriptionRecord("group4",
             new ConsumerGroupMember.Builder("group4-m1")
-                .setSubscribedTopicNames(Arrays.asList("e"))
+                .setSubscribedTopicNames(Collections.singletonList("e"))
                 .build()));
 
         // M1 in group 5 subscribes to f.
         context.replay(RecordHelpers.newMemberSubscriptionRecord("group5",
             new ConsumerGroupMember.Builder("group5-m1")
-                .setSubscribedTopicNames(Arrays.asList("f"))
+                .setSubscribedTopicNames(Collections.singletonList("f"))
                 .build()));
 
         // Ensures that all refresh flags are set to the future.
@@ -4034,7 +2917,7 @@ public void testOnNewMetadataImage() {
             assertTrue(group.hasMetadataExpired(context.time.milliseconds()));
         });
 
-        Arrays.asList("group5").forEach(groupId -> {
+        Collections.singletonList("group5").forEach(groupId -> {
             ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false);
             assertFalse(group.hasMetadataExpired(context.time.milliseconds()));
         });
@@ -4310,7 +3193,7 @@ public void testRevocationTimeoutLifecycle() {
                 .setMemberEpoch(1)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1, 2))))),
@@ -4376,7 +3259,7 @@ public void testRevocationTimeoutLifecycle() {
                 .setMemberEpoch(1)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1))))),
@@ -4448,10 +3331,10 @@ public void testRevocationTimeoutLifecycle() {
                 .setMemberEpoch(1)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0))))),
+                            .setPartitions(Collections.singletonList(0))))),
             result.response()
         );
 
@@ -4481,10 +3364,10 @@ public void testRevocationTimeoutLifecycle() {
                 .setMemberEpoch(3)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0))))),
+                            .setPartitions(Collections.singletonList(0))))),
             result.response()
         );
 
@@ -4542,7 +3425,7 @@ public void testRevocationTimeoutExpiration() {
                 .setMemberEpoch(1)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1, 2))))),
@@ -4606,7 +3489,7 @@ public void testRevocationTimeoutExpiration() {
                 .setMemberEpoch(1)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
+                    .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
                             .setPartitions(Arrays.asList(0, 1))))),
@@ -4657,7 +3540,7 @@ public void testOnLoaded() {
                     .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
-                    .setSubscribedTopicNames(Arrays.asList("foo"))
+                    .setSubscribedTopicNames(Collections.singletonList("foo"))
                     .setServerAssignorName("range")
                     .setAssignedPartitions(mkAssignment(
                         mkTopicAssignment(fooTopicId, 0, 1, 2)))
@@ -4670,7 +3553,7 @@ public void testOnLoaded() {
                     .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
-                    .setSubscribedTopicNames(Arrays.asList("foo"))
+                    .setSubscribedTopicNames(Collections.singletonList("foo"))
                     .setServerAssignorName("range")
                     .setPartitionsPendingAssignment(mkAssignment(
                         mkTopicAssignment(fooTopicId, 3, 4, 5)))
@@ -4699,13 +3582,13 @@ public void testGenerateRecordsOnNewClassicGroup() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.MEMBER_ID_REQUIRED.code(), joinResult.joinFuture.get().errorCode());
 
@@ -4722,14 +3605,14 @@ public void testGenerateRecordsOnNewClassicGroupFailureTransformsError() throws
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withGroupInstanceId("group-instance-id")
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
         assertFalse(joinResult.joinFuture.isDone());
 
         // Simulate a failed write to the log.
@@ -4778,7 +3661,7 @@ public void testReplayGroupMetadataRecords(boolean useDefaultRebalanceTimeout) {
             ));
         });
 
-        Record groupMetadataRecord = newGroupMetadataRecord("group-id",
+        Record groupMetadataRecord = GroupMetadataManagerTestContext.newGroupMetadataRecord("group-id",
             new GroupMetadataValue()
                 .setMembers(members)
                 .setGeneration(1)
@@ -4829,8 +3712,8 @@ public void testOnLoadedExceedGroupMaxSizeTriggersRebalance() {
             Collections.singletonList("foo"))).array();
         List<GroupMetadataValue.MemberMetadata> members = new ArrayList<>();
 
-        IntStream.range(0, 2).forEach(i -> {
-            members.add(new GroupMetadataValue.MemberMetadata()
+        IntStream.range(0, 2).forEach(i -> members.add(
+            new GroupMetadataValue.MemberMetadata()
                 .setMemberId("member-" + i)
                 .setGroupInstanceId("group-instance-id-" + i)
                 .setSubscription(subscription)
@@ -4839,10 +3722,9 @@ public void testOnLoadedExceedGroupMaxSizeTriggersRebalance() {
                 .setClientHost("host-" + i)
                 .setSessionTimeout(4000)
                 .setRebalanceTimeout(9000)
-            );
-        });
+        ));
 
-        Record groupMetadataRecord = newGroupMetadataRecord("group-id",
+        Record groupMetadataRecord = GroupMetadataManagerTestContext.newGroupMetadataRecord("group-id",
             new GroupMetadataValue()
                 .setMembers(members)
                 .setGeneration(1)
@@ -4869,8 +3751,8 @@ public void testOnLoadedSchedulesClassicGroupMemberHeartbeats() {
             Collections.singletonList("foo"))).array();
         List<GroupMetadataValue.MemberMetadata> members = new ArrayList<>();
 
-        IntStream.range(0, 2).forEach(i -> {
-            members.add(new GroupMetadataValue.MemberMetadata()
+        IntStream.range(0, 2).forEach(i -> members.add(
+            new GroupMetadataValue.MemberMetadata()
                 .setMemberId("member-" + i)
                 .setGroupInstanceId("group-instance-id-" + i)
                 .setSubscription(subscription)
@@ -4879,10 +3761,10 @@ public void testOnLoadedSchedulesClassicGroupMemberHeartbeats() {
                 .setClientHost("host-" + i)
                 .setSessionTimeout(4000)
                 .setRebalanceTimeout(9000)
-            );
-        });
+            )
+        );
 
-        Record groupMetadataRecord = newGroupMetadataRecord("group-id",
+        Record groupMetadataRecord = GroupMetadataManagerTestContext.newGroupMetadataRecord("group-id",
             new GroupMetadataValue()
                 .setMembers(members)
                 .setGeneration(1)
@@ -4911,7 +3793,7 @@ public void testJoinGroupShouldReceiveErrorIfGroupOverMaxSize() throws Exception
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -4919,12 +3801,12 @@ public void testJoinGroupShouldReceiveErrorIfGroupOverMaxSize() throws Exception
             .build();
 
         IntStream.range(0, 10).forEach(i -> {
-            JoinResult joinResult = context.sendClassicGroupJoin(request);
+            GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
             assertFalse(joinResult.joinFuture.isDone());
             assertTrue(joinResult.records.isEmpty());
         });
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.GROUP_MAX_SIZE_REACHED.code(), joinResult.joinFuture.get().errorCode());
@@ -4941,7 +3823,7 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndRequiredKnownMember() {
 
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -4949,7 +3831,7 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndRequiredKnownMember() {
 
         // First round of join requests. Generate member ids. All requests will be accepted
         // as the group is still Empty.
-        List<JoinResult> firstRoundJoinResults = IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> firstRoundJoinResults = IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
             request,
             requiredKnownMemberId
         )).collect(Collectors.toList());
@@ -4962,18 +3844,18 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndRequiredKnownMember() {
 
         // Second round of join requests with the generated member ids.
         // One of them will fail, reaching group max size.
-        List<JoinResult> secondRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> secondRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
             request.setMemberId(memberId),
             requiredKnownMemberId
         )).collect(Collectors.toList());
 
         // Advance clock by group initial rebalance delay to complete first inital delayed join.
         // This will extend the initial rebalance as new members have joined.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
         // Advance clock by group initial rebalance delay to complete second inital delayed join.
         // Since there are no new members that joined since the previous delayed join,
         // the join group phase will complete.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
 
         verifyClassicGroupJoinResponses(secondRoundJoinResults, groupMaxSize, Errors.GROUP_MAX_SIZE_REACHED);
         assertEquals(groupMaxSize, group.size());
@@ -4981,7 +3863,7 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndRequiredKnownMember() {
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
         // Members that were accepted can rejoin while others are rejected in CompletingRebalance state.
-        List<JoinResult> thirdRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> thirdRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
             request.setMemberId(memberId),
             requiredKnownMemberId
         )).collect(Collectors.toList());
@@ -5000,14 +3882,14 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndNotRequiredKnownMember() {
 
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
         // First round of join requests. This will trigger a rebalance.
-        List<JoinResult> firstRoundJoinResults = IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> firstRoundJoinResults = IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
             request,
             requiredKnownMemberId
         )).collect(Collectors.toList());
@@ -5018,16 +3900,16 @@ public void testDynamicMembersJoinGroupWithMaxSizeAndNotRequiredKnownMember() {
 
         // Advance clock by group initial rebalance delay to complete first inital delayed join.
         // This will extend the initial rebalance as new members have joined.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
         // Advance clock by group initial rebalance delay to complete second inital delayed join.
         // Since there are no new members that joined since the previous delayed join,
         // we will complete the rebalance.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
 
         List<String> memberIds = verifyClassicGroupJoinResponses(firstRoundJoinResults, groupMaxSize, Errors.GROUP_MAX_SIZE_REACHED);
 
         // Members that were accepted can rejoin while others are rejected in CompletingRebalance state.
-        List<JoinResult> secondRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> secondRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
             request.setMemberId(memberId),
             requiredKnownMemberId
         )).collect(Collectors.toList());
@@ -5053,16 +3935,16 @@ public void testStaticMembersJoinGroupWithMaxSize() {
 
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
         // First round of join requests. This will trigger a rebalance.
-        List<JoinResult> firstRoundJoinResults = groupInstanceIds.stream()
-            .map(instanceId -> context.sendClassicGroupJoin(request.setGroupInstanceId(instanceId)))
-            .collect(Collectors.toList());
+        List<GroupMetadataManagerTestContext.JoinResult> firstRoundJoinResults = groupInstanceIds.stream()
+                                                                                                 .map(instanceId -> context.sendClassicGroupJoin(request.setGroupInstanceId(instanceId)))
+                                                                                                 .collect(Collectors.toList());
 
         assertEquals(groupMaxSize, group.size());
         assertEquals(groupMaxSize, group.numAwaitingJoinResponse());
@@ -5070,17 +3952,17 @@ public void testStaticMembersJoinGroupWithMaxSize() {
 
         // Advance clock by group initial rebalance delay to complete first inital delayed join.
         // This will extend the initial rebalance as new members have joined.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
         // Advance clock by group initial rebalance delay to complete second inital delayed join.
         // Since there are no new members that joined since the previous delayed join,
         // we will complete the rebalance.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
 
         List<String> memberIds = verifyClassicGroupJoinResponses(firstRoundJoinResults, groupMaxSize, Errors.GROUP_MAX_SIZE_REACHED);
 
         // Members which were accepted can rejoin, others are rejected, while
         // completing rebalance
-        List<JoinResult> secondRoundJoinResults =  IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> secondRoundJoinResults =  IntStream.range(0, groupMaxSize + 1).mapToObj(i -> context.sendClassicGroupJoin(
             request
                 .setMemberId(memberIds.get(i))
                 .setGroupInstanceId(groupInstanceIds.get(i))
@@ -5103,16 +3985,16 @@ public void testDynamicMembersCanRejoinGroupWithMaxSizeWhileRebalancing() {
 
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
         // First round of join requests. Generate member ids.
-        List<JoinResult> firstRoundJoinResults =  IntStream.range(0, groupMaxSize + 1)
-            .mapToObj(__ -> context.sendClassicGroupJoin(request, requiredKnownMemberId))
-            .collect(Collectors.toList());
+        List<GroupMetadataManagerTestContext.JoinResult> firstRoundJoinResults =  IntStream.range(0, groupMaxSize + 1)
+                                                                                           .mapToObj(__ -> context.sendClassicGroupJoin(request, requiredKnownMemberId))
+                                                                                           .collect(Collectors.toList());
 
         assertEquals(0, group.size());
         assertEquals(groupMaxSize + 1, group.numPendingJoinMembers());
@@ -5124,7 +4006,7 @@ public void testDynamicMembersCanRejoinGroupWithMaxSizeWhileRebalancing() {
         // Second round of join requests with the generated member ids.
         // One of them will fail, reaching group max size.
         memberIds.forEach(memberId -> {
-            JoinResult joinResult = context.sendClassicGroupJoin(
+            GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
                 request.setMemberId(memberId),
                 requiredKnownMemberId
             );
@@ -5136,18 +4018,18 @@ public void testDynamicMembersCanRejoinGroupWithMaxSizeWhileRebalancing() {
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Members can rejoin while rebalancing
-        List<JoinResult> thirdRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
+        List<GroupMetadataManagerTestContext.JoinResult> thirdRoundJoinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
             request.setMemberId(memberId),
             requiredKnownMemberId
         )).collect(Collectors.toList());
 
         // Advance clock by group initial rebalance delay to complete first inital delayed join.
         // This will extend the initial rebalance as new members have joined.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
         // Advance clock by group initial rebalance delay to complete second inital delayed join.
         // Since there are no new members that joined since the previous delayed join,
         // we will complete the rebalance.
-        assertNoOrEmptyResult(context.sleep(50));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(50));
 
         verifyClassicGroupJoinResponses(thirdRoundJoinResults, groupMaxSize, Errors.GROUP_MAX_SIZE_REACHED);
         assertEquals(groupMaxSize, group.size());
@@ -5171,25 +4053,23 @@ public void testLastJoiningMembersAreKickedOutWhenRejoiningGroupWithMaxSize() {
             .mapToObj(i -> group.generateMemberId("client-id", Optional.empty()))
             .collect(Collectors.toList());
 
-        memberIds.forEach(memberId -> {
-            group.add(
-                new ClassicGroupMember(
-                    memberId,
-                    Optional.empty(),
-                    "client-id",
-                    "client-host",
-                    10000,
-                    5000,
-                    "consumer",
-                    toProtocols("range")
-                )
-            );
-        });
+        memberIds.forEach(memberId -> group.add(
+            new ClassicGroupMember(
+                memberId,
+                Optional.empty(),
+                "client-id",
+                "client-host",
+                10000,
+                5000,
+                "consumer",
+                GroupMetadataManagerTestContext.toProtocols("range")
+            )
+        ));
 
         context.groupMetadataManager.prepareRebalance(group, "test");
 
-        List<JoinResult> joinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        List<GroupMetadataManagerTestContext.JoinResult> joinResults = memberIds.stream().map(memberId -> context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(memberId)
                 .withDefaultProtocolTypeAndProtocols()
@@ -5202,7 +4082,7 @@ public void testLastJoiningMembersAreKickedOutWhenRejoiningGroupWithMaxSize() {
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Advance clock by rebalance timeout to complete join phase.
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
 
         verifyClassicGroupJoinResponses(joinResults, groupMaxSize, Errors.GROUP_MAX_SIZE_REACHED);
 
@@ -5223,13 +4103,13 @@ public void testJoinGroupSessionTimeoutTooSmall() throws Exception {
             .withClassicGroupMinSessionTimeoutMs(minSessionTimeout)
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withSessionTimeoutMs(minSessionTimeout - 1)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
         assertTrue(joinResult.joinFuture.isDone());
         assertTrue(joinResult.records.isEmpty());
         assertEquals(Errors.INVALID_SESSION_TIMEOUT.code(), joinResult.joinFuture.get().errorCode());
@@ -5242,13 +4122,13 @@ public void testJoinGroupSessionTimeoutTooLarge() throws Exception {
             .withClassicGroupMaxSessionTimeoutMs(maxSessionTimeout)
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withSessionTimeoutMs(maxSessionTimeout + 1)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5260,19 +4140,19 @@ public void testJoinGroupUnknownMemberNewGroup() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId("member-id")
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.UNKNOWN_MEMBER_ID.code(), joinResult.joinFuture.get().errorCode());
 
         // Static member
-        request = new JoinGroupRequestBuilder()
+        request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId("member-id")
             .withGroupInstanceId("group-instance-id")
@@ -5291,7 +4171,7 @@ public void testClassicGroupJoinInconsistentProtocolType() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -5299,14 +4179,14 @@ public void testClassicGroupJoinInconsistentProtocolType() throws Exception {
 
         context.joinClassicGroupAsDynamicMemberAndCompleteJoin(request);
 
-        request = new JoinGroupRequestBuilder()
+        request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("connect")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5319,14 +4199,14 @@ public void testJoinGroupWithEmptyProtocolType() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL.code(), joinResult.joinFuture.get().errorCode());
@@ -5344,14 +4224,14 @@ public void testJoinGroupWithEmptyGroupProtocol() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
             .withProtocols(new JoinGroupRequestProtocolCollection(0))
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL.code(), joinResult.joinFuture.get().errorCode());
@@ -5369,7 +4249,7 @@ public void testNewMemberJoinExpiration() throws Exception {
 
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -5386,7 +4266,7 @@ public void testNewMemberJoinExpiration() throws Exception {
         assertEquals(0, group.allMembers().stream().filter(ClassicGroupMember::isNew).count());
 
         // Send second join group request for a new dynamic member.
-        JoinResult secondJoinResult = context.sendClassicGroupJoin(request
+        GroupMetadataManagerTestContext.JoinResult secondJoinResult = context.sendClassicGroupJoin(request
             .setSessionTimeoutMs(5000)
             .setRebalanceTimeoutMs(5000)
         );
@@ -5400,7 +4280,7 @@ public void testNewMemberJoinExpiration() throws Exception {
         assertNotEquals(firstMemberId, newMember.memberId());
 
         // Advance clock by new member join timeout to expire the second member.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupNewMemberJoinTimeoutMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupNewMemberJoinTimeoutMs));
 
         assertTrue(secondJoinResult.joinFuture.isDone());
         JoinGroupResponseData secondResponse = secondJoinResult.joinFuture.get();
@@ -5417,24 +4297,24 @@ public void testJoinGroupInconsistentGroupProtocol() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
 
-        JoinResult otherJoinResult = context.sendClassicGroupJoin(request.setProtocols(toProtocols("roundrobin")));
+        GroupMetadataManagerTestContext.JoinResult otherJoinResult = context.sendClassicGroupJoin(request.setProtocols(GroupMetadataManagerTestContext.toProtocols("roundrobin")));
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(otherJoinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
         assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL.code(), otherJoinResult.joinFuture.get().errorCode());
@@ -5446,13 +4326,13 @@ public void testJoinGroupSecondJoinInconsistentProtocol() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5471,12 +4351,12 @@ public void testJoinGroupSecondJoinInconsistentProtocol() throws Exception {
         assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL.code(), joinResult.joinFuture.get().errorCode());
 
         // Sending consistent protocol should be accepted
-        joinResult = context.sendClassicGroupJoin(request.setProtocols(toProtocols("range")), true);
+        joinResult = context.sendClassicGroupJoin(request.setProtocols(GroupMetadataManagerTestContext.toProtocols("range")), true);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
 
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
@@ -5488,7 +4368,7 @@ public void testStaticMemberJoinAsFirstMember() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withGroupInstanceId("group-instance-id")
@@ -5504,7 +4384,7 @@ public void testStaticMemberRejoinWithExplicitUnknownMemberId() throws Exception
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withGroupInstanceId("group-instance-id")
@@ -5516,7 +4396,7 @@ public void testStaticMemberRejoinWithExplicitUnknownMemberId() throws Exception
         JoinGroupResponseData response = context.joinClassicGroupAndCompleteJoin(request, false, true);
         assertEquals(Errors.NONE.code(), response.errorCode());
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request.setMemberId("unknown-member-id")
         );
 
@@ -5531,7 +4411,7 @@ public void testJoinGroupUnknownConsumerExistingGroup() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -5542,7 +4422,7 @@ public void testJoinGroupUnknownConsumerExistingGroup() throws Exception {
         JoinGroupResponseData response = context.joinClassicGroupAsDynamicMemberAndCompleteJoin(request);
         assertEquals(Errors.NONE.code(), response.errorCode());
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request.setMemberId("other-member-id")
         );
 
@@ -5558,13 +4438,13 @@ public void testJoinGroupUnknownConsumerNewDeadGroup() throws Exception {
         ClassicGroup group = context.createClassicGroup("group-id");
         group.transitionTo(DEAD);
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5576,13 +4456,13 @@ public void testJoinGroupProtocolTypeIsNotProvidedWhenAnErrorOccurs() throws Exc
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId("member-id")
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5597,25 +4477,25 @@ public void testJoinGroupReturnsTheProtocolType() throws Exception {
         context.createClassicGroup("group-id");
 
         // Leader joins
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertFalse(leaderJoinResult.joinFuture.isDone());
 
         // Member joins
-        JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(memberJoinResult.records.isEmpty());
         assertFalse(memberJoinResult.joinFuture.isDone());
 
         // Complete join group phase
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
         assertTrue(leaderJoinResult.joinFuture.isDone());
         assertTrue(memberJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
@@ -5630,21 +4510,21 @@ public void testDelayInitialRebalanceByGroupInitialRebalanceDelayOnEmptyGroup()
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2));
         assertFalse(joinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2 + 1));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2 + 1));
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
     }
@@ -5655,32 +4535,32 @@ public void testResetRebalanceDelayWhenNewMemberJoinsGroupDuringInitialRebalance
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .withRebalanceTimeoutMs(context.classicGroupInitialRebalanceDelayMs * 3)
             .build();
 
-        JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(firstMemberJoinResult.records.isEmpty());
         assertFalse(firstMemberJoinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs - 1));
-        JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs - 1));
+        GroupMetadataManagerTestContext.JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(secondMemberJoinResult.records.isEmpty());
         assertFalse(secondMemberJoinResult.joinFuture.isDone());
-        assertNoOrEmptyResult(context.sleep(2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2));
 
         // Advance clock past initial rebalance delay and verify futures are not completed.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2 + 1));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2 + 1));
         assertFalse(firstMemberJoinResult.joinFuture.isDone());
         assertFalse(secondMemberJoinResult.joinFuture.isDone());
 
         // Advance clock beyond recomputed delay and make sure the futures have completed.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs / 2));
         assertTrue(firstMemberJoinResult.joinFuture.isDone());
         assertTrue(secondMemberJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), firstMemberJoinResult.joinFuture.get().errorCode());
@@ -5693,38 +4573,38 @@ public void testDelayRebalanceUptoRebalanceTimeout() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .withRebalanceTimeoutMs(context.classicGroupInitialRebalanceDelayMs * 2)
             .build();
 
-        JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(firstMemberJoinResult.records.isEmpty());
         assertFalse(firstMemberJoinResult.joinFuture.isDone());
 
-        JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(secondMemberJoinResult.records.isEmpty());
         assertFalse(secondMemberJoinResult.joinFuture.isDone());
 
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs + 1));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs + 1));
 
-        JoinResult thirdMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult thirdMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(thirdMemberJoinResult.records.isEmpty());
         assertFalse(thirdMemberJoinResult.joinFuture.isDone());
 
         // Advance clock right before rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs - 1));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs - 1));
         assertFalse(firstMemberJoinResult.joinFuture.isDone());
         assertFalse(secondMemberJoinResult.joinFuture.isDone());
         assertFalse(thirdMemberJoinResult.joinFuture.isDone());
 
         // Advance clock beyond rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(1));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(1));
         assertTrue(firstMemberJoinResult.joinFuture.isDone());
         assertTrue(secondMemberJoinResult.joinFuture.isDone());
         assertTrue(thirdMemberJoinResult.joinFuture.isDone());
@@ -5740,7 +4620,7 @@ public void testJoinGroupReplaceStaticMember() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withGroupInstanceId("group-instance-id")
@@ -5749,7 +4629,7 @@ public void testJoinGroupReplaceStaticMember() throws Exception {
             .build();
 
         // Send join group as static member.
-        JoinResult oldMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult oldMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(oldMemberJoinResult.records.isEmpty());
         assertFalse(oldMemberJoinResult.joinFuture.isDone());
@@ -5757,7 +4637,7 @@ public void testJoinGroupReplaceStaticMember() throws Exception {
         assertEquals(1, group.size());
 
         // Replace static member with new member id. Old member id should be fenced.
-        JoinResult newMemberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult newMemberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(newMemberJoinResult.records.isEmpty());
         assertFalse(newMemberJoinResult.joinFuture.isDone());
@@ -5767,7 +4647,7 @@ public void testJoinGroupReplaceStaticMember() throws Exception {
         assertEquals(1, group.size());
 
         // Complete join for new member.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
         assertTrue(newMemberJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), newMemberJoinResult.joinFuture.get().errorCode());
         assertEquals(0, group.numAwaitingJoinResponse());
@@ -5780,14 +4660,14 @@ public void testHeartbeatExpirationShouldRemovePendingMember() throws Exception
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .withSessionTimeoutMs(1000)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5796,7 +4676,7 @@ public void testHeartbeatExpirationShouldRemovePendingMember() throws Exception
         assertEquals(1, group.numPendingJoinMembers());
 
         // Advance clock by session timeout. Pending member should be removed from group as heartbeat expires.
-        assertNoOrEmptyResult(context.sleep(1000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(1000));
         assertEquals(0, group.numPendingJoinMembers());
     }
 
@@ -5808,13 +4688,13 @@ public void testHeartbeatExpirationShouldRemoveMember() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
@@ -5828,9 +4708,10 @@ public void testHeartbeatExpirationShouldRemoveMember() throws Exception {
         assertEquals(1, timeouts.size());
         timeouts.forEach(timeout -> {
             assertEquals(classicGroupHeartbeatKey("group-id", memberId), timeout.key);
-            assertEquals(Collections.singletonList(
-                newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
-                timeout.result.records());
+            assertEquals(
+                Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
+                timeout.result.records()
+            );
         });
 
         assertTrue(joinResult.joinFuture.isDone());
@@ -5844,7 +4725,7 @@ public void testExistingMemberJoinDeadGroup() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -5858,7 +4739,7 @@ public void testExistingMemberJoinDeadGroup() throws Exception {
 
         group.transitionTo(DEAD);
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -5871,13 +4752,13 @@ public void testJoinGroupExistingPendingMemberWithGroupInstanceIdThrowsException
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertEquals(Errors.MEMBER_ID_REQUIRED.code(), joinResult.joinFuture.get().errorCode());
@@ -5899,9 +4780,9 @@ public void testJoinGroupExistingMemberUpdatedMetadataTriggersRebalance() throws
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestProtocolCollection protocols = toProtocols("range");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("range");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
@@ -5918,10 +4799,10 @@ public void testJoinGroupExistingMemberUpdatedMetadataTriggersRebalance() throws
         assertTrue(group.isInState(COMPLETING_REBALANCE));
         assertEquals(1, group.generationId());
 
-        protocols = toProtocols("range", "roundrobin");
+        protocols = GroupMetadataManagerTestContext.toProtocols("range", "roundrobin");
 
         // Send updated member metadata. This should trigger a rebalance and complete the join phase.
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request
                 .setMemberId(memberId)
                 .setProtocols(protocols)
@@ -5943,7 +4824,7 @@ public void testJoinGroupAsExistingLeaderTriggersRebalanceInStableState() throws
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -5960,7 +4841,7 @@ public void testJoinGroupAsExistingLeaderTriggersRebalanceInStableState() throws
 
         group.transitionTo(STABLE);
         // Sending join group as leader should trigger a rebalance.
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request.setMemberId(memberId)
         );
 
@@ -5978,11 +4859,11 @@ public void testJoinGroupAsExistingMemberWithUpdatedMetadataTriggersRebalanceInS
         ClassicGroup group = context.createClassicGroup("group-id");
 
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
         JoinGroupResponseData leaderResponse = context.joinClassicGroupAsDynamicMemberAndCompleteJoin(request);
@@ -5991,13 +4872,13 @@ public void testJoinGroupAsExistingMemberWithUpdatedMetadataTriggersRebalanceInS
         assertEquals(1, group.generationId());
 
         // Member joins.
-        JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(memberJoinResult.records.isEmpty());
         assertFalse(memberJoinResult.joinFuture.isDone());
 
         // Leader also rejoins. Completes join group phase.
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(request.setMemberId(leaderId));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(request.setMemberId(leaderId));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -6013,7 +4894,7 @@ public void testJoinGroupAsExistingMemberWithUpdatedMetadataTriggersRebalanceInS
         // Member rejoins with updated metadata. This should trigger a rebalance.
         String memberId = memberJoinResult.joinFuture.get().memberId();
 
-        JoinGroupRequestProtocolCollection protocols = toProtocols("range", "roundrobin");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("range", "roundrobin");
 
         memberJoinResult = context.sendClassicGroupJoin(
             request
@@ -6042,7 +4923,7 @@ public void testJoinGroupExistingMemberDoesNotTriggerRebalanceInStableState() th
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6054,13 +4935,13 @@ public void testJoinGroupExistingMemberDoesNotTriggerRebalanceInStableState() th
         assertEquals(1, group.generationId());
 
         // Member joins.
-        JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(memberJoinResult.records.isEmpty());
         assertFalse(memberJoinResult.joinFuture.isDone());
 
         // Leader also rejoins. Completes join group phase.
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(request.setMemberId(leaderId));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(request.setMemberId(leaderId));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -6090,7 +4971,7 @@ public void testJoinGroupExistingMemberInEmptyState() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6106,7 +4987,7 @@ public void testJoinGroupExistingMemberInEmptyState() throws Exception {
         group.transitionTo(PREPARING_REBALANCE);
         group.transitionTo(EMPTY);
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request.setMemberId(memberId));
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request.setMemberId(memberId));
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -6120,7 +5001,7 @@ public void testCompleteJoinRemoveNotYetRejoinedDynamicMembers() throws Exceptio
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6133,7 +5014,7 @@ public void testCompleteJoinRemoveNotYetRejoinedDynamicMembers() throws Exceptio
         assertEquals(1, group.generationId());
 
         // Add new member. This triggers a rebalance.
-        JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult memberJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(memberJoinResult.records.isEmpty());
         assertFalse(memberJoinResult.joinFuture.isDone());
@@ -6141,7 +5022,7 @@ public void testCompleteJoinRemoveNotYetRejoinedDynamicMembers() throws Exceptio
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Advance clock by rebalance timeout. This will expire the leader as it has not rejoined.
-        assertNoOrEmptyResult(context.sleep(1000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(1000));
 
         assertTrue(memberJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), memberJoinResult.joinFuture.get().errorCode());
@@ -6156,7 +5037,7 @@ public void testCompleteJoinPhaseInEmptyStateSkipsRebalance() {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6164,7 +5045,7 @@ public void testCompleteJoinPhaseInEmptyStateSkipsRebalance() {
             .withRebalanceTimeoutMs(1000)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
@@ -6174,7 +5055,7 @@ public void testCompleteJoinPhaseInEmptyStateSkipsRebalance() {
         group.transitionTo(DEAD);
 
         // Advance clock by initial rebalance delay to complete join phase.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
         assertEquals(0, group.generationId());
     }
 
@@ -6184,7 +5065,7 @@ public void testCompleteJoinPhaseNoMembersRejoinedExtendsJoinPhase() throws Exce
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("first-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -6199,7 +5080,7 @@ public void testCompleteJoinPhaseNoMembersRejoinedExtendsJoinPhase() throws Exce
         String firstMemberId = firstMemberResponse.memberId();
 
         // Second member joins and group goes into rebalancing state.
-        JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(
             request.setGroupInstanceId("second-instance-id")
         );
 
@@ -6207,7 +5088,7 @@ public void testCompleteJoinPhaseNoMembersRejoinedExtendsJoinPhase() throws Exce
         assertFalse(secondMemberJoinResult.joinFuture.isDone());
 
         // First static member rejoins and completes join phase.
-        JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult firstMemberJoinResult = context.sendClassicGroupJoin(
             request.setMemberId(firstMemberId).setGroupInstanceId("first-instance-id"));
 
         assertTrue(firstMemberJoinResult.records.isEmpty());
@@ -6229,9 +5110,9 @@ public void testCompleteJoinPhaseNoMembersRejoinedExtendsJoinPhase() throws Exce
 
         // Advance clock by rebalance timeout to complete join phase. As long as both members have not
         // rejoined, we extend the join phase.
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
         assertEquals(10000, context.timer.timeout("join-group-id").deadlineMs - context.time.milliseconds());
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
         assertEquals(10000, context.timer.timeout("join-group-id").deadlineMs - context.time.milliseconds());
 
         assertTrue(group.isInState(PREPARING_REBALANCE));
@@ -6276,12 +5157,12 @@ public void testReplaceStaticMemberInStableStateNoError(
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("group-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
         JoinGroupResponseData response = context.joinClassicGroupAndCompleteJoin(request, true, supportSkippingAssignment);
@@ -6296,9 +5177,9 @@ public void testReplaceStaticMemberInStableStateNoError(
         group.transitionTo(STABLE);
 
         // Static member rejoins with UNKNOWN_MEMBER_ID. This should update the log with the generated member id.
-        JoinGroupRequestProtocolCollection protocols = toProtocols("range", "roundrobin");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("range", "roundrobin");
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request
                 .setProtocols(protocols)
                 .setRebalanceTimeoutMs(7000)
@@ -6308,7 +5189,7 @@ public void testReplaceStaticMemberInStableStateNoError(
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             joinResult.records
         );
         assertFalse(joinResult.joinFuture.isDone());
@@ -6360,12 +5241,12 @@ public void testReplaceStaticMemberInStableStateWithUpdatedProtocolTriggersRebal
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("group-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
-            .withProtocols(toProtocols("range", "roundrobin"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range", "roundrobin"))
             .build();
 
         JoinGroupResponseData response = context.joinClassicGroupAndCompleteJoin(request, true, true);
@@ -6378,8 +5259,8 @@ public void testReplaceStaticMemberInStableStateWithUpdatedProtocolTriggersRebal
         group.transitionTo(STABLE);
 
         // Static member rejoins with UNKNOWN_MEMBER_ID. The selected protocol changes and triggers a rebalance.
-        JoinResult joinResult = context.sendClassicGroupJoin(
-            request.setProtocols(toProtocols("roundrobin"))
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
+            request.setProtocols(GroupMetadataManagerTestContext.toProtocols("roundrobin"))
         );
 
         assertTrue(joinResult.records.isEmpty());
@@ -6397,9 +5278,9 @@ public void testReplaceStaticMemberInStableStateErrors() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestProtocolCollection protocols = toProtocols("range");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("range");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("group-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -6425,7 +5306,7 @@ public void testReplaceStaticMemberInStableStateErrors() throws Exception {
                 .setMetadata(ConsumerProtocol.serializeSubscription(new ConsumerPartitionAssignor.Subscription(
                     Collections.singletonList("bar"))).array()));
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request
                 .setProtocols(protocols)
                 .setRebalanceTimeoutMs(7000)
@@ -6435,7 +5316,7 @@ public void testReplaceStaticMemberInStableStateErrors() throws Exception {
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             joinResult.records
         );
         assertFalse(joinResult.joinFuture.isDone());
@@ -6478,12 +5359,12 @@ public void testReplaceStaticMemberInStableStateSucceeds(
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("group-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolType("consumer")
-            .withProtocols(toProtocols("range"))
+            .withProtocols(GroupMetadataManagerTestContext.toProtocols("range"))
             .build();
 
         JoinGroupResponseData response = context.joinClassicGroupAndCompleteJoin(
@@ -6502,16 +5383,16 @@ public void testReplaceStaticMemberInStableStateSucceeds(
         group.transitionTo(STABLE);
 
         // Static member rejoins with UNKNOWN_MEMBER_ID and the append succeeds.
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request
-                .setProtocols(toProtocols("range", "roundrobin"))
+                .setProtocols(GroupMetadataManagerTestContext.toProtocols("range", "roundrobin"))
                 .setRebalanceTimeoutMs(7000)
                 .setSessionTimeoutMs(6000),
             true,
             supportSkippingAssignment);
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             joinResult.records
         );
         assertFalse(joinResult.joinFuture.isDone());
@@ -6538,7 +5419,7 @@ public void testReplaceStaticMemberInStableStateSucceeds(
         assertEquals(Optional.of("group-instance-id"), newMember.groupInstanceId());
         assertEquals(7000, newMember.rebalanceTimeoutMs());
         assertEquals(6000, newMember.sessionTimeoutMs());
-        assertEquals(toProtocols("range", "roundrobin"), newMember.supportedProtocols());
+        assertEquals(GroupMetadataManagerTestContext.toProtocols("range", "roundrobin"), newMember.supportedProtocols());
         assertEquals(1, group.size());
         assertEquals(1, group.generationId());
         assertTrue(group.isInState(STABLE));
@@ -6550,7 +5431,7 @@ public void testReplaceStaticMemberInCompletingRebalanceStateTriggersRebalance()
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("group-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -6565,7 +5446,7 @@ public void testReplaceStaticMemberInCompletingRebalanceStateTriggersRebalance()
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
         // Static member rejoins with UNKNOWN_MEMBER_ID and triggers a rebalance.
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -6597,8 +5478,8 @@ public void testNewMemberTimeoutCompletion() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
@@ -6610,7 +5491,7 @@ public void testNewMemberTimeoutCompletion() throws Exception {
         assertFalse(joinResult.joinFuture.isDone());
 
         // Advance clock by initial rebalance delay to complete join phase.
-        assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupInitialRebalanceDelayMs));
 
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
@@ -6619,8 +5500,8 @@ public void testNewMemberTimeoutCompletion() throws Exception {
 
         assertEquals(0, group.allMembers().stream().filter(ClassicGroupMember::isNew).count());
 
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(joinResult.joinFuture.get().memberId())
                 .withGenerationId(joinResult.joinFuture.get().generationId())
@@ -6635,12 +5516,12 @@ public void testNewMemberTimeoutCompletion() throws Exception {
         assertEquals(1, group.size());
 
         // Make sure the NewMemberTimeout is not still in effect, and the member is not kicked
-        assertNoOrEmptyResult(context.sleep(context.classicGroupNewMemberJoinTimeoutMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupNewMemberJoinTimeoutMs));
         assertEquals(1, group.size());
 
         // Member should be removed as heartbeat expires. The group is now empty.
         List<ExpiredTimeout<Void, Record>> timeouts = context.sleep(5000);
-        List<Record> expectedRecords = Collections.singletonList(newGroupMetadataRecord(
+        List<Record> expectedRecords = Collections.singletonList(GroupMetadataManagerTestContext.newGroupMetadataRecord(
             group.groupId(),
             new GroupMetadataValue()
                 .setMembers(Collections.emptyList())
@@ -6673,7 +5554,7 @@ public void testNewMemberFailureAfterJoinGroupCompletion() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6688,8 +5569,8 @@ public void testNewMemberFailureAfterJoinGroupCompletion() throws Exception {
         assertEquals(memberId, joinResponse.leader());
         assertEquals(1, joinResponse.generationId());
 
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(memberId)
                 .withGenerationId(1)
@@ -6705,8 +5586,8 @@ public void testNewMemberFailureAfterJoinGroupCompletion() throws Exception {
         assertTrue(group.isInState(STABLE));
         assertEquals(1, group.generationId());
 
-        JoinResult otherJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(UNKNOWN_MEMBER_ID));
-        JoinResult joinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(memberId));
+        GroupMetadataManagerTestContext.JoinResult otherJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(UNKNOWN_MEMBER_ID));
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(memberId));
 
         assertTrue(otherJoinResult.records.isEmpty());
         assertTrue(joinResult.records.isEmpty());
@@ -6723,7 +5604,7 @@ public void testStaticMemberFenceDuplicateRejoinedFollower() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -6731,7 +5612,7 @@ public void testStaticMemberFenceDuplicateRejoinedFollower() throws Exception {
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // A third member joins. Trigger a rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -6742,7 +5623,7 @@ public void testStaticMemberFenceDuplicateRejoinedFollower() throws Exception {
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Old follower rejoins group will be matching current member.id.
-        JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
             request
                 .setMemberId(rebalanceResult.followerId)
                 .setGroupInstanceId("follower-instance-id")
@@ -6784,7 +5665,7 @@ public void testStaticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged()
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -6792,7 +5673,7 @@ public void testStaticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged()
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Known leader rejoins will trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
@@ -6800,14 +5681,14 @@ public void testStaticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged()
             .withRebalanceTimeoutMs(10000)
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertFalse(leaderJoinResult.joinFuture.isDone());
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Old follower rejoins group will match current member.id.
-        JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
             request
                 .setMemberId(rebalanceResult.followerId)
                 .setGroupInstanceId("follower-instance-id")
@@ -6863,8 +5744,8 @@ public void testStaticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged()
         // Duplicate follower joins group with unknown member id will trigger member.id replacement,
         // and will also trigger a rebalance under CompletingRebalance state; the old follower sync callback
         // will return fenced exception while broker replaces the member identity with the duplicate follower joins.
-        SyncResult oldFollowerSyncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult oldFollowerSyncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("follower-instance-id")
                 .withGenerationId(oldFollowerJoinResult.joinFuture.get().generationId())
@@ -6875,7 +5756,7 @@ public void testStaticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged()
         assertTrue(oldFollowerSyncResult.records.isEmpty());
         assertFalse(oldFollowerSyncResult.syncFuture.isDone());
 
-        JoinResult duplicateFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult duplicateFollowerJoinResult = context.sendClassicGroupJoin(
             request
                 .setMemberId(UNKNOWN_MEMBER_ID)
                 .setGroupInstanceId("follower-instance-id")
@@ -6906,7 +5787,7 @@ public void testStaticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -6914,7 +5795,7 @@ public void testStaticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Known leader rejoins will trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
@@ -6922,14 +5803,14 @@ public void testStaticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(
             .withRebalanceTimeoutMs(10000)
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertFalse(leaderJoinResult.joinFuture.isDone());
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Duplicate follower joins group will trigger member id replacement.
-        JoinResult duplicateFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult duplicateFollowerJoinResult = context.sendClassicGroupJoin(
             request
                 .setMemberId(UNKNOWN_MEMBER_ID)
                 .setGroupInstanceId("follower-instance-id")
@@ -6939,7 +5820,7 @@ public void testStaticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(
         assertTrue(duplicateFollowerJoinResult.joinFuture.isDone());
 
         // Old follower rejoins group will fail because member id is already updated.
-        JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
             request.setMemberId(rebalanceResult.followerId)
         );
 
@@ -7009,7 +5890,7 @@ public void testStaticMemberRejoinWithKnownMemberId() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withGroupInstanceId("group-instance-id")
@@ -7021,7 +5902,7 @@ public void testStaticMemberRejoinWithKnownMemberId() throws Exception {
 
         String memberId = joinResponse.memberId();
 
-        JoinResult rejoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult rejoinResult = context.sendClassicGroupJoin(
             request.setMemberId(memberId)
         );
 
@@ -7031,8 +5912,8 @@ public void testStaticMemberRejoinWithKnownMemberId() throws Exception {
         assertTrue(rejoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), rejoinResult.joinFuture.get().errorCode());
 
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(memberId)
                 .withGenerationId(joinResponse.generationId())
@@ -7056,7 +5937,7 @@ public void testStaticMemberRejoinWithLeaderIdAndUnknownMemberId(
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7065,21 +5946,21 @@ public void testStaticMemberRejoinWithLeaderIdAndUnknownMemberId(
 
         // A static leader rejoin with unknown id will not trigger rebalance, and no assignment will be returned.
         // As the group was in Stable state and the member id was updated, this will generate records.
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolSuperset()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             joinRequest,
             true,
             supportSkippingAssignment
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             joinResult.records
         );
         // Simulate a successful write to the log.
@@ -7110,7 +5991,7 @@ public void testStaticMemberRejoinWithLeaderIdAndUnknownMemberId(
             supportSkippingAssignment ? mkSet("leader-instance-id", "follower-instance-id") : Collections.emptySet()
         );
 
-        JoinResult oldLeaderJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult oldLeaderJoinResult = context.sendClassicGroupJoin(
             joinRequest.setMemberId(rebalanceResult.leaderId),
             true,
             supportSkippingAssignment
@@ -7121,14 +6002,14 @@ public void testStaticMemberRejoinWithLeaderIdAndUnknownMemberId(
         assertEquals(Errors.FENCED_INSTANCE_ID.code(), oldLeaderJoinResult.joinFuture.get().errorCode());
 
         // Old leader will get fenced.
-        SyncGroupRequestData oldLeaderSyncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData oldLeaderSyncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withGenerationId(rebalanceResult.generationId)
             .withMemberId(rebalanceResult.leaderId)
             .build();
 
-        SyncResult oldLeaderSyncResult = context.sendClassicGroupSync(oldLeaderSyncRequest);
+        GroupMetadataManagerTestContext.SyncResult oldLeaderSyncResult = context.sendClassicGroupSync(oldLeaderSyncRequest);
 
         assertTrue(oldLeaderSyncResult.records.isEmpty());
         assertTrue(oldLeaderSyncResult.syncFuture.isDone());
@@ -7136,7 +6017,7 @@ public void testStaticMemberRejoinWithLeaderIdAndUnknownMemberId(
 
         // Calling sync on old leader id will fail because that leader id is no longer valid and replaced.
         SyncGroupRequestData newLeaderSyncRequest = oldLeaderSyncRequest.setGroupInstanceId(null);
-        SyncResult newLeaderSyncResult = context.sendClassicGroupSync(newLeaderSyncRequest);
+        GroupMetadataManagerTestContext.SyncResult newLeaderSyncResult = context.sendClassicGroupSync(newLeaderSyncRequest);
 
         assertTrue(newLeaderSyncResult.records.isEmpty());
         assertTrue(newLeaderSyncResult.syncFuture.isDone());
@@ -7148,7 +6029,7 @@ public void testStaticMemberRejoinWithLeaderIdAndKnownMemberId() throws Exceptio
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7156,7 +6037,7 @@ public void testStaticMemberRejoinWithLeaderIdAndKnownMemberId() throws Exceptio
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Known static leader rejoin will trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
@@ -7192,7 +6073,7 @@ public void testStaticMemberRejoinWithLeaderIdAndUnexpectedDeadGroup() throws Ex
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7200,14 +6081,14 @@ public void testStaticMemberRejoinWithLeaderIdAndUnexpectedDeadGroup() throws Ex
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
         group.transitionTo(DEAD);
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -7219,7 +6100,7 @@ public void testStaticMemberRejoinWithLeaderIdAndUnexpectedEmptyGroup() throws E
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7229,14 +6110,14 @@ public void testStaticMemberRejoinWithLeaderIdAndUnexpectedEmptyGroup() throws E
         group.transitionTo(PREPARING_REBALANCE);
         group.transitionTo(EMPTY);
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request, true, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -7250,7 +6131,7 @@ public void testStaticMemberRejoinWithFollowerIdAndChangeOfProtocol() throws Exc
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id",
@@ -7260,8 +6141,8 @@ public void testStaticMemberRejoinWithFollowerIdAndChangeOfProtocol() throws Exc
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // A static follower rejoin with changed protocol will trigger rebalance.
-        JoinGroupRequestProtocolCollection protocols = toProtocols("roundrobin");
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("roundrobin");
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(rebalanceResult.followerId)
@@ -7270,13 +6151,13 @@ public void testStaticMemberRejoinWithFollowerIdAndChangeOfProtocol() throws Exc
             .withSessionTimeoutMs(sessionTimeoutMs)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
 
         // Old leader hasn't joined in the meantime, triggering a re-election.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
 
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
@@ -7309,7 +6190,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWithSele
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id",
@@ -7321,9 +6202,9 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWithSele
         assertNotEquals("roundrobin", group.selectProtocol());
 
         // A static follower rejoin with changed protocol will trigger rebalance.
-        JoinGroupRequestProtocolCollection protocols = toProtocols("roundrobin");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("roundrobin");
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -7332,13 +6213,13 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWithSele
             .withSessionTimeoutMs(sessionTimeoutMs)
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
         assertTrue(joinResult.records.isEmpty());
         assertFalse(joinResult.joinFuture.isDone());
         assertEquals("roundrobin", group.selectProtocol());
 
         // Old leader hasn't joined in the meantime, triggering a re-election.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
         assertTrue(group.hasStaticMember("leader-instance-id"));
@@ -7369,30 +6250,30 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
         );
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
-        JoinGroupRequestProtocolCollection protocols = toProtocols(group.selectProtocol());
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols(group.selectProtocol());
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocols(protocols)
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             followerJoinResult.records
         );
         // Simulate a failed write to the log.
@@ -7424,7 +6305,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         assertTrue(followerJoinResult.records.isEmpty());
 
         // Join with leader and complete join phase.
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(
             request.setGroupInstanceId("leader-instance-id")
                 .setMemberId(rebalanceResult.leaderId)
         );
@@ -7436,20 +6317,20 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
 
         // Sync with leader and receive assignment.
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
             .withGenerationId(rebalanceResult.generationId + 1)
             .build();
 
-        SyncResult leaderSyncResult = context.sendClassicGroupSync(syncRequest);
+        GroupMetadataManagerTestContext.SyncResult leaderSyncResult = context.sendClassicGroupSync(syncRequest);
 
         // Simulate a successful write to the log. This will update the group with the new (empty) assignment.
         leaderSyncResult.appendFuture.complete(null);
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             leaderSyncResult.records
         );
 
@@ -7458,7 +6339,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         assertEquals(Errors.NONE.code(), leaderSyncResult.syncFuture.get().errorCode());
 
         // Sync with old member id will also not fail as the member id is not updated due to persistence failure
-        SyncResult oldMemberSyncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult oldMemberSyncResult = context.sendClassicGroupSync(
             syncRequest
                 .setGroupInstanceId("follower-instance-id")
                 .setMemberId(rebalanceResult.followerId)
@@ -7474,7 +6355,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7483,23 +6364,23 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
 
         // A static follower rejoin with protocol changing to leader protocol subset won't trigger rebalance if updated
         // group's selectProtocol remain unchanged.
-        JoinGroupRequestProtocolCollection protocols = toProtocols(group.selectProtocol());
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols(group.selectProtocol());
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocols(protocols)
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             followerJoinResult.records
         );
 
@@ -7534,7 +6415,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
         assertEquals(Errors.FENCED_INSTANCE_ID.code(), followerJoinResult.joinFuture.get().errorCode());
 
         // Sync with old member id will fail because the member id is updated
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withGenerationId(rebalanceResult.generationId)
@@ -7542,7 +6423,7 @@ public void testStaticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSel
             .withAssignment(Collections.emptyList())
             .build();
 
-        SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
 
         assertTrue(syncResult.records.isEmpty());
         assertTrue(syncResult.syncFuture.isDone());
@@ -7564,7 +6445,7 @@ public void testStaticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollower
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7572,7 +6453,7 @@ public void testStaticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollower
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // A static leader rejoin with known member id will trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.leaderId)
@@ -7581,7 +6462,7 @@ public void testStaticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollower
             .withSessionTimeoutMs(5000)
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -7591,7 +6472,7 @@ public void testStaticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollower
         assertFalse(leaderJoinResult.joinFuture.isDone());
 
         // Rebalance completes immediately after follower rejoins.
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request.setGroupInstanceId("follower-instance-id")
                 .setMemberId(rebalanceResult.followerId),
             true,
@@ -7642,7 +6523,7 @@ public void testStaticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollower
         );
 
         // The follower protocol changed from protocolSuperset to general protocols.
-        JoinGroupRequestProtocolCollection protocols = toProtocols("range");
+        JoinGroupRequestProtocolCollection protocols = GroupMetadataManagerTestContext.toProtocols("range");
 
         followerJoinResult = context.sendClassicGroupJoin(
             request.setGroupInstanceId("follower-instance-id")
@@ -7689,7 +6570,7 @@ public void testStaticMemberRejoinAsFollowerWithUnknownMemberId() throws Excepti
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7697,21 +6578,21 @@ public void testStaticMemberRejoinAsFollowerWithUnknownMemberId() throws Excepti
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // A static follower rejoin with no protocol change will not trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolSuperset()
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
         );
 
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             followerJoinResult.records
         );
         // Simulate a successful write to log.
@@ -7741,8 +6622,8 @@ public void testStaticMemberRejoinAsFollowerWithUnknownMemberId() throws Excepti
         );
         assertNotEquals(rebalanceResult.followerId, followerJoinResult.joinFuture.get().memberId());
 
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("follower-instance-id")
                 .withGenerationId(rebalanceResult.generationId)
@@ -7761,7 +6642,7 @@ public void testStaticMemberRejoinAsFollowerWithKnownMemberIdAndNoProtocolChange
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7769,14 +6650,14 @@ public void testStaticMemberRejoinAsFollowerWithKnownMemberIdAndNoProtocolChange
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // A static follower rejoin with no protocol change will not trigger rebalance.
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(rebalanceResult.followerId)
             .withProtocolSuperset()
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -7813,20 +6694,20 @@ public void testStaticMemberRejoinAsFollowerWithMismatchedInstanceId() throws Ex
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
         );
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(rebalanceResult.followerId)
             .withProtocolSuperset()
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -7842,20 +6723,20 @@ public void testStaticMemberRejoinAsLeaderWithMismatchedInstanceId() throws Exce
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
         );
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(rebalanceResult.leaderId)
             .withProtocolSuperset()
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -7877,13 +6758,13 @@ public void testStaticMemberSyncAsLeaderWithInvalidMemberId() throws Exception {
             "follower-instance-id"
         );
 
-        SyncGroupRequestData request = new SyncGroupRequestBuilder()
+        SyncGroupRequestData request = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId("invalid-member-id")
             .build();
 
-        SyncResult leaderSyncResult = context.sendClassicGroupSync(request);
+        GroupMetadataManagerTestContext.SyncResult leaderSyncResult = context.sendClassicGroupSync(request);
 
         assertTrue(leaderSyncResult.records.isEmpty());
         assertTrue(leaderSyncResult.syncFuture.isDone());
@@ -7895,7 +6776,7 @@ public void testGetDifferentStaticMemberIdAfterEachRejoin() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
@@ -7904,21 +6785,21 @@ public void testGetDifferentStaticMemberIdAfterEachRejoin() throws Exception {
 
         String lastMemberId = rebalanceResult.leaderId;
         for (int i = 0; i < 5; i++) {
-            JoinGroupRequestData request = new JoinGroupRequestBuilder()
+            JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("leader-instance-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withProtocolSuperset()
                 .build();
 
-            JoinResult leaderJoinResult = context.sendClassicGroupJoin(
+            GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(
                 request,
                 true,
                 true
             );
 
             assertEquals(
-                Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+                Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
                 leaderJoinResult.records
             );
             // Simulate a successful write to log.
@@ -7937,20 +6818,20 @@ public void testStaticMemberJoinWithUnknownInstanceIdAndKnownMemberId() throws E
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
         );
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("unknown-instance-id")
             .withMemberId(rebalanceResult.leaderId)
             .withProtocolSuperset()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -7976,7 +6857,7 @@ public void testStaticMemberReJoinWithIllegalStateAsUnknownMember() throws Excep
         group.transitionTo(PREPARING_REBALANCE);
         group.transitionTo(EMPTY);
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("follower-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -8001,7 +6882,7 @@ public void testStaticMemberFollowerFailToRejoinBeforeRebalanceTimeout() throws
             .build();
 
         // Increase session timeout so that the follower won't be evicted when rebalance timeout is reached.
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id",
@@ -8013,14 +6894,14 @@ public void testStaticMemberFollowerFailToRejoinBeforeRebalanceTimeout() throws
         String newMemberInstanceId = "new-member-instance-id";
         String leaderId = rebalanceResult.leaderId;
 
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId(newMemberInstanceId)
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolSuperset()
             .build();
 
-        JoinResult newMemberJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult newMemberJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -8030,7 +6911,7 @@ public void testStaticMemberFollowerFailToRejoinBeforeRebalanceTimeout() throws
         assertFalse(newMemberJoinResult.joinFuture.isDone());
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(
             request
                 .setGroupInstanceId("leader-instance-id")
                 .setMemberId(leaderId),
@@ -8042,7 +6923,7 @@ public void testStaticMemberFollowerFailToRejoinBeforeRebalanceTimeout() throws
         assertFalse(leaderJoinResult.joinFuture.isDone());
 
         // Advance clock by rebalance timeout to complete join phase.
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
 
         assertTrue(leaderJoinResult.joinFuture.isDone());
         assertTrue(newMemberJoinResult.joinFuture.isDone());
@@ -8090,7 +6971,7 @@ public void testStaticMemberLeaderFailToRejoinBeforeRebalanceTimeout() throws Ex
             .build();
 
         // Increase session timeout so that the leader won't be evicted when rebalance timeout is reached.
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id",
@@ -8100,14 +6981,14 @@ public void testStaticMemberLeaderFailToRejoinBeforeRebalanceTimeout() throws Ex
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         String newMemberInstanceId = "new-member-instance-id";
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId(newMemberInstanceId)
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolSuperset()
             .build();
 
-        JoinResult newMemberJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult newMemberJoinResult = context.sendClassicGroupJoin(
             request,
             true,
             true
@@ -8117,7 +6998,7 @@ public void testStaticMemberLeaderFailToRejoinBeforeRebalanceTimeout() throws Ex
         assertFalse(newMemberJoinResult.joinFuture.isDone());
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
-        JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult oldFollowerJoinResult = context.sendClassicGroupJoin(
             request
                 .setGroupInstanceId("follower-instance-id")
                 .setMemberId(rebalanceResult.followerId),
@@ -8129,7 +7010,7 @@ public void testStaticMemberLeaderFailToRejoinBeforeRebalanceTimeout() throws Ex
         assertFalse(oldFollowerJoinResult.joinFuture.isDone());
 
         // Advance clock by rebalance timeout to complete join phase.
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
 
         assertTrue(oldFollowerJoinResult.joinFuture.isDone());
         assertTrue(newMemberJoinResult.joinFuture.isDone());
@@ -8236,26 +7117,26 @@ private void testSyncGroupProtocolTypeAndNameWith(
         context.createClassicGroup("group-id");
 
         // JoinGroup(leader) with the Protocol Type of the group
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withProtocolSuperset()
             .build();
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest);
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest);
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertFalse(leaderJoinResult.joinFuture.isDone());
 
         // JoinGroup(follower) with the Protocol Type of the group
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest.setGroupInstanceId("follower-instance-id"));
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest.setGroupInstanceId("follower-instance-id"));
 
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
 
         // Advance clock by rebalance timeout to complete join phase.
-        assertNoOrEmptyResult(context.sleep(10000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(10000));
         assertTrue(leaderJoinResult.joinFuture.isDone());
         assertTrue(followerJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
@@ -8268,7 +7149,7 @@ private void testSyncGroupProtocolTypeAndNameWith(
         // SyncGroup with the provided Protocol Type and Name
         List<SyncGroupRequestAssignment> assignment = new ArrayList<>();
         assignment.add(new SyncGroupRequestAssignment().setMemberId(leaderId));
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(leaderId)
             .withProtocolType(protocolType.orElse(null))
@@ -8277,7 +7158,7 @@ private void testSyncGroupProtocolTypeAndNameWith(
             .withAssignment(assignment)
             .build();
 
-        SyncResult leaderSyncResult = context.sendClassicGroupSync(syncRequest);
+        GroupMetadataManagerTestContext.SyncResult leaderSyncResult = context.sendClassicGroupSync(syncRequest);
         // Simulate a successful write to the log.
         leaderSyncResult.appendFuture.complete(null);
 
@@ -8286,7 +7167,7 @@ private void testSyncGroupProtocolTypeAndNameWith(
         assertEquals(expectedProtocolType.orElse(null), leaderSyncResult.syncFuture.get().protocolType());
         assertEquals(expectedProtocolName.orElse(null), leaderSyncResult.syncFuture.get().protocolName());
 
-        SyncResult followerSyncResult = context.sendClassicGroupSync(syncRequest.setMemberId(followerId));
+        GroupMetadataManagerTestContext.SyncResult followerSyncResult = context.sendClassicGroupSync(syncRequest.setMemberId(followerId));
 
         assertTrue(followerSyncResult.records.isEmpty());
         assertTrue(followerSyncResult.syncFuture.isDone());
@@ -8301,8 +7182,8 @@ public void testSyncGroupFromUnknownGroup() throws Exception {
             .build();
 
         // SyncGroup with the provided Protocol Type and Name
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId("member-id")
                 .withGenerationId(1)
@@ -8321,7 +7202,7 @@ public void testSyncGroupFromUnknownMember() throws Exception {
         context.createClassicGroup("group-id");
 
         JoinGroupResponseData joinResponse = context.joinClassicGroupAndCompleteJoin(
-            new JoinGroupRequestBuilder()
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("leader-instance-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
@@ -8336,14 +7217,14 @@ public void testSyncGroupFromUnknownMember() throws Exception {
 
         List<SyncGroupRequestAssignment> assignment = new ArrayList<>();
         assignment.add(new SyncGroupRequestAssignment().setMemberId(memberId));
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(memberId)
             .withGenerationId(generationId)
             .withAssignment(assignment)
             .build();
 
-        SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
 
         // Simulate a successful write to log.
         syncResult.appendFuture.complete(null);
@@ -8367,7 +7248,7 @@ public void testSyncGroupFromIllegalGeneration() throws Exception {
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -8380,8 +7261,8 @@ public void testSyncGroupFromIllegalGeneration() throws Exception {
         int generationId = joinResponse.generationId();
 
         // Send the sync group with an invalid generation
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(memberId)
                 .withGenerationId(generationId + 1)
@@ -8399,7 +7280,7 @@ public void testSyncGroupAsLeaderAppendFailureTransformsError() throws Exception
             .build();
         context.createClassicGroup("group-id");
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -8409,8 +7290,8 @@ public void testSyncGroupAsLeaderAppendFailureTransformsError() throws Exception
         JoinGroupResponseData joinResponse = context.joinClassicGroupAndCompleteJoin(joinRequest, true, true);
 
         // Send the sync group with an invalid generation
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(joinResponse.memberId())
                 .withGenerationId(1)
@@ -8434,18 +7315,18 @@ public void testJoinGroupFromUnchangedFollowerDoesNotRebalance() throws Exceptio
             .build();
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
 
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -8479,7 +7360,7 @@ public void testLeaderFailureInSyncGroup() throws Exception {
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -8487,12 +7368,12 @@ public void testLeaderFailureInSyncGroup() throws Exception {
             .withSessionTimeoutMs(5000)
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
 
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -8509,7 +7390,7 @@ public void testLeaderFailureInSyncGroup() throws Exception {
 
         // With no leader SyncGroup, the follower's sync request should fail with an error indicating
         // that it should rejoin
-        SyncResult followerSyncResult = context.sendClassicGroupSync(new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult followerSyncResult = context.sendClassicGroupSync(new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(followerId)
             .withGenerationId(nextGenerationId)
@@ -8542,7 +7423,7 @@ public void testSyncGroupFollowerAfterLeader() throws Exception {
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -8550,12 +7431,12 @@ public void testSyncGroupFollowerAfterLeader() throws Exception {
             .withSessionTimeoutMs(5000)
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(UNKNOWN_MEMBER_ID));
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(UNKNOWN_MEMBER_ID));
 
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -8583,14 +7464,14 @@ public void testSyncGroupFollowerAfterLeader() throws Exception {
             .setAssignment(followerAssignment)
         );
 
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(leaderJoinResponse.memberId())
             .withGenerationId(leaderJoinResponse.generationId())
             .withAssignment(assignment)
             .build();
 
-        SyncResult syncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
             syncRequest.setGenerationId(nextGenerationId)
         );
 
@@ -8602,7 +7483,7 @@ public void testSyncGroupFollowerAfterLeader() throws Exception {
         assertEquals(leaderAssignment, syncResult.syncFuture.get().assignment());
 
         // Sync group with follower to get new assignment.
-        SyncResult followerSyncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult followerSyncResult = context.sendClassicGroupSync(
             syncRequest
                 .setMemberId(followerId)
                 .setGenerationId(nextGenerationId)
@@ -8626,7 +7507,7 @@ public void testSyncGroupLeaderAfterFollower() throws Exception {
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -8634,12 +7515,12 @@ public void testSyncGroupLeaderAfterFollower() throws Exception {
             .withSessionTimeoutMs(5000)
             .build();
 
-        JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
+        GroupMetadataManagerTestContext.JoinResult followerJoinResult = context.sendClassicGroupJoin(joinRequest);
 
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
 
-        JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
+        GroupMetadataManagerTestContext.JoinResult leaderJoinResult = context.sendClassicGroupJoin(joinRequest.setMemberId(leaderJoinResponse.memberId()));
 
         assertTrue(leaderJoinResult.records.isEmpty());
         assertTrue(leaderJoinResult.joinFuture.isDone());
@@ -8657,13 +7538,13 @@ public void testSyncGroupLeaderAfterFollower() throws Exception {
         byte[] followerAssignment = new byte[]{1};
 
         // Sync group with follower to get new assignment.
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(leaderJoinResponse.memberId())
             .withGenerationId(leaderJoinResponse.generationId())
             .build();
 
-        SyncResult followerSyncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult followerSyncResult = context.sendClassicGroupSync(
             syncRequest
                 .setMemberId(followerId)
                 .setGenerationId(nextGenerationId)
@@ -8683,7 +7564,7 @@ public void testSyncGroupLeaderAfterFollower() throws Exception {
             .setAssignment(followerAssignment)
         );
 
-        SyncResult syncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
             syncRequest
                 .setMemberId(leaderJoinResponse.memberId())
                 .setGenerationId(nextGenerationId)
@@ -8721,8 +7602,8 @@ public void testJoinGroupFromUnchangedLeaderShouldRebalance() throws Exception {
 
         // Join group from the leader should force the group to rebalance, which allows the
         // leader to push new assignment when local metadata changes
-        JoinResult leaderRejoinResult = context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        GroupMetadataManagerTestContext.JoinResult leaderRejoinResult = context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(leaderJoinResponse.memberId())
                 .withDefaultProtocolTypeAndProtocols()
@@ -8746,13 +7627,13 @@ public void testJoinGroupCompletionWhenPendingMemberJoins() throws Exception {
         JoinGroupResponseData pendingMemberResponse = context.setupGroupWithPendingMember(group).pendingMemberResponse;
 
         // Compete join group for the pending member
-        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(pendingMemberResponse.memberId())
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -8775,7 +7656,7 @@ public void testJoinGroupCompletionWhenPendingMemberTimesOut() throws Exception
         // Advancing clock by > 2500 (session timeout for the third member)
         // and < 5000 (for first and second members). This will force the coordinator to attempt join
         // completion on heartbeat expiration (since we are in PendingRebalance stage).
-        assertNoOrEmptyResult(context.sleep(3000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(3000));
         assertTrue(group.isInState(COMPLETING_REBALANCE));
         assertEquals(2, group.allMembers().size());
         assertEquals(0, group.numPendingJoinMembers());
@@ -8787,8 +7668,8 @@ public void testGenerationIdIncrementsOnRebalance() throws Exception {
             .build();
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
 
-        JoinResult joinResult = context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(leaderJoinResponse.memberId())
                 .withDefaultProtocolTypeAndProtocols()
@@ -8806,14 +7687,14 @@ public void testStaticMemberHeartbeatLeaderWithInvalidMemberId() throws Exceptio
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
 
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id"
         );
 
-        SyncResult syncResult = context.sendClassicGroupSync(
-            new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("leader-instance-id")
                 .withMemberId(rebalanceResult.leaderId)
@@ -8883,7 +7764,7 @@ public void testHeartbeatEmptyGroup() {
             10000,
             5000,
             "consumer",
-            toProtocols("range")
+            GroupMetadataManagerTestContext.toProtocols("range")
         ));
 
         HeartbeatRequestData heartbeatRequest = new HeartbeatRequestData()
@@ -8915,13 +7796,13 @@ public void testHeartbeatDuringPreparingRebalance() throws Exception {
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
 
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .build();
 
-        JoinResult joinResult = context.sendClassicGroupJoin(joinRequest, true);
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(joinRequest, true);
 
         assertTrue(joinResult.records.isEmpty());
         assertTrue(joinResult.joinFuture.isDone());
@@ -8950,7 +7831,7 @@ public void testHeartbeatDuringCompletingRebalance() throws Exception {
         ClassicGroup group = context.createClassicGroup("group-id");
 
         JoinGroupResponseData leaderJoinResponse =
-            context.joinClassicGroupAsDynamicMemberAndCompleteJoin(new JoinGroupRequestBuilder()
+            context.joinClassicGroupAsDynamicMemberAndCompleteJoin(new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
@@ -8975,14 +7856,12 @@ public void testHeartbeatIllegalGeneration() throws Exception {
             .build();
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
 
-        assertThrows(IllegalGenerationException.class, () -> {
-            context.sendClassicGroupHeartbeat(
-                new HeartbeatRequestData()
-                    .setGroupId("group-id")
-                    .setMemberId(leaderJoinResponse.memberId())
-                    .setGenerationId(leaderJoinResponse.generationId() + 1)
-            );
-        });
+        assertThrows(IllegalGenerationException.class, () -> context.sendClassicGroupHeartbeat(
+            new HeartbeatRequestData()
+                .setGroupId("group-id")
+                .setMemberId(leaderJoinResponse.memberId())
+                .setGenerationId(leaderJoinResponse.generationId() + 1)
+        ));
     }
 
     @Test
@@ -9026,7 +7905,7 @@ public void testClassicGroupMemberHeartbeatMaintainsSession() throws Exception {
         JoinGroupResponseData leaderJoinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteRebalance("group-id");
 
         // Advance clock by 1/2 of session timeout.
-        assertNoOrEmptyResult(context.sleep(2500));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
 
         HeartbeatRequestData heartbeatRequest = new HeartbeatRequestData()
             .setGroupId("group-id")
@@ -9036,7 +7915,7 @@ public void testClassicGroupMemberHeartbeatMaintainsSession() throws Exception {
         HeartbeatResponseData heartbeatResponse = context.sendClassicGroupHeartbeat(heartbeatRequest);
         assertEquals(Errors.NONE.code(), heartbeatResponse.errorCode());
 
-        assertNoOrEmptyResult(context.sleep(2500));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
 
         heartbeatResponse = context.sendClassicGroupHeartbeat(heartbeatRequest);
         assertEquals(Errors.NONE.code(), heartbeatResponse.errorCode());
@@ -9051,8 +7930,8 @@ public void testClassicGroupMemberSessionTimeoutDuringRebalance() throws Excepti
 
         // Add a new member. This should trigger a rebalance. The new member has the
         // 'classicGroupNewMemberJoinTimeoutMs` session timeout, so it has a longer expiration than the existing member.
-        JoinResult otherJoinResult = context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        GroupMetadataManagerTestContext.JoinResult otherJoinResult = context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
@@ -9066,7 +7945,7 @@ public void testClassicGroupMemberSessionTimeoutDuringRebalance() throws Excepti
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
         // Advance clock by 1/2 of session timeout.
-        assertNoOrEmptyResult(context.sleep(2500));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
 
         HeartbeatRequestData heartbeatRequest = new HeartbeatRequestData()
             .setGroupId("group-id")
@@ -9077,12 +7956,12 @@ public void testClassicGroupMemberSessionTimeoutDuringRebalance() throws Excepti
         assertEquals(Errors.REBALANCE_IN_PROGRESS.code(), heartbeatResponse.errorCode());
 
         // Advance clock by first member's session timeout.
-        assertNoOrEmptyResult(context.sleep(5000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(5000));
 
         assertThrows(UnknownMemberIdException.class, () -> context.sendClassicGroupHeartbeat(heartbeatRequest));
 
         // Advance clock by remaining rebalance timeout to complete join phase.
-        assertNoOrEmptyResult(context.sleep(2500));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
 
         assertTrue(otherJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), otherJoinResult.joinFuture.get().errorCode());
@@ -9098,7 +7977,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         ClassicGroup group = context.createClassicGroup("group-id");
 
         // Create a group with a single member
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -9115,13 +7994,13 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         assertEquals(1, firstGenerationId);
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
-        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+        SyncGroupRequestData syncRequest = new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(firstMemberId)
             .withGenerationId(firstGenerationId)
             .build();
 
-        SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(syncRequest);
 
         // Simulate a successful write to the log.
         syncResult.appendFuture.complete(null);
@@ -9132,7 +8011,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
 
         // Add a new dynamic member. This should trigger a rebalance. The new member has the
         // 'classicGroupNewMemberJoinTimeoutMs` session timeout, so it has a longer expiration than the existing member.
-        JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult secondMemberJoinResult = context.sendClassicGroupJoin(
             joinRequest
                 .setMemberId(UNKNOWN_MEMBER_ID)
                 .setGroupInstanceId(null)
@@ -9150,7 +8029,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
             .setGenerationId(firstGenerationId);
 
         for (int i = 0; i < 2; i++) {
-            assertNoOrEmptyResult(context.sleep(2500));
+            GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
             HeartbeatResponseData heartbeatResponse = context.sendClassicGroupHeartbeat(firstMemberHeartbeatRequest);
             assertEquals(Errors.REBALANCE_IN_PROGRESS.code(), heartbeatResponse.errorCode());
         }
@@ -9158,7 +8037,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         // Advance clock by remaining rebalance timeout to complete join phase.
         // The second member will become the leader. However, as the first member is a static member
         // it will not be kicked out.
-        assertNoOrEmptyResult(context.sleep(8000));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(8000));
 
         assertTrue(secondMemberJoinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), secondMemberJoinResult.joinFuture.get().errorCode());
@@ -9188,7 +8067,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         // Now session timeout the unjoined (first) member. Still keeping the new member.
         List<Errors> expectedErrors = Arrays.asList(Errors.NONE, Errors.NONE, Errors.REBALANCE_IN_PROGRESS);
         for (Errors expectedError : expectedErrors) {
-            assertNoOrEmptyResult(context.sleep(2000));
+            GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2000));
             HeartbeatResponseData heartbeatResponse = context.sendClassicGroupHeartbeat(
                 firstMemberHeartbeatRequest
                     .setMemberId(otherMemberId)
@@ -9200,7 +8079,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         assertEquals(1, group.size());
         assertTrue(group.isInState(PREPARING_REBALANCE));
 
-        JoinResult otherMemberRejoinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult otherMemberRejoinResult = context.sendClassicGroupJoin(
             joinRequest
                 .setMemberId(otherMemberId)
                 .setGroupInstanceId(null)
@@ -9213,7 +8092,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         assertEquals(3, otherMemberRejoinResult.joinFuture.get().generationId());
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
-        SyncResult otherMemberResyncResult = context.sendClassicGroupSync(
+        GroupMetadataManagerTestContext.SyncResult otherMemberResyncResult = context.sendClassicGroupSync(
             syncRequest
                 .setGroupInstanceId(null)
                 .setMemberId(otherMemberId)
@@ -9230,7 +8109,7 @@ public void testRebalanceCompletesBeforeMemberJoins() throws Exception {
         // The joined member should get heart beat response with no error. Let the new member keep
         // heartbeating for a while to verify that no new rebalance is triggered unexpectedly.
         for (int i = 0; i < 20; i++) {
-            assertNoOrEmptyResult(context.sleep(2000));
+            GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2000));
             HeartbeatResponseData heartbeatResponse = context.sendClassicGroupHeartbeat(
                 firstMemberHeartbeatRequest
                     .setMemberId(otherMemberId)
@@ -9268,7 +8147,7 @@ public void testSecondMemberPartiallyJoinAndTimeout() throws Exception {
         ClassicGroup group = context.createClassicGroup("group-id");
 
         // Create a group with a single member
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withGroupInstanceId("leader-instance-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
@@ -9285,7 +8164,7 @@ public void testSecondMemberPartiallyJoinAndTimeout() throws Exception {
         assertEquals(1, firstGenerationId);
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
-        SyncResult syncResult = context.sendClassicGroupSync(new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(firstMemberId)
             .withGenerationId(firstGenerationId)
@@ -9299,7 +8178,7 @@ public void testSecondMemberPartiallyJoinAndTimeout() throws Exception {
         assertTrue(group.isInState(STABLE));
 
         // Add a new dynamic pending member.
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             joinRequest
                 .setMemberId(UNKNOWN_MEMBER_ID)
                 .setGroupInstanceId(null)
@@ -9321,7 +8200,7 @@ public void testSecondMemberPartiallyJoinAndTimeout() throws Exception {
             .setGenerationId(firstGenerationId);
 
         for (int i = 0; i < 2; i++) {
-            assertNoOrEmptyResult(context.sleep(2500));
+            GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(2500));
             HeartbeatResponseData heartbeatResponse = context.sendClassicGroupHeartbeat(heartbeatRequest);
             assertEquals(Errors.NONE.code(), heartbeatResponse.errorCode());
         }
@@ -9347,7 +8226,7 @@ public void testRebalanceTimesOutWhenSyncRequestIsNotReceived() throws Exception
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Advance clock by 1/2 rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // Heartbeats to ensure that heartbeating does not interfere with the
         // delayed sync operation.
@@ -9360,7 +8239,7 @@ public void testRebalanceTimesOutWhenSyncRequestIsNotReceived() throws Exception
         ExpiredTimeout<Void, Record> timeout = timeouts.get(0);
         assertEquals(classicGroupSyncKey("group-id"), timeout.key);
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             timeout.result.records()
         );
 
@@ -9385,14 +8264,14 @@ public void testRebalanceTimesOutWhenSyncRequestIsNotReceivedFromFollowers() thr
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Advance clock by 1/2 rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // Heartbeats to ensure that heartbeating does not interfere with the
         // delayed sync operation.
         joinResponses.forEach(response -> context.verifyHeartbeat(group.groupId(), response, Errors.NONE));
 
         // Leader sends a sync group request.
-        SyncResult syncResult = context.sendClassicGroupSync(new SyncGroupRequestBuilder()
+        GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
             .withGroupId("group-id")
             .withGenerationId(1)
             .withMemberId(joinResponses.get(0).memberId())
@@ -9437,7 +8316,7 @@ public void testRebalanceTimesOutWhenSyncRequestIsNotReceivedFromLeaders() throw
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // Advance clock by 1/2 rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // Heartbeats to ensure that heartbeating does not interfere with the
         // delayed sync operation.
@@ -9446,8 +8325,8 @@ public void testRebalanceTimesOutWhenSyncRequestIsNotReceivedFromLeaders() throw
         // Followers send sync group requests.
         List<CompletableFuture<SyncGroupResponseData>> followerSyncFutures = joinResponses.subList(1, 3).stream()
             .map(response -> {
-                SyncResult syncResult = context.sendClassicGroupSync(
-                    new SyncGroupRequestBuilder()
+                GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+                    new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                         .withGroupId("group-id")
                         .withGenerationId(1)
                         .withMemberId(response.memberId())
@@ -9498,7 +8377,7 @@ public void testRebalanceDoesNotTimeOutWhenAllSyncAreReceived() throws Exception
         String leaderId = joinResponses.get(0).memberId();
 
         // Advance clock by 1/2 rebalance timeout.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // Heartbeats to ensure that heartbeating does not interfere with the
         // delayed sync operation.
@@ -9506,8 +8385,8 @@ public void testRebalanceDoesNotTimeOutWhenAllSyncAreReceived() throws Exception
 
         // All members send sync group requests.
         List<CompletableFuture<SyncGroupResponseData>> syncFutures = joinResponses.stream().map(response -> {
-            SyncResult syncResult = context.sendClassicGroupSync(
-                new SyncGroupRequestBuilder()
+            GroupMetadataManagerTestContext.SyncResult syncResult = context.sendClassicGroupSync(
+                new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
                     .withGroupId("group-id")
                     .withGenerationId(1)
                     .withMemberId(response.memberId())
@@ -9516,7 +8395,7 @@ public void testRebalanceDoesNotTimeOutWhenAllSyncAreReceived() throws Exception
 
             if (response.memberId().equals(leaderId)) {
                 assertEquals(
-                    Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+                    Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
                     syncResult.records
                 );
 
@@ -9534,13 +8413,13 @@ public void testRebalanceDoesNotTimeOutWhenAllSyncAreReceived() throws Exception
         }
 
         // Advance clock by 1/2 rebalance timeout. Pending sync should already have been cancelled.
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // All member heartbeats should succeed.
         joinResponses.forEach(response -> context.verifyHeartbeat(group.groupId(), response, Errors.NONE));
 
         // Advance clock a bit more
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs / 2));
 
         // All member heartbeats should succeed.
         joinResponses.forEach(response -> context.verifyHeartbeat(group.groupId(), response, Errors.NONE));
@@ -9555,7 +8434,7 @@ public void testHeartbeatDuringRebalanceCausesRebalanceInProgress() throws Excep
         ClassicGroup group = context.createClassicGroup("group-id");
 
         // First start up a group (with a slightly larger timeout to give us time to heartbeat when the rebalance starts)
-        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+        JoinGroupRequestData joinRequest = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
             .withGroupId("group-id")
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
@@ -9570,7 +8449,7 @@ public void testHeartbeatDuringRebalanceCausesRebalanceInProgress() throws Excep
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
         // Then join with a new consumer to trigger a rebalance
-        JoinResult joinResult = context.sendClassicGroupJoin(
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
             joinRequest.setMemberId(UNKNOWN_MEMBER_ID)
         );
 
@@ -9601,7 +8480,7 @@ public void testListGroups() {
             .build();
 
         // Create one classic group record.
-        context.replay(newGroupMetadataRecord(
+        context.replay(GroupMetadataManagerTestContext.newGroupMetadataRecord(
             classicGroupId,
             new GroupMetadataValue()
                 .setMembers(Collections.emptyList())
@@ -9760,7 +8639,7 @@ public void testConsumerGroupDescribeNoErrors() {
             new ConsumerGroupDescribeResponseData.DescribedGroup()
                 .setGroupEpoch(epoch)
                 .setGroupId(consumerGroupIds.get(1))
-                .setMembers(Arrays.asList(
+                .setMembers(Collections.singletonList(
                     memberBuilder.build().asConsumerGroupDescribeMember(
                         new Assignment(Collections.emptyMap()),
                         new MetadataImageBuilder().build().topics()
@@ -9872,7 +8751,7 @@ public void testDescribeGroupStable() {
             .setProtocol("range")
             .setCurrentStateTimestamp(context.time.milliseconds());
 
-        context.replay(newGroupMetadataRecord(
+        context.replay(GroupMetadataManagerTestContext.newGroupMetadataRecord(
             "group-id",
             groupMetadataValue,
             MetadataVersion.latestTesting()
@@ -9921,7 +8800,7 @@ public void testDescribeGroupRebalancing() throws Exception {
             .setProtocol("range")
             .setCurrentStateTimestamp(context.time.milliseconds());
 
-        context.replay(newGroupMetadataRecord(
+        context.replay(GroupMetadataManagerTestContext.newGroupMetadataRecord(
             "group-id",
             groupMetadataValue,
             MetadataVersion.latestTesting()
@@ -9961,182 +8840,13 @@ public void testDescribeGroupsGroupIdNotFoundException() {
         context.verifyDescribeGroupsReturnsDeadGroup("group-id");
     }
 
-    public static <T> void assertUnorderedListEquals(
-        List<T> expected,
-        List<T> actual
-    ) {
-        assertEquals(new HashSet<>(expected), new HashSet<>(actual));
-    }
-
-    private static void assertResponseEquals(
-        ConsumerGroupHeartbeatResponseData expected,
-        ConsumerGroupHeartbeatResponseData actual
-    ) {
-        if (!responseEquals(expected, actual)) {
-            assertionFailure()
-                .expected(expected)
-                .actual(actual)
-                .buildAndThrow();
-        }
-    }
-
-    private static boolean responseEquals(
-        ConsumerGroupHeartbeatResponseData expected,
-        ConsumerGroupHeartbeatResponseData actual
-    ) {
-        if (expected.throttleTimeMs() != actual.throttleTimeMs()) return false;
-        if (expected.errorCode() != actual.errorCode()) return false;
-        if (!Objects.equals(expected.errorMessage(), actual.errorMessage())) return false;
-        if (!Objects.equals(expected.memberId(), actual.memberId())) return false;
-        if (expected.memberEpoch() != actual.memberEpoch()) return false;
-        if (expected.heartbeatIntervalMs() != actual.heartbeatIntervalMs()) return false;
-        // Unordered comparison of the assignments.
-        return responseAssignmentEquals(expected.assignment(), actual.assignment());
-    }
-
-    private static boolean responseAssignmentEquals(
-        ConsumerGroupHeartbeatResponseData.Assignment expected,
-        ConsumerGroupHeartbeatResponseData.Assignment actual
-    ) {
-        if (expected == actual) return true;
-        if (expected == null) return false;
-        if (actual == null) return false;
-
-        return Objects.equals(fromAssignment(expected.topicPartitions()), fromAssignment(actual.topicPartitions()));
-    }
-
-    private static Map<Uuid, Set<Integer>> fromAssignment(
-        List<ConsumerGroupHeartbeatResponseData.TopicPartitions> assignment
-    ) {
-        if (assignment == null) return null;
-
-        Map<Uuid, Set<Integer>> assignmentMap = new HashMap<>();
-        assignment.forEach(topicPartitions -> {
-            assignmentMap.put(topicPartitions.topicId(), new HashSet<>(topicPartitions.partitions()));
-        });
-        return assignmentMap;
-    }
-
-    private static void assertRecordsEquals(
-        List<Record> expectedRecords,
-        List<Record> actualRecords
-    ) {
-        try {
-            assertEquals(expectedRecords.size(), actualRecords.size());
-
-            for (int i = 0; i < expectedRecords.size(); i++) {
-                Record expectedRecord = expectedRecords.get(i);
-                Record actualRecord = actualRecords.get(i);
-                assertRecordEquals(expectedRecord, actualRecord);
-            }
-        } catch (AssertionFailedError e) {
-            assertionFailure()
-                .expected(expectedRecords)
-                .actual(actualRecords)
-                .buildAndThrow();
-        }
-    }
-
-    private static void assertRecordEquals(
-        Record expected,
-        Record actual
-    ) {
-        try {
-            assertApiMessageAndVersionEquals(expected.key(), actual.key());
-            assertApiMessageAndVersionEquals(expected.value(), actual.value());
-        } catch (AssertionFailedError e) {
-            assertionFailure()
-                .expected(expected)
-                .actual(actual)
-                .buildAndThrow();
-        }
-    }
-
-    private static void assertApiMessageAndVersionEquals(
-        ApiMessageAndVersion expected,
-        ApiMessageAndVersion actual
-    ) {
-        if (expected == actual) return;
-
-        assertEquals(expected.version(), actual.version());
-
-        if (actual.message() instanceof ConsumerGroupCurrentMemberAssignmentValue) {
-            // The order of the topics stored in ConsumerGroupCurrentMemberAssignmentValue is not
-            // always guaranteed. Therefore, we need a special comparator.
-            ConsumerGroupCurrentMemberAssignmentValue expectedValue =
-                (ConsumerGroupCurrentMemberAssignmentValue) expected.message();
-            ConsumerGroupCurrentMemberAssignmentValue actualValue =
-                (ConsumerGroupCurrentMemberAssignmentValue) actual.message();
-
-            assertEquals(expectedValue.memberEpoch(), actualValue.memberEpoch());
-            assertEquals(expectedValue.previousMemberEpoch(), actualValue.previousMemberEpoch());
-            assertEquals(expectedValue.targetMemberEpoch(), actualValue.targetMemberEpoch());
-            assertEquals(expectedValue.error(), actualValue.error());
-            assertEquals(expectedValue.metadataVersion(), actualValue.metadataVersion());
-            assertEquals(expectedValue.metadataBytes(), actualValue.metadataBytes());
-
-            // We transform those to Maps before comparing them.
-            assertEquals(fromTopicPartitions(expectedValue.assignedPartitions()),
-                fromTopicPartitions(actualValue.assignedPartitions()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingRevocation()),
-                fromTopicPartitions(actualValue.partitionsPendingRevocation()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingAssignment()),
-                fromTopicPartitions(actualValue.partitionsPendingAssignment()));
-        } else {
-            assertEquals(expected.message(), actual.message());
-        }
-    }
-
-    private static Map<Uuid, Set<Integer>> fromTopicPartitions(
-        List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> assignment
-    ) {
-        Map<Uuid, Set<Integer>> assignmentMap = new HashMap<>();
-        assignment.forEach(topicPartitions -> {
-            assignmentMap.put(topicPartitions.topicId(), new HashSet<>(topicPartitions.partitions()));
-        });
-        return assignmentMap;
-    }
-
-    private static List<JoinGroupResponseMember> toJoinResponseMembers(ClassicGroup group) {
-        List<JoinGroupResponseMember> members = new ArrayList<>();
-        String protocolName = group.protocolName().get();
-        group.allMembers().forEach(member -> {
-            members.add(
-                new JoinGroupResponseMember()
-                    .setMemberId(member.memberId())
-                    .setGroupInstanceId(member.groupInstanceId().orElse(""))
-                    .setMetadata(member.metadata(protocolName))
-            );
-        });
-
-        return members;
-    }
-
-    private static void checkJoinGroupResponse(
-        JoinGroupResponseData expectedResponse,
-        JoinGroupResponseData actualResponse,
-        ClassicGroup group,
-        ClassicGroupState expectedState,
-        Set<String> expectedGroupInstanceIds
-    ) {
-        assertEquals(expectedResponse, actualResponse);
-        assertTrue(group.isInState(expectedState));
-
-        Set<String> groupInstanceIds = actualResponse.members()
-            .stream()
-            .map(JoinGroupResponseData.JoinGroupResponseMember::groupInstanceId)
-            .collect(Collectors.toSet());
-
-        assertEquals(expectedGroupInstanceIds, groupInstanceIds);
-    }
-
     @Test
     public void testGroupStuckInRebalanceTimeoutDueToNonjoinedStaticMember() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
         int longSessionTimeoutMs = 10000;
         int rebalanceTimeoutMs = 5000;
-        RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
+        GroupMetadataManagerTestContext.RebalanceResult rebalanceResult = context.staticMembersJoinAndRebalance(
             "group-id",
             "leader-instance-id",
             "follower-instance-id",
@@ -10146,8 +8856,8 @@ public void testGroupStuckInRebalanceTimeoutDueToNonjoinedStaticMember() throws
         ClassicGroup group = context.groupMetadataManager.getOrMaybeCreateClassicGroup("group-id", false);
 
         // New member joins
-        JoinResult joinResult = context.sendClassicGroupJoin(
-            new JoinGroupRequestBuilder()
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withProtocolSuperset()
@@ -10156,7 +8866,7 @@ public void testGroupStuckInRebalanceTimeoutDueToNonjoinedStaticMember() throws
         );
 
         // The new dynamic member has been elected as leader
-        assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
+        GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(rebalanceTimeoutMs));
         assertTrue(joinResult.joinFuture.isDone());
         assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
         assertEquals(joinResult.joinFuture.get().leader(), joinResult.joinFuture.get().memberId());
@@ -10271,7 +8981,7 @@ public void testLeaveGroupUnknownMemberIdExistingGroup() throws Exception {
         context.createClassicGroup("group-id");
 
         context.joinClassicGroupAsDynamicMemberAndCompleteJoin(
-            new JoinGroupRequestBuilder()
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
@@ -10328,7 +9038,7 @@ public void testValidLeaveGroup() throws Exception {
         ClassicGroup group = context.createClassicGroup("group-id");
 
         JoinGroupResponseData joinResponse = context.joinClassicGroupAsDynamicMemberAndCompleteJoin(
-            new JoinGroupRequestBuilder()
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
@@ -10345,7 +9055,7 @@ public void testValidLeaveGroup() throws Exception {
                 ))
         );
         assertEquals(
-            Collections.singletonList(newGroupMetadataRecordWithCurrentState(group, MetadataVersion.latestTesting())),
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
             leaveResult.records()
         );
         // Simulate a successful write to the log.
@@ -10369,7 +9079,7 @@ public void testLeaveGroupWithFencedInstanceId() throws Exception {
         context.createClassicGroup("group-id");
 
         context.joinClassicGroupAndCompleteJoin(
-            new JoinGroupRequestBuilder()
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("group-instance-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
@@ -10407,7 +9117,7 @@ public void testLeaveGroupStaticMemberWithUnknownMemberId() throws Exception {
         context.createClassicGroup("group-id");
 
         context.joinClassicGroupAndCompleteJoin(
-            new JoinGroupRequestBuilder()
+            new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
                 .withGroupId("group-id")
                 .withGroupInstanceId("group-instance-id")
                 .withMemberId(UNKNOWN_MEMBER_ID)
@@ -10601,7 +9311,7 @@ public void testJoinedMemberPendingMemberBatchLeaveGroup() throws Exception {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
-        PendingMemberGroupResult pendingMemberGroupResult = context.setupGroupWithPendingMember(group);
+        GroupMetadataManagerTestContext.PendingMemberGroupResult pendingMemberGroupResult = context.setupGroupWithPendingMember(group);
 
         CoordinatorResult<LeaveGroupResponseData, Record> leaveResult = context.sendClassicGroupLeave(
             new LeaveGroupRequestData()
@@ -10638,7 +9348,7 @@ public void testJoinedMemberPendingMemberBatchLeaveGroupWithUnknownMember() thro
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
         ClassicGroup group = context.createClassicGroup("group-id");
-        PendingMemberGroupResult pendingMemberGroupResult = context.setupGroupWithPendingMember(group);
+        GroupMetadataManagerTestContext.PendingMemberGroupResult pendingMemberGroupResult = context.setupGroupWithPendingMember(group);
 
         CoordinatorResult<LeaveGroupResponseData, Record> leaveResult = context.sendClassicGroupLeave(
             new LeaveGroupRequestData()
@@ -10680,7 +9390,7 @@ public void testJoinedMemberPendingMemberBatchLeaveGroupWithUnknownMember() thro
     public void testClassicGroupDelete() {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
-        ClassicGroup group = context.createClassicGroup("group-id");
+        context.createClassicGroup("group-id");
 
         List<Record> expectedRecords = Collections.singletonList(RecordHelpers.newGroupMetadataTombstoneRecord("group-id"));
         List<Record> records = new ArrayList<>();
@@ -10713,7 +9423,7 @@ public void testClassicGroupMaybeDelete() {
     public void testConsumerGroupDelete() {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group-id", true);
+        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group-id", true);
 
         List<Record> expectedRecords = Arrays.asList(
             RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id"),
@@ -10839,12 +9549,12 @@ public void testOnClassicGroupStateTransitionOnLoading() {
 
         // Even if there are more group metadata records loaded than tombstone records, the last replayed record
         // (tombstone in this test) is the latest state of the group. Hence, the overall metric count should be 0.
-        IntStream.range(0, 5).forEach(__ -> {
-            context.replay(RecordHelpers.newGroupMetadataRecord(group, Collections.emptyMap(), MetadataVersion.LATEST_PRODUCTION));
-        });
-        IntStream.range(0, 4).forEach(__ -> {
-            context.replay(RecordHelpers.newGroupMetadataTombstoneRecord("group-id"));
-        });
+        IntStream.range(0, 5).forEach(__ ->
+            context.replay(RecordHelpers.newGroupMetadataRecord(group, Collections.emptyMap(), MetadataVersion.LATEST_PRODUCTION))
+        );
+        IntStream.range(0, 4).forEach(__ ->
+            context.replay(RecordHelpers.newGroupMetadataTombstoneRecord("group-id"))
+        );
 
         verify(context.metrics, times(1)).onClassicGroupStateTransition(null, EMPTY);
         verify(context.metrics, times(1)).onClassicGroupStateTransition(EMPTY, null);
@@ -10883,9 +9593,9 @@ public void testOnConsumerGroupStateTransitionOnLoading() {
 
         // Even if there are more group epoch records loaded than tombstone records, the last replayed record
         // (tombstone in this test) is the latest state of the group. Hence, the overall metric count should be 0.
-        IntStream.range(0, 5).forEach(__ -> {
-            context.replay(RecordHelpers.newGroupEpochRecord("group-id", 0));
-        });
+        IntStream.range(0, 5).forEach(__ ->
+            context.replay(RecordHelpers.newGroupEpochRecord("group-id", 0))
+        );
         context.replay(RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id"));
         context.replay(RecordHelpers.newGroupEpochTombstoneRecord("group-id"));
         IntStream.range(0, 3).forEach(__ -> {
@@ -10897,19 +9607,45 @@ public void testOnConsumerGroupStateTransitionOnLoading() {
         verify(context.metrics, times(1)).onConsumerGroupStateTransition(ConsumerGroup.ConsumerGroupState.EMPTY, null);
     }
 
-    private static void assertNoOrEmptyResult(List<ExpiredTimeout<Void, Record>> timeouts) {
-        assertTrue(timeouts.size() <= 1);
-        timeouts.forEach(timeout -> assertEquals(EMPTY_RESULT, timeout.result));
+    private static void checkJoinGroupResponse(
+        JoinGroupResponseData expectedResponse,
+        JoinGroupResponseData actualResponse,
+        ClassicGroup group,
+        ClassicGroupState expectedState,
+        Set<String> expectedGroupInstanceIds
+    ) {
+        assertEquals(expectedResponse, actualResponse);
+        assertTrue(group.isInState(expectedState));
+
+        Set<String> groupInstanceIds = actualResponse.members()
+                                                     .stream()
+                                                     .map(JoinGroupResponseMember::groupInstanceId)
+                                                     .collect(Collectors.toSet());
+
+        assertEquals(expectedGroupInstanceIds, groupInstanceIds);
+    }
+
+    private static List<JoinGroupResponseMember> toJoinResponseMembers(ClassicGroup group) {
+        List<JoinGroupResponseMember> members = new ArrayList<>();
+        String protocolName = group.protocolName().get();
+        group.allMembers().forEach(member -> members.add(
+            new JoinGroupResponseMember()
+                .setMemberId(member.memberId())
+                .setGroupInstanceId(member.groupInstanceId().orElse(""))
+                .setMetadata(member.metadata(protocolName))
+        ));
+
+        return members;
     }
 
     private static List<String> verifyClassicGroupJoinResponses(
-        List<JoinResult> joinResults,
+        List<GroupMetadataManagerTestContext.JoinResult> joinResults,
         int expectedSuccessCount,
         Errors expectedFailure
     ) {
         int successCount = 0;
         List<String> memberIds = new ArrayList<>();
-        for (JoinResult joinResult : joinResults) {
+        for (GroupMetadataManagerTestContext.JoinResult joinResult : joinResults) {
             if (!joinResult.joinFuture.isDone()) {
                 fail("All responseFutures should be completed.");
             }
@@ -10929,240 +9665,4 @@ private static List<String> verifyClassicGroupJoinResponses(
         assertEquals(expectedSuccessCount, successCount);
         return memberIds;
     }
-
-    private static JoinGroupRequestProtocolCollection toProtocols(String... protocolNames) {
-        JoinGroupRequestProtocolCollection protocols = new JoinGroupRequestProtocolCollection(0);
-        List<String> topicNames = Arrays.asList("foo", "bar", "baz");
-        for (int i = 0; i < protocolNames.length; i++) {
-            protocols.add(new JoinGroupRequestProtocol()
-                .setName(protocolNames[i])
-                .setMetadata(ConsumerProtocol.serializeSubscription(new ConsumerPartitionAssignor.Subscription(
-                    Collections.singletonList(topicNames.get(i % topicNames.size())))).array())
-            );
-        }
-        return protocols;
-    }
-
-    private static Record newGroupMetadataRecord(
-        String groupId,
-        GroupMetadataValue value,
-        MetadataVersion metadataVersion
-    ) {
-        return new Record(
-            new ApiMessageAndVersion(
-                new GroupMetadataKey()
-                    .setGroup(groupId),
-                (short) 2
-            ),
-            new ApiMessageAndVersion(
-                value,
-                metadataVersion.groupMetadataValueVersion()
-            )
-        );
-    }
-
-    private static Record newGroupMetadataRecordWithCurrentState(
-        ClassicGroup group,
-        MetadataVersion metadataVersion
-    ) {
-        return RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), metadataVersion);
-    }
-
-    private static class JoinGroupRequestBuilder {
-        String groupId = null;
-        String groupInstanceId = null;
-        String memberId = null;
-        String protocolType = "consumer";
-        JoinGroupRequestProtocolCollection protocols = new JoinGroupRequestProtocolCollection(0);
-        int sessionTimeoutMs = 500;
-        int rebalanceTimeoutMs = 500;
-        String reason = null;
-
-        JoinGroupRequestBuilder withGroupId(String groupId) {
-            this.groupId = groupId;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withGroupInstanceId(String groupInstanceId) {
-            this.groupInstanceId = groupInstanceId;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withMemberId(String memberId) {
-            this.memberId = memberId;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withDefaultProtocolTypeAndProtocols() {
-            this.protocols = toProtocols("range");
-            return this;
-        }
-
-        JoinGroupRequestBuilder withProtocolSuperset() {
-            this.protocols = toProtocols("range", "roundrobin");
-            return this;
-        }
-
-        JoinGroupRequestBuilder withProtocolType(String protocolType) {
-            this.protocolType = protocolType;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withProtocols(JoinGroupRequestProtocolCollection protocols) {
-            this.protocols = protocols;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withRebalanceTimeoutMs(int rebalanceTimeoutMs) {
-            this.rebalanceTimeoutMs = rebalanceTimeoutMs;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withSessionTimeoutMs(int sessionTimeoutMs) {
-            this.sessionTimeoutMs = sessionTimeoutMs;
-            return this;
-        }
-
-        JoinGroupRequestBuilder withReason(String reason) {
-            this.reason = reason;
-            return this;
-        }
-
-        JoinGroupRequestData build() {
-            return new JoinGroupRequestData()
-                .setGroupId(groupId)
-                .setGroupInstanceId(groupInstanceId)
-                .setMemberId(memberId)
-                .setProtocolType(protocolType)
-                .setProtocols(protocols)
-                .setRebalanceTimeoutMs(rebalanceTimeoutMs)
-                .setSessionTimeoutMs(sessionTimeoutMs)
-                .setReason(reason);
-        }
-    }
-
-    private static class SyncGroupRequestBuilder {
-        String groupId = null;
-        String groupInstanceId = null;
-        String memberId = null;
-        String protocolType = "consumer";
-        String protocolName = "range";
-        int generationId = 0;
-        List<SyncGroupRequestAssignment> assignment = Collections.emptyList();
-
-        SyncGroupRequestBuilder withGroupId(String groupId) {
-            this.groupId = groupId;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withGroupInstanceId(String groupInstanceId) {
-            this.groupInstanceId = groupInstanceId;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withMemberId(String memberId) {
-            this.memberId = memberId;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withGenerationId(int generationId) {
-            this.generationId = generationId;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withProtocolType(String protocolType) {
-            this.protocolType = protocolType;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withProtocolName(String protocolName) {
-            this.protocolName = protocolName;
-            return this;
-        }
-
-        SyncGroupRequestBuilder withAssignment(List<SyncGroupRequestAssignment> assignment) {
-            this.assignment = assignment;
-            return this;
-        }
-
-
-        SyncGroupRequestData build() {
-            return new SyncGroupRequestData()
-                .setGroupId(groupId)
-                .setGroupInstanceId(groupInstanceId)
-                .setMemberId(memberId)
-                .setGenerationId(generationId)
-                .setProtocolType(protocolType)
-                .setProtocolName(protocolName)
-                .setAssignments(assignment);
-        }
-    }
-
-    private static class RebalanceResult {
-        int generationId;
-        String leaderId;
-        byte[] leaderAssignment;
-        String followerId;
-        byte[] followerAssignment;
-
-        RebalanceResult(
-            int generationId,
-            String leaderId,
-            byte[] leaderAssignment,
-            String followerId,
-            byte[] followerAssignment
-        ) {
-            this.generationId = generationId;
-            this.leaderId = leaderId;
-            this.leaderAssignment = leaderAssignment;
-            this.followerId = followerId;
-            this.followerAssignment = followerAssignment;
-        }
-    }
-
-    private static class PendingMemberGroupResult {
-        String leaderId;
-        String followerId;
-        JoinGroupResponseData pendingMemberResponse;
-
-        public PendingMemberGroupResult(
-            String leaderId,
-            String followerId,
-            JoinGroupResponseData pendingMemberResponse
-        ) {
-            this.leaderId = leaderId;
-            this.followerId = followerId;
-            this.pendingMemberResponse = pendingMemberResponse;
-        }
-    }
-
-    private static class JoinResult {
-        CompletableFuture<JoinGroupResponseData> joinFuture;
-        List<Record> records;
-        CompletableFuture<Void> appendFuture;
-
-        public JoinResult(
-            CompletableFuture<JoinGroupResponseData> joinFuture,
-            CoordinatorResult<Void, Record> coordinatorResult
-        ) {
-            this.joinFuture = joinFuture;
-            this.records = coordinatorResult.records();
-            this.appendFuture = coordinatorResult.appendFuture();
-        }
-    }
-
-    private static class SyncResult {
-        CompletableFuture<SyncGroupResponseData> syncFuture;
-        List<Record> records;
-        CompletableFuture<Void> appendFuture;
-
-        public SyncResult(
-            CompletableFuture<SyncGroupResponseData> syncFuture,
-            CoordinatorResult<Void, Record> coordinatorResult
-        ) {
-            this.syncFuture = syncFuture;
-            this.records = coordinatorResult.records();
-            this.appendFuture = coordinatorResult.appendFuture();
-        }
-    }
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
new file mode 100644
index 0000000000000..d2d00c1582405
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -0,0 +1,1277 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group;
+
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
+import org.apache.kafka.common.errors.UnknownMemberIdException;
+import org.apache.kafka.common.message.ConsumerGroupDescribeResponseData;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData;
+import org.apache.kafka.common.message.HeartbeatRequestData;
+import org.apache.kafka.common.message.HeartbeatResponseData;
+import org.apache.kafka.common.message.JoinGroupRequestData;
+import org.apache.kafka.common.message.JoinGroupResponseData;
+import org.apache.kafka.common.message.LeaveGroupRequestData;
+import org.apache.kafka.common.message.LeaveGroupResponseData;
+import org.apache.kafka.common.message.ListGroupsResponseData;
+import org.apache.kafka.common.message.SyncGroupRequestData;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+import org.apache.kafka.common.network.ClientInformation;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ApiMessage;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.RequestContext;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.coordinator.group.assignor.PartitionAssignor;
+import org.apache.kafka.coordinator.group.classic.ClassicGroup;
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberValue;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataKey;
+import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataValue;
+import org.apache.kafka.coordinator.group.generated.GroupMetadataKey;
+import org.apache.kafka.coordinator.group.generated.GroupMetadataValue;
+import org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetricsShard;
+import org.apache.kafka.coordinator.group.runtime.CoordinatorResult;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.server.common.ApiMessageAndVersion;
+import org.apache.kafka.server.common.MetadataVersion;
+import org.apache.kafka.timeline.SnapshotRegistry;
+
+import java.net.InetAddress;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.EMPTY_RESULT;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRevocationTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupSessionTimeoutKey;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.COMPLETING_REBALANCE;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.DEAD;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.EMPTY;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.PREPARING_REBALANCE;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.STABLE;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.mock;
+
+public class GroupMetadataManagerTestContext {
+
+    public static void assertNoOrEmptyResult(List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> timeouts) {
+        assertTrue(timeouts.size() <= 1);
+        timeouts.forEach(timeout -> assertEquals(EMPTY_RESULT, timeout.result));
+    }
+
+    public static JoinGroupRequestData.JoinGroupRequestProtocolCollection toProtocols(String... protocolNames) {
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection protocols = new JoinGroupRequestData.JoinGroupRequestProtocolCollection(0);
+        List<String> topicNames = Arrays.asList("foo", "bar", "baz");
+        for (int i = 0; i < protocolNames.length; i++) {
+            protocols.add(new JoinGroupRequestData.JoinGroupRequestProtocol()
+                .setName(protocolNames[i])
+                .setMetadata(ConsumerProtocol.serializeSubscription(new ConsumerPartitionAssignor.Subscription(
+                    Collections.singletonList(topicNames.get(i % topicNames.size())))).array())
+            );
+        }
+        return protocols;
+    }
+
+    public static Record newGroupMetadataRecord(
+        String groupId,
+        GroupMetadataValue value,
+        MetadataVersion metadataVersion
+    ) {
+        return new Record(
+            new ApiMessageAndVersion(
+                new GroupMetadataKey()
+                    .setGroup(groupId),
+                (short) 2
+            ),
+            new ApiMessageAndVersion(
+                value,
+                metadataVersion.groupMetadataValueVersion()
+            )
+        );
+    }
+
+    public static class RebalanceResult {
+        int generationId;
+        String leaderId;
+        byte[] leaderAssignment;
+        String followerId;
+        byte[] followerAssignment;
+
+        RebalanceResult(
+            int generationId,
+            String leaderId,
+            byte[] leaderAssignment,
+            String followerId,
+            byte[] followerAssignment
+        ) {
+            this.generationId = generationId;
+            this.leaderId = leaderId;
+            this.leaderAssignment = leaderAssignment;
+            this.followerId = followerId;
+            this.followerAssignment = followerAssignment;
+        }
+    }
+
+    public static class PendingMemberGroupResult {
+        String leaderId;
+        String followerId;
+        JoinGroupResponseData pendingMemberResponse;
+
+        public PendingMemberGroupResult(
+            String leaderId,
+            String followerId,
+            JoinGroupResponseData pendingMemberResponse
+        ) {
+            this.leaderId = leaderId;
+            this.followerId = followerId;
+            this.pendingMemberResponse = pendingMemberResponse;
+        }
+    }
+
+    public static class JoinResult {
+        CompletableFuture<JoinGroupResponseData> joinFuture;
+        List<Record> records;
+        CompletableFuture<Void> appendFuture;
+
+        public JoinResult(
+            CompletableFuture<JoinGroupResponseData> joinFuture,
+            CoordinatorResult<Void, Record> coordinatorResult
+        ) {
+            this.joinFuture = joinFuture;
+            this.records = coordinatorResult.records();
+            this.appendFuture = coordinatorResult.appendFuture();
+        }
+    }
+
+    public static class SyncResult {
+        CompletableFuture<SyncGroupResponseData> syncFuture;
+        List<Record> records;
+        CompletableFuture<Void> appendFuture;
+
+        public SyncResult(
+            CompletableFuture<SyncGroupResponseData> syncFuture,
+            CoordinatorResult<Void, Record> coordinatorResult
+        ) {
+            this.syncFuture = syncFuture;
+            this.records = coordinatorResult.records();
+            this.appendFuture = coordinatorResult.appendFuture();
+        }
+    }
+
+    public static class JoinGroupRequestBuilder {
+        String groupId = null;
+        String groupInstanceId = null;
+        String memberId = null;
+        String protocolType = "consumer";
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection protocols = new JoinGroupRequestData.JoinGroupRequestProtocolCollection(0);
+        int sessionTimeoutMs = 500;
+        int rebalanceTimeoutMs = 500;
+        String reason = null;
+
+        JoinGroupRequestBuilder withGroupId(String groupId) {
+            this.groupId = groupId;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withGroupInstanceId(String groupInstanceId) {
+            this.groupInstanceId = groupInstanceId;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withMemberId(String memberId) {
+            this.memberId = memberId;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withDefaultProtocolTypeAndProtocols() {
+            this.protocols = toProtocols("range");
+            return this;
+        }
+
+        JoinGroupRequestBuilder withProtocolSuperset() {
+            this.protocols = toProtocols("range", "roundrobin");
+            return this;
+        }
+
+        JoinGroupRequestBuilder withProtocolType(String protocolType) {
+            this.protocolType = protocolType;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withProtocols(JoinGroupRequestData.JoinGroupRequestProtocolCollection protocols) {
+            this.protocols = protocols;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withRebalanceTimeoutMs(int rebalanceTimeoutMs) {
+            this.rebalanceTimeoutMs = rebalanceTimeoutMs;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withSessionTimeoutMs(int sessionTimeoutMs) {
+            this.sessionTimeoutMs = sessionTimeoutMs;
+            return this;
+        }
+
+        JoinGroupRequestBuilder withReason(String reason) {
+            this.reason = reason;
+            return this;
+        }
+
+        JoinGroupRequestData build() {
+            return new JoinGroupRequestData()
+                .setGroupId(groupId)
+                .setGroupInstanceId(groupInstanceId)
+                .setMemberId(memberId)
+                .setProtocolType(protocolType)
+                .setProtocols(protocols)
+                .setRebalanceTimeoutMs(rebalanceTimeoutMs)
+                .setSessionTimeoutMs(sessionTimeoutMs)
+                .setReason(reason);
+        }
+    }
+
+    public static class SyncGroupRequestBuilder {
+        String groupId = null;
+        String groupInstanceId = null;
+        String memberId = null;
+        String protocolType = "consumer";
+        String protocolName = "range";
+        int generationId = 0;
+        List<SyncGroupRequestData.SyncGroupRequestAssignment> assignment = Collections.emptyList();
+
+        SyncGroupRequestBuilder withGroupId(String groupId) {
+            this.groupId = groupId;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withGroupInstanceId(String groupInstanceId) {
+            this.groupInstanceId = groupInstanceId;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withMemberId(String memberId) {
+            this.memberId = memberId;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withGenerationId(int generationId) {
+            this.generationId = generationId;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withProtocolType(String protocolType) {
+            this.protocolType = protocolType;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withProtocolName(String protocolName) {
+            this.protocolName = protocolName;
+            return this;
+        }
+
+        SyncGroupRequestBuilder withAssignment(List<SyncGroupRequestData.SyncGroupRequestAssignment> assignment) {
+            this.assignment = assignment;
+            return this;
+        }
+
+
+        SyncGroupRequestData build() {
+            return new SyncGroupRequestData()
+                .setGroupId(groupId)
+                .setGroupInstanceId(groupInstanceId)
+                .setMemberId(memberId)
+                .setGenerationId(generationId)
+                .setProtocolType(protocolType)
+                .setProtocolName(protocolName)
+                .setAssignments(assignment);
+        }
+    }
+
+    public static class Builder {
+        final private MockTime time = new MockTime();
+        final private MockCoordinatorTimer<Void, Record> timer = new MockCoordinatorTimer<>(time);
+        final private LogContext logContext = new LogContext();
+        final private SnapshotRegistry snapshotRegistry = new SnapshotRegistry(logContext);
+        private MetadataImage metadataImage;
+        private List<PartitionAssignor> consumerGroupAssignors = Collections.singletonList(new MockPartitionAssignor("range"));
+        final private List<ConsumerGroupBuilder> consumerGroupBuilders = new ArrayList<>();
+        private int consumerGroupMaxSize = Integer.MAX_VALUE;
+        private int consumerGroupMetadataRefreshIntervalMs = Integer.MAX_VALUE;
+        private int classicGroupMaxSize = Integer.MAX_VALUE;
+        private int classicGroupInitialRebalanceDelayMs = 3000;
+        final private int classicGroupNewMemberJoinTimeoutMs = 5 * 60 * 1000;
+        private int classicGroupMinSessionTimeoutMs = 10;
+        private int classicGroupMaxSessionTimeoutMs = 10 * 60 * 1000;
+        final private GroupCoordinatorMetricsShard metrics = mock(GroupCoordinatorMetricsShard.class);
+
+        public Builder withMetadataImage(MetadataImage metadataImage) {
+            this.metadataImage = metadataImage;
+            return this;
+        }
+
+        public Builder withAssignors(List<PartitionAssignor> assignors) {
+            this.consumerGroupAssignors = assignors;
+            return this;
+        }
+
+        public Builder withConsumerGroup(ConsumerGroupBuilder builder) {
+            this.consumerGroupBuilders.add(builder);
+            return this;
+        }
+
+        public Builder withConsumerGroupMaxSize(int consumerGroupMaxSize) {
+            this.consumerGroupMaxSize = consumerGroupMaxSize;
+            return this;
+        }
+
+        public Builder withConsumerGroupMetadataRefreshIntervalMs(int consumerGroupMetadataRefreshIntervalMs) {
+            this.consumerGroupMetadataRefreshIntervalMs = consumerGroupMetadataRefreshIntervalMs;
+            return this;
+        }
+
+        public Builder withClassicGroupMaxSize(int classicGroupMaxSize) {
+            this.classicGroupMaxSize = classicGroupMaxSize;
+            return this;
+        }
+
+        public Builder withClassicGroupInitialRebalanceDelayMs(int classicGroupInitialRebalanceDelayMs) {
+            this.classicGroupInitialRebalanceDelayMs = classicGroupInitialRebalanceDelayMs;
+            return this;
+        }
+
+        public Builder withClassicGroupMinSessionTimeoutMs(int classicGroupMinSessionTimeoutMs) {
+            this.classicGroupMinSessionTimeoutMs = classicGroupMinSessionTimeoutMs;
+            return this;
+        }
+
+        public Builder withClassicGroupMaxSessionTimeoutMs(int classicGroupMaxSessionTimeoutMs) {
+            this.classicGroupMaxSessionTimeoutMs = classicGroupMaxSessionTimeoutMs;
+            return this;
+        }
+
+        public GroupMetadataManagerTestContext build() {
+            if (metadataImage == null) metadataImage = MetadataImage.EMPTY;
+            if (consumerGroupAssignors == null) consumerGroupAssignors = Collections.emptyList();
+
+            GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext(
+                time,
+                timer,
+                snapshotRegistry,
+                metrics,
+                new GroupMetadataManager.Builder()
+                    .withSnapshotRegistry(snapshotRegistry)
+                    .withLogContext(logContext)
+                    .withTime(time)
+                    .withTimer(timer)
+                    .withMetadataImage(metadataImage)
+                    .withConsumerGroupHeartbeatInterval(5000)
+                    .withConsumerGroupSessionTimeout(45000)
+                    .withConsumerGroupMaxSize(consumerGroupMaxSize)
+                    .withConsumerGroupAssignors(consumerGroupAssignors)
+                    .withConsumerGroupMetadataRefreshIntervalMs(consumerGroupMetadataRefreshIntervalMs)
+                    .withClassicGroupMaxSize(classicGroupMaxSize)
+                    .withClassicGroupMinSessionTimeoutMs(classicGroupMinSessionTimeoutMs)
+                    .withClassicGroupMaxSessionTimeoutMs(classicGroupMaxSessionTimeoutMs)
+                    .withClassicGroupInitialRebalanceDelayMs(classicGroupInitialRebalanceDelayMs)
+                    .withClassicGroupNewMemberJoinTimeoutMs(classicGroupNewMemberJoinTimeoutMs)
+                    .withGroupCoordinatorMetricsShard(metrics)
+                    .build(),
+                classicGroupInitialRebalanceDelayMs,
+                classicGroupNewMemberJoinTimeoutMs
+            );
+
+            consumerGroupBuilders.forEach(builder -> builder.build(metadataImage.topics()).forEach(context::replay));
+
+            context.commit();
+
+            return context;
+        }
+    }
+
+    final MockTime time;
+    final MockCoordinatorTimer<Void, Record> timer;
+    final SnapshotRegistry snapshotRegistry;
+    final GroupCoordinatorMetricsShard metrics;
+    final GroupMetadataManager groupMetadataManager;
+    final int classicGroupInitialRebalanceDelayMs;
+    final int classicGroupNewMemberJoinTimeoutMs;
+
+    long lastCommittedOffset = 0L;
+    long lastWrittenOffset = 0L;
+
+    public GroupMetadataManagerTestContext(
+        MockTime time,
+        MockCoordinatorTimer<Void, Record> timer,
+        SnapshotRegistry snapshotRegistry,
+        GroupCoordinatorMetricsShard metrics,
+        GroupMetadataManager groupMetadataManager,
+        int classicGroupInitialRebalanceDelayMs,
+        int classicGroupNewMemberJoinTimeoutMs
+    ) {
+        this.time = time;
+        this.timer = timer;
+        this.snapshotRegistry = snapshotRegistry;
+        this.metrics = metrics;
+        this.groupMetadataManager = groupMetadataManager;
+        this.classicGroupInitialRebalanceDelayMs = classicGroupInitialRebalanceDelayMs;
+        this.classicGroupNewMemberJoinTimeoutMs = classicGroupNewMemberJoinTimeoutMs;
+        snapshotRegistry.getOrCreateSnapshot(lastWrittenOffset);
+    }
+
+    public void commit() {
+        long lastCommittedOffset = this.lastCommittedOffset;
+        this.lastCommittedOffset = lastWrittenOffset;
+        snapshotRegistry.deleteSnapshotsUpTo(lastCommittedOffset);
+    }
+
+    public void rollback() {
+        lastWrittenOffset = lastCommittedOffset;
+        snapshotRegistry.revertToSnapshot(lastCommittedOffset);
+    }
+
+    public ConsumerGroup.ConsumerGroupState consumerGroupState(
+        String groupId
+    ) {
+        return groupMetadataManager
+            .getOrMaybeCreateConsumerGroup(groupId, false)
+            .state();
+    }
+
+    public ConsumerGroupMember.MemberState consumerGroupMemberState(
+        String groupId,
+        String memberId
+    ) {
+        return groupMetadataManager
+            .getOrMaybeCreateConsumerGroup(groupId, false)
+            .getOrMaybeCreateMember(memberId, false)
+            .state();
+    }
+
+    public CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGroupHeartbeat(
+        ConsumerGroupHeartbeatRequestData request
+    ) {
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.CONSUMER_GROUP_HEARTBEAT,
+                ApiKeys.CONSUMER_GROUP_HEARTBEAT.latestVersion(),
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result = groupMetadataManager.consumerGroupHeartbeat(
+            context,
+            request
+        );
+
+        result.records().forEach(this::replay);
+        return result;
+    }
+
+    public List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> sleep(long ms) {
+        time.sleep(ms);
+        List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> timeouts = timer.poll();
+        timeouts.forEach(timeout -> {
+            if (timeout.result.replayRecords()) {
+                timeout.result.records().forEach(this::replay);
+            }
+        });
+        return timeouts;
+    }
+
+    public void assertSessionTimeout(
+        String groupId,
+        String memberId,
+        long delayMs
+    ) {
+        MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
+            timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
+        assertNotNull(timeout);
+        assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
+    }
+
+    public void assertNoSessionTimeout(
+        String groupId,
+        String memberId
+    ) {
+        MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
+            timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
+        assertNull(timeout);
+    }
+
+    public MockCoordinatorTimer.ScheduledTimeout<Void, Record> assertRevocationTimeout(
+        String groupId,
+        String memberId,
+        long delayMs
+    ) {
+        MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
+            timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
+        assertNotNull(timeout);
+        assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
+        return timeout;
+    }
+
+    public void assertNoRevocationTimeout(
+        String groupId,
+        String memberId
+    ) {
+        MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
+            timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
+        assertNull(timeout);
+    }
+
+    ClassicGroup createClassicGroup(String groupId) {
+        return groupMetadataManager.getOrMaybeCreateClassicGroup(groupId, true);
+    }
+
+    public JoinResult sendClassicGroupJoin(
+        JoinGroupRequestData request
+    ) {
+        return sendClassicGroupJoin(request, false);
+    }
+
+    public JoinResult sendClassicGroupJoin(
+        JoinGroupRequestData request,
+        boolean requireKnownMemberId
+    ) {
+        return sendClassicGroupJoin(request, requireKnownMemberId, false);
+    }
+
+    public JoinResult sendClassicGroupJoin(
+        JoinGroupRequestData request,
+        boolean requireKnownMemberId,
+        boolean supportSkippingAssignment
+    ) {
+        // requireKnownMemberId is true: version >= 4 (See JoinGroupRequest#requiresKnownMemberId())
+        // supportSkippingAssignment is true: version >= 9 (See JoinGroupRequest#supportsSkippingAssignment())
+        short joinGroupVersion = 3;
+
+        if (requireKnownMemberId) {
+            joinGroupVersion = 4;
+            if (supportSkippingAssignment) {
+                joinGroupVersion = ApiKeys.JOIN_GROUP.latestVersion();
+            }
+        }
+
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.JOIN_GROUP,
+                joinGroupVersion,
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        CompletableFuture<JoinGroupResponseData> responseFuture = new CompletableFuture<>();
+        CoordinatorResult<Void, Record> coordinatorResult = groupMetadataManager.classicGroupJoin(
+            context,
+            request,
+            responseFuture
+        );
+
+        return new JoinResult(responseFuture, coordinatorResult);
+    }
+
+    public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteRebalance(
+        String groupId
+    ) throws Exception {
+        ClassicGroup group = createClassicGroup(groupId);
+
+        JoinGroupResponseData leaderJoinResponse =
+            joinClassicGroupAsDynamicMemberAndCompleteJoin(new JoinGroupRequestBuilder()
+                .withGroupId("group-id")
+                .withMemberId(UNKNOWN_MEMBER_ID)
+                .withDefaultProtocolTypeAndProtocols()
+                .withRebalanceTimeoutMs(10000)
+                .withSessionTimeoutMs(5000)
+                .build());
+
+        assertEquals(1, leaderJoinResponse.generationId());
+        assertTrue(group.isInState(COMPLETING_REBALANCE));
+
+        SyncResult syncResult = sendClassicGroupSync(new SyncGroupRequestBuilder()
+            .withGroupId("group-id")
+            .withMemberId(leaderJoinResponse.memberId())
+            .withGenerationId(leaderJoinResponse.generationId())
+            .build());
+
+        assertEquals(
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
+            syncResult.records
+        );
+        // Simulate a successful write to the log.
+        syncResult.appendFuture.complete(null);
+
+        assertTrue(syncResult.syncFuture.isDone());
+        assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
+        assertTrue(group.isInState(STABLE));
+
+        return leaderJoinResponse;
+    }
+
+    public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteJoin(
+        JoinGroupRequestData request
+    ) throws ExecutionException, InterruptedException {
+        boolean requireKnownMemberId = true;
+        String newMemberId = request.memberId();
+
+        if (request.memberId().equals(UNKNOWN_MEMBER_ID)) {
+            // Since member id is required, we need another round to get the successful join group result.
+            JoinResult firstJoinResult = sendClassicGroupJoin(
+                request,
+                requireKnownMemberId
+            );
+            assertTrue(firstJoinResult.records.isEmpty());
+            assertTrue(firstJoinResult.joinFuture.isDone());
+            assertEquals(Errors.MEMBER_ID_REQUIRED.code(), firstJoinResult.joinFuture.get().errorCode());
+            newMemberId = firstJoinResult.joinFuture.get().memberId();
+        }
+
+        // Second round
+        JoinGroupRequestData secondRequest = new JoinGroupRequestData()
+            .setGroupId(request.groupId())
+            .setMemberId(newMemberId)
+            .setProtocolType(request.protocolType())
+            .setProtocols(request.protocols())
+            .setSessionTimeoutMs(request.sessionTimeoutMs())
+            .setRebalanceTimeoutMs(request.rebalanceTimeoutMs())
+            .setReason(request.reason());
+
+        JoinResult secondJoinResult = sendClassicGroupJoin(
+            secondRequest,
+            requireKnownMemberId
+        );
+
+        assertTrue(secondJoinResult.records.isEmpty());
+        List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> timeouts = sleep(classicGroupInitialRebalanceDelayMs);
+        assertEquals(1, timeouts.size());
+        assertTrue(secondJoinResult.joinFuture.isDone());
+        assertEquals(Errors.NONE.code(), secondJoinResult.joinFuture.get().errorCode());
+
+        return secondJoinResult.joinFuture.get();
+    }
+
+    public JoinGroupResponseData joinClassicGroupAndCompleteJoin(
+        JoinGroupRequestData request,
+        boolean requireKnownMemberId,
+        boolean supportSkippingAssignment
+    ) throws ExecutionException, InterruptedException {
+        return joinClassicGroupAndCompleteJoin(
+            request,
+            requireKnownMemberId,
+            supportSkippingAssignment,
+            classicGroupInitialRebalanceDelayMs
+        );
+    }
+
+    public JoinGroupResponseData joinClassicGroupAndCompleteJoin(
+        JoinGroupRequestData request,
+        boolean requireKnownMemberId,
+        boolean supportSkippingAssignment,
+        int advanceClockMs
+    ) throws ExecutionException, InterruptedException {
+        if (requireKnownMemberId && request.groupInstanceId().isEmpty()) {
+            return joinClassicGroupAsDynamicMemberAndCompleteJoin(request);
+        }
+
+        try {
+            JoinResult joinResult = sendClassicGroupJoin(
+                request,
+                requireKnownMemberId,
+                supportSkippingAssignment
+            );
+
+            sleep(advanceClockMs);
+            assertTrue(joinResult.joinFuture.isDone());
+            assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
+            return joinResult.joinFuture.get();
+        } catch (Exception e) {
+            fail("Failed to due: " + e.getMessage());
+        }
+        return null;
+    }
+
+    public SyncResult sendClassicGroupSync(SyncGroupRequestData request) {
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.SYNC_GROUP,
+                ApiKeys.SYNC_GROUP.latestVersion(),
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        CompletableFuture<SyncGroupResponseData> responseFuture = new CompletableFuture<>();
+
+        CoordinatorResult<Void, Record> coordinatorResult = groupMetadataManager.classicGroupSync(
+            context,
+            request,
+            responseFuture
+        );
+
+        return new SyncResult(responseFuture, coordinatorResult);
+    }
+
+    public RebalanceResult staticMembersJoinAndRebalance(
+        String groupId,
+        String leaderInstanceId,
+        String followerInstanceId
+    ) throws Exception {
+        return staticMembersJoinAndRebalance(
+            groupId,
+            leaderInstanceId,
+            followerInstanceId,
+            10000,
+            5000
+        );
+    }
+
+    public RebalanceResult staticMembersJoinAndRebalance(
+        String groupId,
+        String leaderInstanceId,
+        String followerInstanceId,
+        int rebalanceTimeoutMs,
+        int sessionTimeoutMs
+    ) throws Exception {
+        ClassicGroup group = createClassicGroup("group-id");
+
+        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+            .withGroupId(groupId)
+            .withGroupInstanceId(leaderInstanceId)
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withProtocolType("consumer")
+            .withProtocolSuperset()
+            .withRebalanceTimeoutMs(rebalanceTimeoutMs)
+            .withSessionTimeoutMs(sessionTimeoutMs)
+            .build();
+
+        JoinResult leaderJoinResult = sendClassicGroupJoin(joinRequest);
+        JoinResult followerJoinResult = sendClassicGroupJoin(joinRequest.setGroupInstanceId(followerInstanceId));
+
+        assertTrue(leaderJoinResult.records.isEmpty());
+        assertTrue(followerJoinResult.records.isEmpty());
+        assertFalse(leaderJoinResult.joinFuture.isDone());
+        assertFalse(followerJoinResult.joinFuture.isDone());
+
+        // The goal for two timer advance is to let first group initial join complete and set newMemberAdded flag to false. Next advance is
+        // to trigger the rebalance as needed for follower delayed join. One large time advance won't help because we could only populate one
+        // delayed join from purgatory and the new delayed op is created at that time and never be triggered.
+        assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
+        assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
+
+        assertTrue(leaderJoinResult.joinFuture.isDone());
+        assertTrue(followerJoinResult.joinFuture.isDone());
+        assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
+        assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
+        assertEquals(1, leaderJoinResult.joinFuture.get().generationId());
+        assertEquals(1, followerJoinResult.joinFuture.get().generationId());
+        assertEquals(2, group.size());
+        assertEquals(1, group.generationId());
+        assertTrue(group.isInState(COMPLETING_REBALANCE));
+
+        String leaderId = leaderJoinResult.joinFuture.get().memberId();
+        String followerId = followerJoinResult.joinFuture.get().memberId();
+        List<SyncGroupRequestData.SyncGroupRequestAssignment> assignment = new ArrayList<>();
+        assignment.add(new SyncGroupRequestData.SyncGroupRequestAssignment().setMemberId(leaderId)
+                                                                            .setAssignment(new byte[]{1}));
+        assignment.add(new SyncGroupRequestData.SyncGroupRequestAssignment().setMemberId(followerId)
+                                                                            .setAssignment(new byte[]{2}));
+
+        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+            .withGroupId(groupId)
+            .withGroupInstanceId(leaderInstanceId)
+            .withMemberId(leaderId)
+            .withGenerationId(1)
+            .withAssignment(assignment)
+            .build();
+
+        SyncResult leaderSyncResult = sendClassicGroupSync(syncRequest);
+
+        // The generated record should contain the new assignment.
+        Map<String, byte[]> groupAssignment = assignment.stream().collect(Collectors.toMap(
+            SyncGroupRequestData.SyncGroupRequestAssignment::memberId, SyncGroupRequestData.SyncGroupRequestAssignment::assignment
+        ));
+        assertEquals(
+            Collections.singletonList(
+                RecordHelpers.newGroupMetadataRecord(group, groupAssignment, MetadataVersion.latestTesting())),
+            leaderSyncResult.records
+        );
+
+        // Simulate a successful write to the log.
+        leaderSyncResult.appendFuture.complete(null);
+
+        assertTrue(leaderSyncResult.syncFuture.isDone());
+        assertEquals(Errors.NONE.code(), leaderSyncResult.syncFuture.get().errorCode());
+        assertTrue(group.isInState(STABLE));
+
+        SyncResult followerSyncResult = sendClassicGroupSync(
+            syncRequest.setGroupInstanceId(followerInstanceId)
+                       .setMemberId(followerId)
+                       .setAssignments(Collections.emptyList())
+        );
+
+        assertTrue(followerSyncResult.records.isEmpty());
+        assertTrue(followerSyncResult.syncFuture.isDone());
+        assertEquals(Errors.NONE.code(), followerSyncResult.syncFuture.get().errorCode());
+        assertTrue(group.isInState(STABLE));
+
+        assertEquals(2, group.size());
+        assertEquals(1, group.generationId());
+
+        return new RebalanceResult(
+            1,
+            leaderId,
+            leaderSyncResult.syncFuture.get().assignment(),
+            followerId,
+            followerSyncResult.syncFuture.get().assignment()
+        );
+    }
+
+    public PendingMemberGroupResult setupGroupWithPendingMember(ClassicGroup group) throws Exception {
+        // Add the first member
+        JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
+            .withGroupId("group-id")
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withDefaultProtocolTypeAndProtocols()
+            .withRebalanceTimeoutMs(10000)
+            .withSessionTimeoutMs(5000)
+            .build();
+
+        JoinGroupResponseData leaderJoinResponse =
+            joinClassicGroupAsDynamicMemberAndCompleteJoin(joinRequest);
+
+        List<SyncGroupRequestData.SyncGroupRequestAssignment> assignment = new ArrayList<>();
+        assignment.add(new SyncGroupRequestData.SyncGroupRequestAssignment().setMemberId(leaderJoinResponse.memberId()));
+        SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
+            .withGroupId("group-id")
+            .withMemberId(leaderJoinResponse.memberId())
+            .withGenerationId(leaderJoinResponse.generationId())
+            .withAssignment(assignment)
+            .build();
+
+        SyncResult syncResult = sendClassicGroupSync(syncRequest);
+
+        // Now the group is stable, with the one member that joined above
+        assertEquals(
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
+            syncResult.records
+        );
+        // Simulate a successful write to log.
+        syncResult.appendFuture.complete(null);
+
+        assertTrue(syncResult.syncFuture.isDone());
+        assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
+
+        // Start the join for the second member
+        JoinResult followerJoinResult = sendClassicGroupJoin(
+            joinRequest.setMemberId(UNKNOWN_MEMBER_ID)
+        );
+
+        assertTrue(followerJoinResult.records.isEmpty());
+        assertFalse(followerJoinResult.joinFuture.isDone());
+
+        JoinResult leaderJoinResult = sendClassicGroupJoin(
+            joinRequest.setMemberId(leaderJoinResponse.memberId())
+        );
+
+        assertTrue(leaderJoinResult.records.isEmpty());
+        assertTrue(group.isInState(COMPLETING_REBALANCE));
+        assertTrue(leaderJoinResult.joinFuture.isDone());
+        assertTrue(followerJoinResult.joinFuture.isDone());
+        assertEquals(Errors.NONE.code(), leaderJoinResult.joinFuture.get().errorCode());
+        assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
+        assertEquals(leaderJoinResult.joinFuture.get().generationId(), followerJoinResult.joinFuture.get().generationId());
+        assertEquals(leaderJoinResponse.memberId(), leaderJoinResult.joinFuture.get().leader());
+        assertEquals(leaderJoinResponse.memberId(), followerJoinResult.joinFuture.get().leader());
+
+        int nextGenerationId = leaderJoinResult.joinFuture.get().generationId();
+        String followerId = followerJoinResult.joinFuture.get().memberId();
+
+        // Stabilize the group
+        syncResult = sendClassicGroupSync(syncRequest.setGenerationId(nextGenerationId));
+
+        assertEquals(
+            Collections.singletonList(RecordHelpers.newGroupMetadataRecord(group, group.groupAssignment(), MetadataVersion.latestTesting())),
+            syncResult.records
+        );
+        // Simulate a successful write to log.
+        syncResult.appendFuture.complete(null);
+
+        assertTrue(syncResult.syncFuture.isDone());
+        assertEquals(Errors.NONE.code(), syncResult.syncFuture.get().errorCode());
+        assertTrue(group.isInState(STABLE));
+
+        // Re-join an existing member, to transition the group to PreparingRebalance state.
+        leaderJoinResult = sendClassicGroupJoin(
+            joinRequest.setMemberId(leaderJoinResponse.memberId()));
+
+        assertTrue(leaderJoinResult.records.isEmpty());
+        assertFalse(leaderJoinResult.joinFuture.isDone());
+        assertTrue(group.isInState(PREPARING_REBALANCE));
+
+        // Create a pending member in the group
+        JoinResult pendingMemberJoinResult = sendClassicGroupJoin(
+            joinRequest
+                .setMemberId(UNKNOWN_MEMBER_ID)
+                .setSessionTimeoutMs(2500),
+            true
+        );
+
+        assertTrue(pendingMemberJoinResult.records.isEmpty());
+        assertTrue(pendingMemberJoinResult.joinFuture.isDone());
+        assertEquals(Errors.MEMBER_ID_REQUIRED.code(), pendingMemberJoinResult.joinFuture.get().errorCode());
+        assertEquals(1, group.numPendingJoinMembers());
+
+        // Re-join the second existing member
+        followerJoinResult = sendClassicGroupJoin(
+            joinRequest.setMemberId(followerId).setSessionTimeoutMs(5000)
+        );
+
+        assertTrue(followerJoinResult.records.isEmpty());
+        assertFalse(followerJoinResult.joinFuture.isDone());
+        assertTrue(group.isInState(PREPARING_REBALANCE));
+        assertEquals(2, group.size());
+        assertEquals(1, group.numPendingJoinMembers());
+
+        return new PendingMemberGroupResult(
+            leaderJoinResponse.memberId(),
+            followerId,
+            pendingMemberJoinResult.joinFuture.get()
+        );
+    }
+
+    public void verifySessionExpiration(ClassicGroup group, int timeoutMs) {
+        Set<String> expectedHeartbeatKeys = group.allMembers().stream()
+                                                 .map(member -> classicGroupHeartbeatKey(group.groupId(), member.memberId())).collect(Collectors.toSet());
+
+        // Member should be removed as session expires.
+        List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> timeouts = sleep(timeoutMs);
+        List<Record> expectedRecords = Collections.singletonList(newGroupMetadataRecord(
+            group.groupId(),
+            new GroupMetadataValue()
+                .setMembers(Collections.emptyList())
+                .setGeneration(group.generationId())
+                .setLeader(null)
+                .setProtocolType("consumer")
+                .setProtocol(null)
+                .setCurrentStateTimestamp(time.milliseconds()),
+            MetadataVersion.latestTesting()
+        ));
+
+
+        Set<String> heartbeatKeys = timeouts.stream().map(timeout -> timeout.key).collect(Collectors.toSet());
+        assertEquals(expectedHeartbeatKeys, heartbeatKeys);
+
+        // Only the last member leaving the group should result in the empty group metadata record.
+        int timeoutsSize = timeouts.size();
+        assertEquals(expectedRecords, timeouts.get(timeoutsSize - 1).result.records());
+        assertNoOrEmptyResult(timeouts.subList(0, timeoutsSize - 1));
+        assertTrue(group.isInState(EMPTY));
+        assertEquals(0, group.size());
+    }
+
+    public HeartbeatResponseData sendClassicGroupHeartbeat(
+        HeartbeatRequestData request
+    ) {
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.HEARTBEAT,
+                ApiKeys.HEARTBEAT.latestVersion(),
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        return groupMetadataManager.classicGroupHeartbeat(
+            context,
+            request
+        );
+    }
+
+    public List<ListGroupsResponseData.ListedGroup> sendListGroups(List<String> statesFilter, List<String> typesFilter) {
+        Set<String> statesFilterSet = new HashSet<>(statesFilter);
+        Set<String> typesFilterSet = new HashSet<>(typesFilter);
+        return groupMetadataManager.listGroups(statesFilterSet, typesFilterSet, lastCommittedOffset);
+    }
+
+    public List<ConsumerGroupDescribeResponseData.DescribedGroup> sendConsumerGroupDescribe(List<String> groupIds) {
+        return groupMetadataManager.consumerGroupDescribe(groupIds, lastCommittedOffset);
+    }
+
+    public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(List<String> groupIds) {
+        return groupMetadataManager.describeGroups(groupIds, lastCommittedOffset);
+    }
+
+    public void verifyHeartbeat(
+        String groupId,
+        JoinGroupResponseData joinResponse,
+        Errors expectedError
+    ) {
+        HeartbeatRequestData request = new HeartbeatRequestData()
+            .setGroupId(groupId)
+            .setMemberId(joinResponse.memberId())
+            .setGenerationId(joinResponse.generationId());
+
+        if (expectedError == Errors.UNKNOWN_MEMBER_ID) {
+            assertThrows(UnknownMemberIdException.class, () -> sendClassicGroupHeartbeat(request));
+        } else {
+            HeartbeatResponseData response = sendClassicGroupHeartbeat(request);
+            assertEquals(expectedError.code(), response.errorCode());
+        }
+    }
+
+    public List<JoinGroupResponseData> joinWithNMembers(
+        String groupId,
+        int numMembers,
+        int rebalanceTimeoutMs,
+        int sessionTimeoutMs
+    ) {
+        ClassicGroup group = createClassicGroup(groupId);
+        boolean requireKnownMemberId = true;
+
+        // First join requests
+        JoinGroupRequestData request = new JoinGroupRequestBuilder()
+            .withGroupId(groupId)
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withDefaultProtocolTypeAndProtocols()
+            .withRebalanceTimeoutMs(rebalanceTimeoutMs)
+            .withSessionTimeoutMs(sessionTimeoutMs)
+            .build();
+
+        List<String> memberIds = IntStream.range(0, numMembers).mapToObj(i -> {
+            JoinResult joinResult = sendClassicGroupJoin(request, requireKnownMemberId);
+
+            assertTrue(joinResult.records.isEmpty());
+            assertTrue(joinResult.joinFuture.isDone());
+
+            try {
+                return joinResult.joinFuture.get().memberId();
+            } catch (Exception e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+            return null;
+        }).collect(Collectors.toList());
+
+        // Second join requests
+        List<CompletableFuture<JoinGroupResponseData>> secondJoinFutures = IntStream.range(0, numMembers).mapToObj(i -> {
+            JoinResult joinResult = sendClassicGroupJoin(request.setMemberId(memberIds.get(i)), requireKnownMemberId);
+
+            assertTrue(joinResult.records.isEmpty());
+            assertFalse(joinResult.joinFuture.isDone());
+
+            return joinResult.joinFuture;
+        }).collect(Collectors.toList());
+
+        // Advance clock by initial rebalance delay.
+        assertNoOrEmptyResult(sleep(classicGroupInitialRebalanceDelayMs));
+        secondJoinFutures.forEach(future -> assertFalse(future.isDone()));
+        // Advance clock by rebalance timeout to complete join phase.
+        assertNoOrEmptyResult(sleep(rebalanceTimeoutMs));
+
+        List<JoinGroupResponseData> joinResponses = secondJoinFutures.stream().map(future -> {
+            assertTrue(future.isDone());
+            try {
+                assertEquals(Errors.NONE.code(), future.get().errorCode());
+                return future.get();
+            } catch (Exception e) {
+                fail("Unexpected exception: " + e.getMessage());
+            }
+            return null;
+        }).collect(Collectors.toList());
+
+        assertEquals(numMembers, group.size());
+        assertTrue(group.isInState(COMPLETING_REBALANCE));
+
+        return joinResponses;
+    }
+
+    public CoordinatorResult<LeaveGroupResponseData, Record> sendClassicGroupLeave(
+        LeaveGroupRequestData request
+    ) {
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.LEAVE_GROUP,
+                ApiKeys.LEAVE_GROUP.latestVersion(),
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        return groupMetadataManager.classicGroupLeave(context, request);
+    }
+
+    public void verifyDescribeGroupsReturnsDeadGroup(String groupId) {
+        List<DescribeGroupsResponseData.DescribedGroup> describedGroups =
+            describeGroups(Collections.singletonList(groupId));
+
+        assertEquals(
+            Collections.singletonList(new DescribeGroupsResponseData.DescribedGroup()
+                .setGroupId("group-id")
+                .setGroupState(DEAD.toString())
+            ),
+            describedGroups
+        );
+    }
+
+    private ApiMessage messageOrNull(ApiMessageAndVersion apiMessageAndVersion) {
+        if (apiMessageAndVersion == null) {
+            return null;
+        } else {
+            return apiMessageAndVersion.message();
+        }
+    }
+
+    public void replay(
+        Record record
+    ) {
+        ApiMessageAndVersion key = record.key();
+        ApiMessageAndVersion value = record.value();
+
+        if (key == null) {
+            throw new IllegalStateException("Received a null key in " + record);
+        }
+
+        switch (key.version()) {
+            case GroupMetadataKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (GroupMetadataKey) key.message(),
+                    (GroupMetadataValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupMemberMetadataKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupMemberMetadataKey) key.message(),
+                    (ConsumerGroupMemberMetadataValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupMetadataKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupMetadataKey) key.message(),
+                    (ConsumerGroupMetadataValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupPartitionMetadataKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupPartitionMetadataKey) key.message(),
+                    (ConsumerGroupPartitionMetadataValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupTargetAssignmentMemberKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupTargetAssignmentMemberKey) key.message(),
+                    (ConsumerGroupTargetAssignmentMemberValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupTargetAssignmentMetadataKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupTargetAssignmentMetadataKey) key.message(),
+                    (ConsumerGroupTargetAssignmentMetadataValue) messageOrNull(value)
+                );
+                break;
+
+            case ConsumerGroupCurrentMemberAssignmentKey.HIGHEST_SUPPORTED_VERSION:
+                groupMetadataManager.replay(
+                    (ConsumerGroupCurrentMemberAssignmentKey) key.message(),
+                    (ConsumerGroupCurrentMemberAssignmentValue) messageOrNull(value)
+                );
+                break;
+
+            default:
+                throw new IllegalStateException("Received an unknown record type " + key.version()
+                    + " in " + record);
+        }
+
+        lastWrittenOffset++;
+        snapshotRegistry.getOrCreateSnapshot(lastWrittenOffset);
+    }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MetadataImageBuilder.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MetadataImageBuilder.java
new file mode 100644
index 0000000000000..995f1ee74a50b
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MetadataImageBuilder.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.metadata.PartitionRecord;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord;
+import org.apache.kafka.common.metadata.TopicRecord;
+import org.apache.kafka.image.MetadataDelta;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.MetadataProvenance;
+
+import java.util.Arrays;
+
+public class MetadataImageBuilder {
+    private MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY);
+
+    public MetadataImageBuilder addTopic(
+        Uuid topicId,
+        String topicName,
+        int numPartitions
+    ) {
+        // For testing purposes, the following criteria are used:
+        // - Number of replicas for each partition: 2
+        // - Number of brokers available in the cluster: 4
+        delta.replay(new TopicRecord().setTopicId(topicId).setName(topicName));
+        for (int i = 0; i < numPartitions; i++) {
+            delta.replay(new PartitionRecord()
+                .setTopicId(topicId)
+                .setPartitionId(i)
+                .setReplicas(Arrays.asList(i % 4, (i + 1) % 4)));
+        }
+        return this;
+    }
+
+    /**
+     * Add rack Ids for 4 broker Ids.
+     * <p>
+     * For testing purposes, each broker is mapped
+     * to a rack Id with the same broker Id as a suffix.
+     */
+    public MetadataImageBuilder addRacks() {
+        for (int i = 0; i < 4; i++) {
+            delta.replay(new RegisterBrokerRecord().setBrokerId(i).setRack("rack" + i));
+        }
+        return this;
+    }
+
+    public MetadataImage build() {
+        return delta.apply(MetadataProvenance.EMPTY);
+    }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MockPartitionAssignor.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MockPartitionAssignor.java
new file mode 100644
index 0000000000000..18bd6e51f22d4
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/MockPartitionAssignor.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group;
+
+import org.apache.kafka.coordinator.group.assignor.AssignmentSpec;
+import org.apache.kafka.coordinator.group.assignor.GroupAssignment;
+import org.apache.kafka.coordinator.group.assignor.PartitionAssignor;
+import org.apache.kafka.coordinator.group.assignor.PartitionAssignorException;
+import org.apache.kafka.coordinator.group.assignor.SubscribedTopicDescriber;
+
+public class MockPartitionAssignor implements PartitionAssignor {
+    private final String name;
+    private GroupAssignment prepareGroupAssignment = null;
+
+    MockPartitionAssignor(String name) {
+        this.name = name;
+    }
+
+    public void prepareGroupAssignment(GroupAssignment prepareGroupAssignment) {
+        this.prepareGroupAssignment = prepareGroupAssignment;
+    }
+
+    @Override
+    public String name() {
+        return name;
+    }
+
+    @Override
+    public GroupAssignment assign(AssignmentSpec assignmentSpec, SubscribedTopicDescriber subscribedTopicDescriber) throws PartitionAssignorException {
+        return prepareGroupAssignment;
+    }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
index 5ce52d560fbc3..87f7b471035fc 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
@@ -2362,7 +2362,7 @@ public void testConsumerGroupOffsetDeleteWithErrors() {
             "foo",
             true
         );
-        MetadataImage image = new GroupMetadataManagerTest.MetadataImageBuilder()
+        MetadataImage image = new MetadataImageBuilder()
             .addTopic(Uuid.randomUuid(), "foo", 1)
             .addRacks()
             .build();
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
index cb6c583a736e3..2a15fb7decd9c 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
@@ -53,7 +53,6 @@
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
-import org.opentest4j.AssertionFailedError;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
@@ -72,10 +71,10 @@
 import java.util.Set;
 import java.util.stream.Stream;
 
+import static org.apache.kafka.coordinator.group.Assertions.assertRecordEquals;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkSortedAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkSortedTopicAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
-import static org.apache.kafka.coordinator.group.GroupMetadataManagerTest.assertUnorderedListEquals;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newCurrentAssignmentRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newCurrentAssignmentTombstoneRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newGroupEpochRecord;
@@ -88,10 +87,8 @@
 import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentEpochTombstoneRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentTombstoneRecord;
-import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
 
 public class RecordHelpersTest {
@@ -866,143 +863,4 @@ public static Map<Integer, Set<String>> mkMapOfPartitionRacks(int numPartitions)
         }
         return partitionRacks;
     }
-
-    /**
-     * Asserts whether the two provided lists of records are equal.
-     *
-     * @param expectedRecords   The expected list of records.
-     * @param actualRecords     The actual list of records.
-     */
-    public static void assertRecordsEquals(
-        List<Record> expectedRecords,
-        List<Record> actualRecords
-    ) {
-        try {
-            assertEquals(expectedRecords.size(), actualRecords.size());
-
-            for (int i = 0; i < expectedRecords.size(); i++) {
-                Record expectedRecord = expectedRecords.get(i);
-                Record actualRecord = actualRecords.get(i);
-                assertRecordEquals(expectedRecord, actualRecord);
-            }
-        } catch (AssertionFailedError e) {
-            assertionFailure()
-                .expected(expectedRecords)
-                .actual(actualRecords)
-                .buildAndThrow();
-        }
-    }
-
-    /**
-     * Asserts if the two provided records are equal.
-     *
-     * @param expectedRecord   The expected record.
-     * @param actualRecord     The actual record.
-     */
-    public static void assertRecordEquals(
-        Record expectedRecord,
-        Record actualRecord
-    ) {
-        try {
-            assertApiMessageAndVersionEquals(expectedRecord.key(), actualRecord.key());
-            assertApiMessageAndVersionEquals(expectedRecord.value(), actualRecord.value());
-        } catch (AssertionFailedError e) {
-            assertionFailure()
-                .expected(expectedRecord)
-                .actual(actualRecord)
-                .buildAndThrow();
-        }
-    }
-
-    private static void assertApiMessageAndVersionEquals(
-        ApiMessageAndVersion expected,
-        ApiMessageAndVersion actual
-    ) {
-        if (expected == actual) return;
-
-        assertEquals(expected.version(), actual.version());
-
-        if (actual.message() instanceof ConsumerGroupCurrentMemberAssignmentValue) {
-            // The order of the topics stored in ConsumerGroupCurrentMemberAssignmentValue is not
-            // always guaranteed. Therefore, we need a special comparator.
-            ConsumerGroupCurrentMemberAssignmentValue expectedValue =
-                (ConsumerGroupCurrentMemberAssignmentValue) expected.message();
-            ConsumerGroupCurrentMemberAssignmentValue actualValue =
-                (ConsumerGroupCurrentMemberAssignmentValue) actual.message();
-
-            assertEquals(expectedValue.memberEpoch(), actualValue.memberEpoch());
-            assertEquals(expectedValue.previousMemberEpoch(), actualValue.previousMemberEpoch());
-            assertEquals(expectedValue.targetMemberEpoch(), actualValue.targetMemberEpoch());
-            assertEquals(expectedValue.error(), actualValue.error());
-            assertEquals(expectedValue.metadataVersion(), actualValue.metadataVersion());
-            assertEquals(expectedValue.metadataBytes(), actualValue.metadataBytes());
-
-            // We transform those to Maps before comparing them.
-            assertEquals(fromTopicPartitions(expectedValue.assignedPartitions()),
-                fromTopicPartitions(actualValue.assignedPartitions()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingRevocation()),
-                fromTopicPartitions(actualValue.partitionsPendingRevocation()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingAssignment()),
-                fromTopicPartitions(actualValue.partitionsPendingAssignment()));
-        } else if (actual.message() instanceof ConsumerGroupPartitionMetadataValue) {
-            // The order of the racks stored in the PartitionMetadata of the ConsumerGroupPartitionMetadataValue
-            // is not always guaranteed. Therefore, we need a special comparator.
-            ConsumerGroupPartitionMetadataValue expectedValue =
-                (ConsumerGroupPartitionMetadataValue) expected.message();
-            ConsumerGroupPartitionMetadataValue actualValue =
-                (ConsumerGroupPartitionMetadataValue) actual.message();
-
-            List<ConsumerGroupPartitionMetadataValue.TopicMetadata> expectedTopicMetadataList =
-                expectedValue.topics();
-            List<ConsumerGroupPartitionMetadataValue.TopicMetadata> actualTopicMetadataList =
-                actualValue.topics();
-
-            if (expectedTopicMetadataList.size() != actualTopicMetadataList.size()) {
-                fail("Topic metadata lists have different sizes");
-            }
-
-            for (int i = 0; i < expectedTopicMetadataList.size(); i++) {
-                ConsumerGroupPartitionMetadataValue.TopicMetadata expectedTopicMetadata =
-                    expectedTopicMetadataList.get(i);
-                ConsumerGroupPartitionMetadataValue.TopicMetadata actualTopicMetadata =
-                    actualTopicMetadataList.get(i);
-
-                assertEquals(expectedTopicMetadata.topicId(), actualTopicMetadata.topicId());
-                assertEquals(expectedTopicMetadata.topicName(), actualTopicMetadata.topicName());
-                assertEquals(expectedTopicMetadata.numPartitions(), actualTopicMetadata.numPartitions());
-
-                List<ConsumerGroupPartitionMetadataValue.PartitionMetadata> expectedPartitionMetadataList =
-                    expectedTopicMetadata.partitionMetadata();
-                List<ConsumerGroupPartitionMetadataValue.PartitionMetadata> actualPartitionMetadataList =
-                    actualTopicMetadata.partitionMetadata();
-
-                // If the list is empty, rack information wasn't available for any replica of
-                // the partition and hence, the entry wasn't added to the record.
-                if (expectedPartitionMetadataList.size() != actualPartitionMetadataList.size()) {
-                    fail("Partition metadata lists have different sizes");
-                } else if (!expectedPartitionMetadataList.isEmpty() && !actualPartitionMetadataList.isEmpty()) {
-                    for (int j = 0; j < expectedTopicMetadataList.size(); j++) {
-                        ConsumerGroupPartitionMetadataValue.PartitionMetadata expectedPartitionMetadata =
-                            expectedPartitionMetadataList.get(j);
-                        ConsumerGroupPartitionMetadataValue.PartitionMetadata actualPartitionMetadata =
-                            actualPartitionMetadataList.get(j);
-
-                        assertEquals(expectedPartitionMetadata.partition(), actualPartitionMetadata.partition());
-                        assertUnorderedListEquals(expectedPartitionMetadata.racks(), actualPartitionMetadata.racks());
-                    }
-                }
-            }
-        } else {
-            assertEquals(expected.message(), actual.message());
-        }
-    }
-
-    private static Map<Uuid, Set<Integer>> fromTopicPartitions(
-        List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> assignment
-    ) {
-        Map<Uuid, Set<Integer>> assignmentMap = new HashMap<>();
-        assignment.forEach(topicPartitions ->
-            assignmentMap.put(topicPartitions.topicId(), new HashSet<>(topicPartitions.partitions())));
-        return assignmentMap;
-    }
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupBuilder.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupBuilder.java
new file mode 100644
index 0000000000000..5b24e4c959d47
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupBuilder.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.consumer;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.Record;
+import org.apache.kafka.coordinator.group.RecordHelpers;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.image.TopicsImage;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class ConsumerGroupBuilder {
+    private final String groupId;
+    private final int groupEpoch;
+    private int assignmentEpoch;
+    private final Map<String, ConsumerGroupMember> members = new HashMap<>();
+    private final Map<String, Assignment> assignments = new HashMap<>();
+    private Map<String, TopicMetadata> subscriptionMetadata;
+
+    public ConsumerGroupBuilder(String groupId, int groupEpoch) {
+        this.groupId = groupId;
+        this.groupEpoch = groupEpoch;
+        this.assignmentEpoch = 0;
+    }
+
+    public ConsumerGroupBuilder withMember(ConsumerGroupMember member) {
+        this.members.put(member.memberId(), member);
+        return this;
+    }
+
+    public ConsumerGroupBuilder withSubscriptionMetadata(Map<String, TopicMetadata> subscriptionMetadata) {
+        this.subscriptionMetadata = subscriptionMetadata;
+        return this;
+    }
+
+    public ConsumerGroupBuilder withAssignment(String memberId, Map<Uuid, Set<Integer>> assignment) {
+        this.assignments.put(memberId, new Assignment(assignment));
+        return this;
+    }
+
+    public ConsumerGroupBuilder withAssignmentEpoch(int assignmentEpoch) {
+        this.assignmentEpoch = assignmentEpoch;
+        return this;
+    }
+
+    public List<Record> build(TopicsImage topicsImage) {
+        List<Record> records = new ArrayList<>();
+
+        // Add subscription records for members.
+        members.forEach((memberId, member) ->
+            records.add(RecordHelpers.newMemberSubscriptionRecord(groupId, member))
+        );
+
+        // Add subscription metadata.
+        if (subscriptionMetadata == null) {
+            subscriptionMetadata = new HashMap<>();
+            members.forEach((memberId, member) ->
+                member.subscribedTopicNames().forEach(topicName -> {
+                    TopicImage topicImage = topicsImage.getTopic(topicName);
+                    if (topicImage != null) {
+                        subscriptionMetadata.put(topicName, new TopicMetadata(
+                            topicImage.id(),
+                            topicImage.name(),
+                            topicImage.partitions().size(),
+                            Collections.emptyMap()
+                        ));
+                    }
+                })
+            );
+        }
+
+        if (!subscriptionMetadata.isEmpty()) {
+            records.add(RecordHelpers.newGroupSubscriptionMetadataRecord(groupId, subscriptionMetadata));
+        }
+
+        // Add group epoch record.
+        records.add(RecordHelpers.newGroupEpochRecord(groupId, groupEpoch));
+
+        // Add target assignment records.
+        assignments.forEach((memberId, assignment) ->
+            records.add(RecordHelpers.newTargetAssignmentRecord(groupId, memberId, assignment.partitions()))
+        );
+
+        // Add target assignment epoch.
+        records.add(RecordHelpers.newTargetAssignmentEpochRecord(groupId, assignmentEpoch));
+
+        // Add current assignment records for members.
+        members.forEach((memberId, member) ->
+            records.add(RecordHelpers.newCurrentAssignmentRecord(groupId, member))
+        );
+
+        return records;
+    }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
index 2433c7ef3a92a..a5daf3f92b372 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.message.ConsumerGroupDescribeResponseData;
-import org.apache.kafka.coordinator.group.GroupMetadataManagerTest;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
 import org.apache.kafka.image.MetadataImage;
@@ -320,7 +320,7 @@ public void testAsConsumerGroupDescribeMember() {
         Uuid topicId2 = Uuid.randomUuid();
         Uuid topicId3 = Uuid.randomUuid();
         Uuid topicId4 = Uuid.randomUuid();
-        MetadataImage metadataImage = new GroupMetadataManagerTest.MetadataImageBuilder()
+        MetadataImage metadataImage = new MetadataImageBuilder()
             .addTopic(topicId1, "topic1", 3)
             .addTopic(topicId2, "topic2", 3)
             .addTopic(topicId3, "topic3", 3)
@@ -398,7 +398,7 @@ public void testAsConsumerGroupDescribeWithTargetAssignmentNull() {
             .build();
 
         ConsumerGroupDescribeResponseData.Member consumerGroupDescribeMember = member.asConsumerGroupDescribeMember(
-            null, new GroupMetadataManagerTest.MetadataImageBuilder().build().topics());
+            null, new MetadataImageBuilder().build().topics());
 
         assertEquals(new ConsumerGroupDescribeResponseData.Assignment(), consumerGroupDescribeMember.targetAssignment());
     }
@@ -418,7 +418,7 @@ public void testAsConsumerGroupDescribeWithTopicNameNotFound() {
             .setMemberId(memberId.toString())
             .setSubscribedTopicRegex("");
         ConsumerGroupDescribeResponseData.Member actual = member.asConsumerGroupDescribeMember(null,
-            new GroupMetadataManagerTest.MetadataImageBuilder()
+            new MetadataImageBuilder()
                 .addTopic(Uuid.randomUuid(), "foo", 3)
                 .build().topics()
         );
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
index 59a76ed6dd8f1..7d90122c2dfb4 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
@@ -25,7 +25,7 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.coordinator.group.Group;
-import org.apache.kafka.coordinator.group.GroupMetadataManagerTest;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.OffsetAndMetadata;
 import org.apache.kafka.coordinator.group.OffsetExpirationCondition;
 import org.apache.kafka.coordinator.group.OffsetExpirationConditionImpl;
@@ -583,7 +583,7 @@ public void testUpdateSubscriptionMetadata() {
         Uuid barTopicId = Uuid.randomUuid();
         Uuid zarTopicId = Uuid.randomUuid();
 
-        MetadataImage image = new GroupMetadataManagerTest.MetadataImageBuilder()
+        MetadataImage image = new MetadataImageBuilder()
             .addTopic(fooTopicId, "foo", 1)
             .addTopic(barTopicId, "bar", 2)
             .addTopic(zarTopicId, "zar", 3)
@@ -921,7 +921,7 @@ public void testIsSubscribedToTopic() {
         Uuid fooTopicId = Uuid.randomUuid();
         Uuid barTopicId = Uuid.randomUuid();
 
-        MetadataImage image = new GroupMetadataManagerTest.MetadataImageBuilder()
+        MetadataImage image = new MetadataImageBuilder()
             .addTopic(fooTopicId, "foo", 1)
             .addTopic(barTopicId, "bar", 2)
             .addRacks()
@@ -992,7 +992,7 @@ public void testAsDescribedGroup() {
                     .setSubscribedTopicRegex("")
             ));
         ConsumerGroupDescribeResponseData.DescribedGroup actual = group.asDescribedGroup(1, "",
-            new GroupMetadataManagerTest.MetadataImageBuilder().build().topics());
+            new MetadataImageBuilder().build().topics());
 
         assertEquals(expected, actual);
     }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
index bdd7a354ca4c9..af94fede50257 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
@@ -33,6 +33,7 @@
 import java.util.Optional;
 import java.util.Set;
 
+import static org.apache.kafka.coordinator.group.Assertions.assertUnorderedListEquals;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentEpochRecord;
@@ -382,7 +383,7 @@ public void testAssignmentSwapped() {
 
         assertEquals(3, result.records().size());
 
-        assertUnorderedList(Arrays.asList(
+        assertUnorderedListEquals(Arrays.asList(
             newTargetAssignmentRecord("my-group", "member-1", mkAssignment(
                 mkTopicAssignment(fooTopicId, 4, 5, 6),
                 mkTopicAssignment(barTopicId, 4, 5, 6)
@@ -452,7 +453,7 @@ public void testNewMember() {
 
         assertEquals(4, result.records().size());
 
-        assertUnorderedList(Arrays.asList(
+        assertUnorderedListEquals(Arrays.asList(
             newTargetAssignmentRecord("my-group", "member-1", mkAssignment(
                 mkTopicAssignment(fooTopicId, 1, 2),
                 mkTopicAssignment(barTopicId, 1, 2)
@@ -539,7 +540,7 @@ public void testUpdateMember() {
 
         assertEquals(4, result.records().size());
 
-        assertUnorderedList(Arrays.asList(
+        assertUnorderedListEquals(Arrays.asList(
             newTargetAssignmentRecord("my-group", "member-1", mkAssignment(
                 mkTopicAssignment(fooTopicId, 1, 2),
                 mkTopicAssignment(barTopicId, 1, 2)
@@ -621,7 +622,7 @@ public void testPartialAssignmentUpdate() {
         assertEquals(3, result.records().size());
 
         // Member 1 has no record because its assignment did not change.
-        assertUnorderedList(Arrays.asList(
+        assertUnorderedListEquals(Arrays.asList(
             newTargetAssignmentRecord("my-group", "member-2", mkAssignment(
                 mkTopicAssignment(fooTopicId, 3, 4, 5),
                 mkTopicAssignment(barTopicId, 3, 4, 5)
@@ -695,7 +696,7 @@ public void testDeleteMember() {
 
         assertEquals(3, result.records().size());
 
-        assertUnorderedList(Arrays.asList(
+        assertUnorderedListEquals(Arrays.asList(
             newTargetAssignmentRecord("my-group", "member-1", mkAssignment(
                 mkTopicAssignment(fooTopicId, 1, 2, 3),
                 mkTopicAssignment(barTopicId, 1, 2, 3)
@@ -774,7 +775,7 @@ public void testStaticMemberReplace() {
 
         assertEquals(2, result.records().size());
 
-        assertUnorderedList(Collections.singletonList(
+        assertUnorderedListEquals(Collections.singletonList(
             newTargetAssignmentRecord("my-group", "member-3-a", mkAssignment(
                 mkTopicAssignment(fooTopicId, 5, 6),
                 mkTopicAssignment(barTopicId, 5, 6)
@@ -803,11 +804,4 @@ public void testStaticMemberReplace() {
 
         assertEquals(expectedAssignment, result.targetAssignment());
     }
-    private static <T> void assertUnorderedList(
-        List<T> expected,
-        List<T> actual
-    ) {
-        assertEquals(expected.size(), actual.size());
-        assertEquals(new HashSet<>(expected), new HashSet<>(actual));
-    }
 }

From e8c70fce26626ed2ab90f2728a45f6e55e907ec1 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Wed, 14 Feb 2024 16:09:48 +0100
Subject: [PATCH 037/258] KAFKA-16155: Re-enable testAutoCommitIntercept
 (#15334)

The main bug causing this test to fail as described in the ticket was already fixed.

The test is still flaky if unchanged, because in the new consumer, the assignment can
change in between two polls. Interceptors are only executed inside poll (and have to be,
since they must run as part of the application thread), so we need to modify the
integration test to call poll once after observing that the assignment changed.

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 .../integration/kafka/api/PlaintextConsumerTest.scala | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 23d522c41ae4d..74009fb6e0ac3 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -1338,9 +1338,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     MockProducerInterceptor.resetCounters()
   }
 
-  // This is disabled for the the consumer group until KAFKA-16155 is resolved.
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testAutoCommitIntercept(quorum: String, groupProtocol: String): Unit = {
     val topic2 = "topic2"
     createTopic(topic2, 2, brokerCount)
@@ -1378,6 +1377,14 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // after rebalancing, we should have reset to the committed positions
     assertEquals(10, testConsumer.committed(Set(tp).asJava).get(tp).offset)
     assertEquals(20, testConsumer.committed(Set(tp2).asJava).get(tp2).offset)
+
+    // In both CLASSIC and CONSUMER protocols, interceptors are executed in poll and close.
+    // However, in the CONSUMER protocol, the assignment may be changed outside of a poll, so
+    // we need to poll once to ensure the interceptor is called.
+    if (groupProtocol.toUpperCase == GroupProtocol.CONSUMER.name) {
+      testConsumer.poll(Duration.ZERO);
+    }
+
     assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeRebalance)
 
     // verify commits are intercepted on close

From d378ad39fa44544d07459806bf3f7a9ab3258b3f Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Wed, 14 Feb 2024 23:25:35 -0800
Subject: [PATCH 038/258] MINOR: Fix KafkaAdminClientTest.testClientInstanceId
 (#15370)

This patch tries to address the [flakiness](https://ge.apache.org/scans/tests?search.rootProjectNames=kafka&search.timeZoneId=Europe%2FZurich&tests.container=org.apache.kafka.clients.admin.KafkaAdminClientTest&tests.sortField=FLAKY&tests.test=testClientInstanceId()) of `KafkaAdminClientTest.testClientInstanceId`. The test fails with `org.apache.kafka.common.errors.TimeoutException: Timed out waiting for a node assignment.`. I believe that it is because 10ms is not enough when our CI is busy. Let's try with 1s.

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 .../org/apache/kafka/clients/admin/KafkaAdminClientTest.java    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
index 018820273a2f0..b8b3d54ef439a 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
@@ -7258,7 +7258,7 @@ public void testClientInstanceId() {
                 request -> request instanceof GetTelemetrySubscriptionsRequest,
                 new GetTelemetrySubscriptionsResponse(responseData));
 
-            Uuid result = env.adminClient().clientInstanceId(Duration.ofMillis(10));
+            Uuid result = env.adminClient().clientInstanceId(Duration.ofSeconds(1));
             assertEquals(expected, result);
         }
     }

From 5edf52359a99db801db91927c282af06b5ca4ee0 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Thu, 15 Feb 2024 00:19:43 -0800
Subject: [PATCH 039/258] MINOR: Fix group metadata loading log (#15368)

Spotted the following log:
```
[2024-02-14 09:59:30,103] INFO [GroupCoordinator id=1] Finished loading of metadata from 39 in __consumer_offsets-4ms with epoch 2 where 39ms was spent in the scheduler. Loaded 0 records which total to 0 bytes. (org.apache.kafka.coordinator.group.runtime.CoordinatorRuntime)
```
The partition and the time are incorrect. This patch fixes it.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Bruno Cadonna <bruno@confluent.io>
---
 .../kafka/coordinator/group/runtime/CoordinatorRuntime.java   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
index 5674273b12dab..b0be84c7a900b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
@@ -1634,9 +1634,9 @@ public void scheduleLoadOperation(
                                             ctx.transitionTo(CoordinatorState.ACTIVE);
                                             if (summary != null) {
                                                 runtimeMetrics.recordPartitionLoadSensor(summary.startTimeMs(), summary.endTimeMs());
-                                                log.info("Finished loading of metadata from {} in {}ms with epoch {} where {}ms " +
+                                                log.info("Finished loading of metadata from {} with epoch {} in {}ms where {}ms " +
                                                     "was spent in the scheduler. Loaded {} records which total to {} bytes.",
-                                                    summary.endTimeMs() - summary.startTimeMs(), tp, partitionEpoch,
+                                                    tp, partitionEpoch, summary.endTimeMs() - summary.startTimeMs(),
                                                     summary.schedulerQueueTimeMs(), summary.numRecords(), summary.numBytes()
                                                 );
                                             }

From ff90f78c700c582f9800013faad827c36b45ceb7 Mon Sep 17 00:00:00 2001
From: Mayank Shekhar Narula <42991652+msn-tldr@users.noreply.github.com>
Date: Thu, 15 Feb 2024 17:26:47 +0000
Subject: [PATCH 040/258] KAFKA-16226; Reduce synchronization between producer
 threads (#15323)

As this [JIRA](https://issues.apache.org/jira/browse/KAFKA-16226) explains, there is increased synchronization between application-thread, and the background thread as the background thread started to synchronized methods Metadata.currentLeader() in [original PR](https://github.com/apache/kafka/pull/14384). So this PR does the following changes
1. Changes background thread, i.e. RecordAccumulator's partitionReady(), and drainBatchesForOneNode(), to not use `Metadata.currentLeader()`. Instead rely on `MetadataCache` that is immutable. So access to it is unsynchronized.
2.  This PR repurposes `MetadataCache` as an immutable snapshot of Metadata. This is a wrapper around public `Cluster`. `MetadataCache`'s API/functionality should be extended for internal client usage Vs public `Cluster`. For example, this PR adds `MetadataCache.leaderEpochFor()`
3. Rename `MetadataCache` to `MetadataSnapshot` to make it explicit its immutable.

**Note both `Cluster` and `MetadataCache` are not syncronized, hence reduce synchronization from the hot path for high partition counts.**

Reviewers: Jason Gustafson <jason@confluent.io>
---
 .../org/apache/kafka/clients/Metadata.java    |  65 +--
 ...tadataCache.java => MetadataSnapshot.java} |  77 +--
 .../producer/internals/ProducerBatch.java     |  13 +-
 .../producer/internals/RecordAccumulator.java |  78 +--
 .../clients/producer/internals/Sender.java    |   6 +-
 .../common/requests/MetadataResponse.java     |   6 +-
 ...cheTest.java => MetadataSnapshotTest.java} |  56 ++-
 .../apache/kafka/clients/MetadataTest.java    |   4 +-
 .../producer/internals/ProducerBatchTest.java |  36 +-
 .../internals/RecordAccumulatorTest.java      | 468 ++++++++----------
 .../producer/internals/SenderTest.java        |   9 +-
 .../internals/TransactionManagerTest.java     |  33 +-
 .../java/org/apache/kafka/test/TestUtils.java |  39 ++
 13 files changed, 476 insertions(+), 414 deletions(-)
 rename clients/src/main/java/org/apache/kafka/clients/{MetadataCache.java => MetadataSnapshot.java} (77%)
 rename clients/src/test/java/org/apache/kafka/clients/{MetadataCacheTest.java => MetadataSnapshotTest.java} (77%)

diff --git a/clients/src/main/java/org/apache/kafka/clients/Metadata.java b/clients/src/main/java/org/apache/kafka/clients/Metadata.java
index 607c74eeddb00..30cad44a4bc97 100644
--- a/clients/src/main/java/org/apache/kafka/clients/Metadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/Metadata.java
@@ -75,7 +75,7 @@ public class Metadata implements Closeable {
     private KafkaException fatalException;
     private Set<String> invalidTopics;
     private Set<String> unauthorizedTopics;
-    private MetadataCache cache = MetadataCache.empty();
+    private volatile MetadataSnapshot metadataSnapshot = MetadataSnapshot.empty();
     private boolean needFullUpdate;
     private boolean needPartialUpdate;
     private long equivalentResponseCount;
@@ -123,8 +123,15 @@ public Metadata(long refreshBackoffMs,
     /**
      * Get the current cluster info without blocking
      */
-    public synchronized Cluster fetch() {
-        return cache.cluster();
+    public Cluster fetch() {
+        return metadataSnapshot.cluster();
+    }
+
+    /**
+     * Get the current metadata cache.
+     */
+    public MetadataSnapshot fetchMetadataSnapshot() {
+        return metadataSnapshot;
     }
 
     /**
@@ -265,7 +272,7 @@ public synchronized void addClusterUpdateListener(ClusterResourceListener listen
      */
     synchronized Optional<MetadataResponse.PartitionMetadata> partitionMetadataIfCurrent(TopicPartition topicPartition) {
         Integer epoch = lastSeenLeaderEpochs.get(topicPartition);
-        Optional<MetadataResponse.PartitionMetadata> partitionMetadata = cache.partitionMetadata(topicPartition);
+        Optional<MetadataResponse.PartitionMetadata> partitionMetadata = metadataSnapshot.partitionMetadata(topicPartition);
         if (epoch == null) {
             // old cluster format (no epochs)
             return partitionMetadata;
@@ -278,8 +285,8 @@ synchronized Optional<MetadataResponse.PartitionMetadata> partitionMetadataIfCur
     /**
      * @return a mapping from topic names to topic IDs for all topics with valid IDs in the cache
      */
-    public synchronized Map<String, Uuid> topicIds() {
-        return cache.topicIds();
+    public Map<String, Uuid> topicIds() {
+        return metadataSnapshot.topicIds();
     }
 
     public synchronized LeaderAndEpoch currentLeader(TopicPartition topicPartition) {
@@ -289,14 +296,14 @@ public synchronized LeaderAndEpoch currentLeader(TopicPartition topicPartition)
 
         MetadataResponse.PartitionMetadata partitionMetadata = maybeMetadata.get();
         Optional<Integer> leaderEpochOpt = partitionMetadata.leaderEpoch;
-        Optional<Node> leaderNodeOpt = partitionMetadata.leaderId.flatMap(cache::nodeById);
+        Optional<Node> leaderNodeOpt = partitionMetadata.leaderId.flatMap(metadataSnapshot::nodeById);
         return new LeaderAndEpoch(leaderNodeOpt, leaderEpochOpt);
     }
 
     public synchronized void bootstrap(List<InetSocketAddress> addresses) {
         this.needFullUpdate = true;
         this.updateVersion += 1;
-        this.cache = MetadataCache.bootstrap(addresses);
+        this.metadataSnapshot = MetadataSnapshot.bootstrap(addresses);
     }
 
     /**
@@ -335,22 +342,22 @@ public synchronized void update(int requestVersion, MetadataResponse response, b
         // this count is reset to 0 in updateLatestMetadata()
         this.equivalentResponseCount++;
 
-        String previousClusterId = cache.clusterResource().clusterId();
+        String previousClusterId = metadataSnapshot.clusterResource().clusterId();
 
-        this.cache = handleMetadataResponse(response, isPartialUpdate, nowMs);
+        this.metadataSnapshot = handleMetadataResponse(response, isPartialUpdate, nowMs);
 
-        Cluster cluster = cache.cluster();
+        Cluster cluster = metadataSnapshot.cluster();
         maybeSetMetadataError(cluster);
 
         this.lastSeenLeaderEpochs.keySet().removeIf(tp -> !retainTopic(tp.topic(), false, nowMs));
 
-        String newClusterId = cache.clusterResource().clusterId();
+        String newClusterId = metadataSnapshot.clusterResource().clusterId();
         if (!Objects.equals(previousClusterId, newClusterId)) {
             log.info("Cluster ID: {}", newClusterId);
         }
-        clusterResourceListeners.onUpdate(cache.clusterResource());
+        clusterResourceListeners.onUpdate(metadataSnapshot.clusterResource());
 
-        log.debug("Updated cluster metadata updateVersion {} to {}", this.updateVersion, this.cache);
+        log.debug("Updated cluster metadata updateVersion {} to {}", this.updateVersion, this.metadataSnapshot);
     }
 
     /**
@@ -365,7 +372,7 @@ public synchronized void update(int requestVersion, MetadataResponse response, b
     public synchronized Set<TopicPartition> updatePartitionLeadership(Map<TopicPartition, LeaderIdAndEpoch> partitionLeaders, List<Node> leaderNodes) {
         Map<Integer, Node> newNodes = leaderNodes.stream().collect(Collectors.toMap(Node::id, node -> node));
         // Insert non-overlapping nodes from existing-nodes into new-nodes.
-        this.cache.cluster().nodes().stream().forEach(node -> newNodes.putIfAbsent(node.id(), node));
+        this.metadataSnapshot.cluster().nodes().stream().forEach(node -> newNodes.putIfAbsent(node.id(), node));
 
         // Create partition-metadata for all updated partitions. Exclude updates for partitions -
         // 1. for which the corresponding partition has newer leader in existing metadata.
@@ -388,12 +395,12 @@ public synchronized Set<TopicPartition> updatePartitionLeadership(Map<TopicParti
                 log.debug("For {}, incoming leader({}), the corresponding node information for node-id {} is missing, so ignoring.", partition, newLeader, newLeader.leaderId.get());
                 continue;
             }
-            if (!this.cache.partitionMetadata(partition).isPresent()) {
+            if (!this.metadataSnapshot.partitionMetadata(partition).isPresent()) {
                 log.debug("For {}, incoming leader({}), partition metadata is no longer cached, ignoring.", partition, newLeader);
                 continue;
             }
 
-            MetadataResponse.PartitionMetadata existingMetadata = this.cache.partitionMetadata(partition).get();
+            MetadataResponse.PartitionMetadata existingMetadata = this.metadataSnapshot.partitionMetadata(partition).get();
             MetadataResponse.PartitionMetadata updatedMetadata = new MetadataResponse.PartitionMetadata(
                 existingMetadata.error,
                 partition,
@@ -416,7 +423,7 @@ public synchronized Set<TopicPartition> updatePartitionLeadership(Map<TopicParti
         Set<String> updatedTopics = updatePartitionMetadata.stream().map(MetadataResponse.PartitionMetadata::topic).collect(Collectors.toSet());
 
         // Get topic-ids for updated topics from existing topic-ids.
-        Map<String, Uuid> existingTopicIds = this.cache.topicIds();
+        Map<String, Uuid> existingTopicIds = this.metadataSnapshot.topicIds();
         Map<String, Uuid> topicIdsForUpdatedTopics = updatedTopics.stream()
             .filter(e -> existingTopicIds.containsKey(e))
             .collect(Collectors.toMap(e -> e, e -> existingTopicIds.get(e)));
@@ -429,15 +436,15 @@ public synchronized Set<TopicPartition> updatePartitionLeadership(Map<TopicParti
 
         // Fetch responses can include partition level leader changes, when this happens, we perform a partial
         // metadata update, by keeping the unchanged partition and update the changed partitions.
-        this.cache = cache.mergeWith(
-            cache.clusterResource().clusterId(),
+        this.metadataSnapshot = metadataSnapshot.mergeWith(
+            metadataSnapshot.clusterResource().clusterId(),
             newNodes,
             updatePartitionMetadata,
             Collections.emptySet(), Collections.emptySet(), Collections.emptySet(),
-            cache.cluster().controller(),
+            metadataSnapshot.cluster().controller(),
             topicIdsForUpdatedTopics,
             (topic, isInternal) -> true);
-        clusterResourceListeners.onUpdate(cache.clusterResource());
+        clusterResourceListeners.onUpdate(metadataSnapshot.clusterResource());
 
         return updatePartitionMetadata.stream()
             .map(metadata -> metadata.topicPartition)
@@ -467,7 +474,7 @@ private void checkUnauthorizedTopics(Cluster cluster) {
     /**
      * Transform a MetadataResponse into a new MetadataCache instance.
      */
-    private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, boolean isPartialUpdate, long nowMs) {
+    private MetadataSnapshot handleMetadataResponse(MetadataResponse metadataResponse, boolean isPartialUpdate, long nowMs) {
         // All encountered topics.
         Set<String> topics = new HashSet<>();
 
@@ -478,7 +485,7 @@ private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse,
 
         List<MetadataResponse.PartitionMetadata> partitions = new ArrayList<>();
         Map<String, Uuid> topicIds = new HashMap<>();
-        Map<String, Uuid> oldTopicIds = cache.topicIds();
+        Map<String, Uuid> oldTopicIds = metadataSnapshot.topicIds();
         for (MetadataResponse.TopicMetadata metadata : metadataResponse.topicMetadata()) {
             String topicName = metadata.topic();
             Uuid topicId = metadata.topicId();
@@ -526,11 +533,11 @@ else if (metadata.error() == Errors.TOPIC_AUTHORIZATION_FAILED)
 
         Map<Integer, Node> nodes = metadataResponse.brokersById();
         if (isPartialUpdate)
-            return this.cache.mergeWith(metadataResponse.clusterId(), nodes, partitions,
+            return this.metadataSnapshot.mergeWith(metadataResponse.clusterId(), nodes, partitions,
                 unauthorizedTopics, invalidTopics, internalTopics, metadataResponse.controller(), topicIds,
                 (topic, isInternal) -> !topics.contains(topic) && retainTopic(topic, isInternal, nowMs));
         else
-            return new MetadataCache(metadataResponse.clusterId(), nodes, partitions,
+            return new MetadataSnapshot(metadataResponse.clusterId(), nodes, partitions,
                 unauthorizedTopics, invalidTopics, internalTopics, metadataResponse.controller(), topicIds);
     }
 
@@ -575,7 +582,7 @@ private Optional<MetadataResponse.PartitionMetadata> updateLatestMetadata(
             } else {
                 // Otherwise ignore the new metadata and use the previously cached info
                 log.debug("Got metadata for an older epoch {} (current is {}) for partition {}, not updating", newEpoch, currentEpoch, tp);
-                return cache.partitionMetadata(tp);
+                return metadataSnapshot.partitionMetadata(tp);
             }
         } else {
             // Handle old cluster formats as well as error responses where leader and epoch are missing
@@ -738,8 +745,8 @@ protected MetadataRequest.Builder newMetadataRequestBuilderForNewTopics() {
     /**
      * @return Mapping from topic IDs to topic names for all topics in the cache.
      */
-    public synchronized Map<Uuid, String> topicNames() {
-        return cache.topicNames();
+    public Map<Uuid, String> topicNames() {
+        return metadataSnapshot.topicNames();
     }
 
     protected boolean retainTopic(String topic, boolean isInternal, long nowMs) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java b/clients/src/main/java/org/apache/kafka/clients/MetadataSnapshot.java
similarity index 77%
rename from clients/src/main/java/org/apache/kafka/clients/MetadataCache.java
rename to clients/src/main/java/org/apache/kafka/clients/MetadataSnapshot.java
index 45574c3549c56..fbfb828b31c15 100644
--- a/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java
+++ b/clients/src/main/java/org/apache/kafka/clients/MetadataSnapshot.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients;
 
+import java.util.OptionalInt;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.ClusterResource;
 import org.apache.kafka.common.Node;
@@ -39,10 +40,11 @@
 import java.util.stream.Collectors;
 
 /**
- * An internal mutable cache of nodes, topics, and partitions in the Kafka cluster. This keeps an up-to-date Cluster
+ * An internal immutable snapshot of nodes, topics, and partitions in the Kafka cluster. This keeps an up-to-date Cluster
  * instance which is optimized for read access.
+ * Prefer to extend MetadataSnapshot's API for internal client usage Vs the public {@link Cluster}
  */
-public class MetadataCache {
+public class MetadataSnapshot {
     private final String clusterId;
     private final Map<Integer, Node> nodes;
     private final Set<String> unauthorizedTopics;
@@ -54,7 +56,7 @@ public class MetadataCache {
     private final Map<Uuid, String> topicNames;
     private Cluster clusterInstance;
 
-    MetadataCache(String clusterId,
+    public MetadataSnapshot(String clusterId,
                   Map<Integer, Node> nodes,
                   Collection<PartitionMetadata> partitions,
                   Set<String> unauthorizedTopics,
@@ -65,30 +67,32 @@ public class MetadataCache {
         this(clusterId, nodes, partitions, unauthorizedTopics, invalidTopics, internalTopics, controller, topicIds, null);
     }
 
-    private MetadataCache(String clusterId,
-                          Map<Integer, Node> nodes,
-                          Collection<PartitionMetadata> partitions,
-                          Set<String> unauthorizedTopics,
-                          Set<String> invalidTopics,
-                          Set<String> internalTopics,
-                          Node controller,
-                          Map<String, Uuid> topicIds,
-                          Cluster clusterInstance) {
+    // Visible for testing
+    public MetadataSnapshot(String clusterId,
+        Map<Integer, Node> nodes,
+        Collection<PartitionMetadata> partitions,
+        Set<String> unauthorizedTopics,
+        Set<String> invalidTopics,
+        Set<String> internalTopics,
+        Node controller,
+        Map<String, Uuid> topicIds,
+        Cluster clusterInstance) {
         this.clusterId = clusterId;
-        this.nodes = nodes;
-        this.unauthorizedTopics = unauthorizedTopics;
-        this.invalidTopics = invalidTopics;
-        this.internalTopics = internalTopics;
+        this.nodes = Collections.unmodifiableMap(nodes);
+        this.unauthorizedTopics = Collections.unmodifiableSet(unauthorizedTopics);
+        this.invalidTopics = Collections.unmodifiableSet(invalidTopics);
+        this.internalTopics = Collections.unmodifiableSet(internalTopics);
         this.controller = controller;
         this.topicIds = Collections.unmodifiableMap(topicIds);
         this.topicNames = Collections.unmodifiableMap(
             topicIds.entrySet().stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey))
         );
 
-        this.metadataByPartition = new HashMap<>(partitions.size());
+        Map<TopicPartition, PartitionMetadata> tmpMetadataByPartition = new HashMap<>(partitions.size());
         for (PartitionMetadata p : partitions) {
-            this.metadataByPartition.put(p.topicPartition, p);
+            tmpMetadataByPartition.put(p.topicPartition, p);
         }
+        this.metadataByPartition = Collections.unmodifiableMap(tmpMetadataByPartition);
 
         if (clusterInstance == null) {
             computeClusterView();
@@ -113,7 +117,7 @@ Optional<Node> nodeById(int id) {
         return Optional.ofNullable(nodes.get(id));
     }
 
-    Cluster cluster() {
+    public Cluster cluster() {
         if (clusterInstance == null) {
             throw new IllegalStateException("Cached Cluster instance should not be null, but was.");
         } else {
@@ -121,13 +125,28 @@ Cluster cluster() {
         }
     }
 
+    /**
+     * Get leader-epoch for partition.
+     *
+     * @param tp partition
+     * @return leader-epoch if known, else return OptionalInt.empty()
+     */
+    public OptionalInt leaderEpochFor(TopicPartition tp) {
+        PartitionMetadata partitionMetadata = metadataByPartition.get(tp);
+        if (partitionMetadata == null || !partitionMetadata.leaderEpoch.isPresent()) {
+            return OptionalInt.empty();
+        } else {
+            return OptionalInt.of(partitionMetadata.leaderEpoch.get());
+        }
+    }
+
     ClusterResource clusterResource() {
         return new ClusterResource(clusterId);
     }
 
     /**
-     * Merges the metadata cache's contents with the provided metadata, returning a new metadata cache. The provided
-     * metadata is presumed to be more recent than the cache's metadata, and therefore all overlapping metadata will
+     * Merges the metadata snapshot's contents with the provided metadata, returning a new metadata snapshot. The provided
+     * metadata is presumed to be more recent than the snapshot's metadata, and therefore all overlapping metadata will
      * be overridden.
      *
      * @param newClusterId the new cluster Id
@@ -138,9 +157,9 @@ ClusterResource clusterResource() {
      * @param newController the new controller node
      * @param addTopicIds the mapping from topic name to topic ID, for topics in addPartitions
      * @param retainTopic returns whether a pre-existing topic's metadata should be retained
-     * @return the merged metadata cache
+     * @return the merged metadata snapshot
      */
-    MetadataCache mergeWith(String newClusterId,
+    MetadataSnapshot mergeWith(String newClusterId,
                             Map<Integer, Node> newNodes,
                             Collection<PartitionMetadata> addPartitions,
                             Set<String> addUnauthorizedTopics,
@@ -180,7 +199,7 @@ MetadataCache mergeWith(String newClusterId,
         Set<String> newInvalidTopics = fillSet(addInvalidTopics, invalidTopics, shouldRetainTopic);
         Set<String> newInternalTopics = fillSet(addInternalTopics, internalTopics, shouldRetainTopic);
 
-        return new MetadataCache(newClusterId, newNodes, newMetadataByPartition.values(), newUnauthorizedTopics,
+        return new MetadataSnapshot(newClusterId, newNodes, newMetadataByPartition.values(), newUnauthorizedTopics,
                 newInvalidTopics, newInternalTopics, newController, newTopicIds);
     }
 
@@ -212,26 +231,26 @@ private void computeClusterView() {
                 invalidTopics, internalTopics, controller, topicIds);
     }
 
-    static MetadataCache bootstrap(List<InetSocketAddress> addresses) {
+    static MetadataSnapshot bootstrap(List<InetSocketAddress> addresses) {
         Map<Integer, Node> nodes = new HashMap<>();
         int nodeId = -1;
         for (InetSocketAddress address : addresses) {
             nodes.put(nodeId, new Node(nodeId, address.getHostString(), address.getPort()));
             nodeId--;
         }
-        return new MetadataCache(null, nodes, Collections.emptyList(),
+        return new MetadataSnapshot(null, nodes, Collections.emptyList(),
                 Collections.emptySet(), Collections.emptySet(), Collections.emptySet(),
                 null, Collections.emptyMap(), Cluster.bootstrap(addresses));
     }
 
-    static MetadataCache empty() {
-        return new MetadataCache(null, Collections.emptyMap(), Collections.emptyList(),
+    static MetadataSnapshot empty() {
+        return new MetadataSnapshot(null, Collections.emptyMap(), Collections.emptyList(),
                 Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap(), Cluster.empty());
     }
 
     @Override
     public String toString() {
-        return "MetadataCache{" +
+        return "MetadataSnapshot{" +
                 "clusterId='" + clusterId + '\'' +
                 ", nodes=" + nodes +
                 ", partitions=" + metadataByPartition.values() +
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
index 61432b53ab07c..391cc1b3441b3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
-import java.util.Optional;
+import java.util.OptionalInt;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
@@ -81,7 +81,7 @@ private enum FinalState { ABORTED, FAILED, SUCCEEDED }
     private boolean reopened;
 
     // Tracks the current-leader's epoch to which this batch would be sent, in the current to produce the batch.
-    private Optional<Integer> currentLeaderEpoch;
+    private OptionalInt currentLeaderEpoch;
     // Tracks the attempt in which leader was changed to currentLeaderEpoch for the 1st time.
     private int attemptsWhenLeaderLastChanged;
 
@@ -100,7 +100,7 @@ public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, lon
         this.isSplitBatch = isSplitBatch;
         float compressionRatioEstimation = CompressionRatioEstimator.estimation(topicPartition.topic(),
                                                                                 recordsBuilder.compressionType());
-        this.currentLeaderEpoch = Optional.empty();
+        this.currentLeaderEpoch = OptionalInt.empty();
         this.attemptsWhenLeaderLastChanged = 0;
         recordsBuilder.setEstimatedCompressionRatio(compressionRatioEstimation);
     }
@@ -109,8 +109,9 @@ public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, lon
      * It will update the leader to which this batch will be produced for the ongoing attempt, if a newer leader is known.
      * @param latestLeaderEpoch latest leader's epoch.
      */
-    void maybeUpdateLeaderEpoch(Optional<Integer> latestLeaderEpoch) {
-        if (!currentLeaderEpoch.equals(latestLeaderEpoch)) {
+    void maybeUpdateLeaderEpoch(OptionalInt latestLeaderEpoch) {
+        if (latestLeaderEpoch.isPresent()
+            && (!currentLeaderEpoch.isPresent() || currentLeaderEpoch.getAsInt() < latestLeaderEpoch.getAsInt())) {
             log.trace("For {}, leader will be updated, currentLeaderEpoch: {}, attemptsWhenLeaderLastChanged:{}, latestLeaderEpoch: {}, current attempt: {}",
                 this, currentLeaderEpoch, attemptsWhenLeaderLastChanged, latestLeaderEpoch, attempts);
             attemptsWhenLeaderLastChanged = attempts();
@@ -558,7 +559,7 @@ public boolean sequenceHasBeenReset() {
     }
 
     // VisibleForTesting
-    Optional<Integer> currentLeaderEpoch() {
+    OptionalInt currentLeaderEpoch() {
         return currentLeaderEpoch;
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java
index 5865752c8fef6..013ad32dc7df3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java
@@ -26,13 +26,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.OptionalInt;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.utils.ExponentialBackoff;
@@ -647,27 +648,27 @@ private long batchReady(boolean exhausted, TopicPartition part, Node leader,
     }
 
     /**
-     * Iterate over partitions to see which one have batches ready and collect leaders of those partitions
-     * into the set of ready nodes.  If partition has no leader, add the topic to the set of topics with
-     * no leader.  This function also calculates stats for adaptive partitioning.
+     * Iterate over partitions to see which one have batches ready and collect leaders of those
+     * partitions into the set of ready nodes.  If partition has no leader, add the topic to the set
+     * of topics with no leader.  This function also calculates stats for adaptive partitioning.
      *
-     * @param metadata The cluster metadata
-     * @param nowMs The current time
-     * @param topic The topic
-     * @param topicInfo The topic info
+     * @param metadataSnapshot      The cluster metadata
+     * @param nowMs                 The current time
+     * @param topic                 The topic
+     * @param topicInfo             The topic info
      * @param nextReadyCheckDelayMs The delay for next check
-     * @param readyNodes The set of ready nodes (to be filled in)
-     * @param unknownLeaderTopics The set of topics with no leader (to be filled in)
+     * @param readyNodes            The set of ready nodes (to be filled in)
+     * @param unknownLeaderTopics   The set of topics with no leader (to be filled in)
      * @return The delay for next check
      */
-    private long partitionReady(Metadata metadata, long nowMs, String topic,
+    private long partitionReady(MetadataSnapshot metadataSnapshot, long nowMs, String topic,
                                 TopicInfo topicInfo,
                                 long nextReadyCheckDelayMs, Set<Node> readyNodes, Set<String> unknownLeaderTopics) {
         ConcurrentMap<Integer, Deque<ProducerBatch>> batches = topicInfo.batches;
         // Collect the queue sizes for available partitions to be used in adaptive partitioning.
         int[] queueSizes = null;
         int[] partitionIds = null;
-        if (enableAdaptivePartitioning && batches.size() >= metadata.fetch().partitionsForTopic(topic).size()) {
+        if (enableAdaptivePartitioning && batches.size() >= metadataSnapshot.cluster().partitionsForTopic(topic).size()) {
             // We don't do adaptive partitioning until we scheduled at least a batch for all
             // partitions (i.e. we have the corresponding entries in the batches map), we just
             // do uniform.  The reason is that we build queue sizes from the batches map,
@@ -684,8 +685,7 @@ private long partitionReady(Metadata metadata, long nowMs, String topic,
             // Advance queueSizesIndex so that we properly index available
             // partitions.  Do it here so that it's done for all code paths.
 
-            Metadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(part);
-            Node leader = leaderAndEpoch.leader.orElse(null);
+            Node leader = metadataSnapshot.cluster().leaderFor(part);
             if (leader != null && queueSizes != null) {
                 ++queueSizesIndex;
                 assert queueSizesIndex < queueSizes.length;
@@ -700,9 +700,14 @@ private long partitionReady(Metadata metadata, long nowMs, String topic,
             final int dequeSize;
             final boolean full;
 
-            // This loop is especially hot with large partition counts.
+            OptionalInt leaderEpoch = metadataSnapshot.leaderEpochFor(part);
+
+            // This loop is especially hot with large partition counts. So -
 
-            // We are careful to only perform the minimum required inside the
+            // 1. We should avoid code that increases synchronization between application thread calling
+            // send(), and background thread running runOnce(), see https://issues.apache.org/jira/browse/KAFKA-16226
+
+            // 2. We are careful to only perform the minimum required inside the
             // synchronized block, as this lock is also used to synchronize producer threads
             // attempting to append() to a partition/batch.
 
@@ -715,7 +720,7 @@ private long partitionReady(Metadata metadata, long nowMs, String topic,
                 }
 
                 waitedTimeMs = batch.waitedTimeMs(nowMs);
-                batch.maybeUpdateLeaderEpoch(leaderAndEpoch.epoch);
+                batch.maybeUpdateLeaderEpoch(leaderEpoch);
                 backingOff = shouldBackoff(batch.hasLeaderChangedForTheOngoingRetry(), batch, waitedTimeMs);
                 backoffAttempts = batch.attempts();
                 dequeSize = deque.size();
@@ -776,7 +781,7 @@ private long partitionReady(Metadata metadata, long nowMs, String topic,
      * </ul>
      * </ol>
      */
-    public ReadyCheckResult ready(Metadata metadata, long nowMs) {
+    public ReadyCheckResult ready(MetadataSnapshot metadataSnapshot, long nowMs) {
         Set<Node> readyNodes = new HashSet<>();
         long nextReadyCheckDelayMs = Long.MAX_VALUE;
         Set<String> unknownLeaderTopics = new HashSet<>();
@@ -784,7 +789,7 @@ public ReadyCheckResult ready(Metadata metadata, long nowMs) {
         // cumulative frequency table (used in partitioner).
         for (Map.Entry<String, TopicInfo> topicInfoEntry : this.topicInfoMap.entrySet()) {
             final String topic = topicInfoEntry.getKey();
-            nextReadyCheckDelayMs = partitionReady(metadata, nowMs, topic, topicInfoEntry.getValue(), nextReadyCheckDelayMs, readyNodes, unknownLeaderTopics);
+            nextReadyCheckDelayMs = partitionReady(metadataSnapshot, nowMs, topic, topicInfoEntry.getValue(), nextReadyCheckDelayMs, readyNodes, unknownLeaderTopics);
         }
         return new ReadyCheckResult(readyNodes, nextReadyCheckDelayMs, unknownLeaderTopics);
     }
@@ -861,9 +866,9 @@ private boolean shouldStopDrainBatchesForPartition(ProducerBatch first, TopicPar
         return false;
     }
 
-    private List<ProducerBatch> drainBatchesForOneNode(Metadata metadata, Node node, int maxSize, long now) {
+    private List<ProducerBatch> drainBatchesForOneNode(MetadataSnapshot metadataSnapshot, Node node, int maxSize, long now) {
         int size = 0;
-        List<PartitionInfo> parts = metadata.fetch().partitionsForNode(node.id());
+        List<PartitionInfo> parts = metadataSnapshot.cluster().partitionsForNode(node.id());
         List<ProducerBatch> ready = new ArrayList<>();
         if (parts.isEmpty())
             return ready;
@@ -879,17 +884,12 @@ private List<ProducerBatch> drainBatchesForOneNode(Metadata metadata, Node node,
             // Only proceed if the partition has no in-flight batches.
             if (isMuted(tp))
                 continue;
-            Metadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(tp);
-            // Although a small chance, but skip this partition if leader has changed since the partition -> node assignment obtained from outside the loop.
-            // In this case, skip sending it to the old leader, as it would return aa NO_LEADER_OR_FOLLOWER error.
-            if (!leaderAndEpoch.leader.isPresent())
-                continue;
-            if (!node.equals(leaderAndEpoch.leader.get()))
-                continue;
             Deque<ProducerBatch> deque = getDeque(tp);
             if (deque == null)
                 continue;
 
+            OptionalInt leaderEpoch = metadataSnapshot.leaderEpochFor(tp);
+
             final ProducerBatch batch;
             synchronized (deque) {
                 // invariant: !isMuted(tp,now) && deque != null
@@ -899,7 +899,7 @@ private List<ProducerBatch> drainBatchesForOneNode(Metadata metadata, Node node,
 
                 // first != null
                 // Only drain the batch if it is not during backoff period.
-                first.maybeUpdateLeaderEpoch(leaderAndEpoch.epoch);
+                first.maybeUpdateLeaderEpoch(leaderEpoch);
                 if (shouldBackoff(first.hasLeaderChangedForTheOngoingRetry(), first, first.waitedTimeMs(now)))
                     continue;
 
@@ -962,22 +962,24 @@ private void updateDrainIndex(String idString, int drainIndex) {
     }
 
     /**
-     * Drain all the data for the given nodes and collate them into a list of batches that will fit within the specified
-     * size on a per-node basis. This method attempts to avoid choosing the same topic-node over and over.
+     * Drain all the data for the given nodes and collate them into a list of batches that will fit
+     * within the specified size on a per-node basis. This method attempts to avoid choosing the same
+     * topic-node over and over.
      *
-     * @param metadata The current cluster metadata
-     * @param nodes The list of node to drain
-     * @param maxSize The maximum number of bytes to drain
-     * @param now The current unix time in milliseconds
-     * @return A list of {@link ProducerBatch} for each node specified with total size less than the requested maxSize.
+     * @param metadataSnapshot  The current cluster metadata
+     * @param nodes             The list of node to drain
+     * @param maxSize           The maximum number of bytes to drain
+     * @param now               The current unix time in milliseconds
+     * @return A list of {@link ProducerBatch} for each node specified with total size less than the
+     * requested maxSize.
      */
-    public Map<Integer, List<ProducerBatch>> drain(Metadata metadata, Set<Node> nodes, int maxSize, long now) {
+    public Map<Integer, List<ProducerBatch>> drain(MetadataSnapshot metadataSnapshot, Set<Node> nodes, int maxSize, long now) {
         if (nodes.isEmpty())
             return Collections.emptyMap();
 
         Map<Integer, List<ProducerBatch>> batches = new HashMap<>();
         for (Node node : nodes) {
-            List<ProducerBatch> ready = drainBatchesForOneNode(metadata, node, maxSize, now);
+            List<ProducerBatch> ready = drainBatchesForOneNode(metadataSnapshot, node, maxSize, now);
             batches.put(node.id(), ready);
         }
         return batches;
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
index 01f49ffad0903..2cdc3b876d0b1 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
@@ -25,6 +25,7 @@
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.KafkaClient;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.NetworkClientUtils;
 import org.apache.kafka.clients.RequestCompletionHandler;
 import org.apache.kafka.common.InvalidRecordException;
@@ -361,8 +362,9 @@ private boolean shouldHandleAuthorizationError(RuntimeException exception) {
     }
 
     private long sendProducerData(long now) {
+        MetadataSnapshot metadataSnapshot = metadata.fetchMetadataSnapshot();
         // get the list of partitions with data ready to send
-        RecordAccumulator.ReadyCheckResult result = this.accumulator.ready(metadata, now);
+        RecordAccumulator.ReadyCheckResult result = this.accumulator.ready(metadataSnapshot, now);
 
         // if there are any partitions whose leaders are not known yet, force metadata update
         if (!result.unknownLeaderTopics.isEmpty()) {
@@ -397,7 +399,7 @@ private long sendProducerData(long now) {
         }
 
         // create produce requests
-        Map<Integer, List<ProducerBatch>> batches = this.accumulator.drain(metadata, result.readyNodes, this.maxRequestSize, now);
+        Map<Integer, List<ProducerBatch>> batches = this.accumulator.drain(metadataSnapshot, result.readyNodes, this.maxRequestSize, now);
         addToInflightBatches(batches);
         if (guaranteeMessageOrder) {
             // Mute all the partitions drained
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java
index 47cdd3f0d7e90..bc6c387c5e021 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java
@@ -172,9 +172,9 @@ public static PartitionInfo toPartitionInfo(PartitionMetadata metadata, Map<Inte
         return new PartitionInfo(metadata.topic(),
                 metadata.partition(),
                 metadata.leaderId.map(nodesById::get).orElse(null),
-                convertToNodeArray(metadata.replicaIds, nodesById),
-                convertToNodeArray(metadata.inSyncReplicaIds, nodesById),
-                convertToNodeArray(metadata.offlineReplicaIds, nodesById));
+                (metadata.replicaIds == null) ? null : convertToNodeArray(metadata.replicaIds, nodesById),
+                (metadata.inSyncReplicaIds == null) ? null : convertToNodeArray(metadata.inSyncReplicaIds, nodesById),
+                (metadata.offlineReplicaIds == null) ? null : convertToNodeArray(metadata.offlineReplicaIds, nodesById));
     }
 
     private static Node[] convertToNodeArray(List<Integer> replicaIds, Map<Integer, Node> nodesById) {
diff --git a/clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.java b/clients/src/test/java/org/apache/kafka/clients/MetadataSnapshotTest.java
similarity index 77%
rename from clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.java
rename to clients/src/test/java/org/apache/kafka/clients/MetadataSnapshotTest.java
index f99eb04dfcb97..1012709926a72 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataSnapshotTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients;
 
+import java.util.OptionalInt;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
@@ -37,7 +38,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-public class MetadataCacheTest {
+public class MetadataSnapshotTest {
 
     @Test
     public void testMissingLeaderEndpoint() {
@@ -62,7 +63,7 @@ public void testMissingLeaderEndpoint() {
         nodesById.put(7, new Node(7, "localhost", 2078));
         nodesById.put(8, new Node(8, "localhost", 2079));
 
-        MetadataCache cache = new MetadataCache("clusterId",
+        MetadataSnapshot cache = new MetadataSnapshot("clusterId",
                 nodesById,
                 Collections.singleton(partitionMetadata),
                 Collections.emptySet(),
@@ -107,7 +108,7 @@ public void testMergeWithThatPreExistingPartitionIsRetainedPostMerge() {
         Uuid topic1Id = Uuid.randomUuid();
         topicsIds.put(topic1Partition.topic(), topic1Id);
 
-        MetadataCache cache = new MetadataCache("clusterId",
+        MetadataSnapshot cache = new MetadataSnapshot("clusterId",
             nodesById,
             Collections.singleton(partitionMetadata1),
             Collections.emptySet(),
@@ -155,7 +156,7 @@ public void testTopicNamesCacheBuiltFromTopicIds() {
         topicIds.put("topic1", Uuid.randomUuid());
         topicIds.put("topic2", Uuid.randomUuid());
 
-        MetadataCache cache = new MetadataCache("clusterId",
+        MetadataSnapshot cache = new MetadataSnapshot("clusterId",
                 Collections.singletonMap(6, new Node(6, "localhost", 2077)),
                 Collections.emptyList(),
                 Collections.emptySet(),
@@ -174,7 +175,7 @@ public void testTopicNamesCacheBuiltFromTopicIds() {
     public void testEmptyTopicNamesCacheBuiltFromTopicIds() {
         Map<String, Uuid> topicIds = new HashMap<>();
 
-        MetadataCache cache = new MetadataCache("clusterId",
+        MetadataSnapshot cache = new MetadataSnapshot("clusterId",
                 Collections.singletonMap(6, new Node(6, "localhost", 2077)),
                 Collections.emptyList(),
                 Collections.emptySet(),
@@ -185,4 +186,49 @@ public void testEmptyTopicNamesCacheBuiltFromTopicIds() {
         assertEquals(Collections.emptyMap(), cache.topicNames());
     }
 
+    @Test
+    public void testLeaderEpochFor() {
+        // Setup partition 0 with a leader-epoch of 10.
+        TopicPartition topicPartition1 = new TopicPartition("topic", 0);
+        MetadataResponse.PartitionMetadata partitionMetadata1 = new MetadataResponse.PartitionMetadata(
+            Errors.NONE,
+            topicPartition1,
+            Optional.of(5),
+            Optional.of(10),
+            Arrays.asList(5, 6, 7),
+            Arrays.asList(5, 6, 7),
+            Collections.emptyList());
+
+        // Setup partition 1 with an unknown leader epoch.
+        TopicPartition topicPartition2 = new TopicPartition("topic", 1);
+        MetadataResponse.PartitionMetadata partitionMetadata2 = new MetadataResponse.PartitionMetadata(
+            Errors.NONE,
+            topicPartition2,
+            Optional.of(5),
+            Optional.empty(),
+            Arrays.asList(5, 6, 7),
+            Arrays.asList(5, 6, 7),
+            Collections.emptyList());
+
+        Map<Integer, Node> nodesById = new HashMap<>();
+        nodesById.put(5, new Node(5, "localhost", 2077));
+        nodesById.put(6, new Node(6, "localhost", 2078));
+        nodesById.put(7, new Node(7, "localhost", 2079));
+
+        MetadataSnapshot cache = new MetadataSnapshot("clusterId",
+            nodesById,
+            Arrays.asList(partitionMetadata1, partitionMetadata2),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            null,
+            Collections.emptyMap());
+
+        assertEquals(OptionalInt.of(10), cache.leaderEpochFor(topicPartition1));
+
+        assertEquals(OptionalInt.empty(), cache.leaderEpochFor(topicPartition2));
+
+        assertEquals(OptionalInt.empty(), cache.leaderEpochFor(new TopicPartition("topic_missing", 0)));
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
index d7c91fdffadac..600fc23ecb988 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
@@ -604,11 +604,11 @@ public void testClusterCopy() {
 
         // Sentinel instances
         InetSocketAddress address = InetSocketAddress.createUnresolved("localhost", 0);
-        Cluster fromMetadata = MetadataCache.bootstrap(Collections.singletonList(address)).cluster();
+        Cluster fromMetadata = MetadataSnapshot.bootstrap(Collections.singletonList(address)).cluster();
         Cluster fromCluster = Cluster.bootstrap(Collections.singletonList(address));
         assertEquals(fromMetadata, fromCluster);
 
-        Cluster fromMetadataEmpty = MetadataCache.empty().cluster();
+        Cluster fromMetadataEmpty = MetadataSnapshot.empty().cluster();
         Cluster fromClusterEmpty = Cluster.empty();
         assertEquals(fromMetadataEmpty, fromClusterEmpty);
     }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerBatchTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerBatchTest.java
index 24629b612b298..7f98e08e534b1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerBatchTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerBatchTest.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
-import java.util.Optional;
+import java.util.OptionalInt;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.KafkaException;
@@ -276,41 +276,55 @@ public void testWithLeaderChangesAcrossRetries() {
         ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 1), memoryRecordsBuilder, now);
 
         // Starting state for the batch, no attempt made to send it yet.
-        assertEquals(Optional.empty(), batch.currentLeaderEpoch());
+        assertEquals(OptionalInt.empty(), batch.currentLeaderEpoch());
         assertEquals(0, batch.attemptsWhenLeaderLastChanged()); // default value
-        batch.maybeUpdateLeaderEpoch(Optional.empty());
+        batch.maybeUpdateLeaderEpoch(OptionalInt.empty());
         assertFalse(batch.hasLeaderChangedForTheOngoingRetry());
 
         // 1st attempt[Not a retry] to send the batch.
         // Check leader isn't flagged as a new leader.
         int batchLeaderEpoch = 100;
-        batch.maybeUpdateLeaderEpoch(Optional.of(batchLeaderEpoch));
+        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
         assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader is assigned for 1st time");
-        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().get());
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
         assertEquals(0, batch.attemptsWhenLeaderLastChanged());
 
         // 2nd attempt[1st retry] to send the batch to a new leader.
         // Check leader change is detected.
         batchLeaderEpoch = 101;
         batch.reenqueued(0);
-        batch.maybeUpdateLeaderEpoch(Optional.of(batchLeaderEpoch));
+        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
         assertTrue(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has changed");
-        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().get());
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
         assertEquals(1, batch.attemptsWhenLeaderLastChanged());
 
         // 2nd attempt[1st retry] still ongoing, yet to be made.
         // Check same leaderEpoch(101) is still considered as a leader-change.
-        batch.maybeUpdateLeaderEpoch(Optional.of(batchLeaderEpoch));
+        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
         assertTrue(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has changed");
-        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().get());
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
         assertEquals(1, batch.attemptsWhenLeaderLastChanged());
 
         // 3rd attempt[2nd retry] to the same leader-epoch(101).
         // Check same leaderEpoch(101) as not detected as a leader-change.
         batch.reenqueued(0);
-        batch.maybeUpdateLeaderEpoch(Optional.of(batchLeaderEpoch));
+        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
         assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
-        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().get());
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
+        assertEquals(1, batch.attemptsWhenLeaderLastChanged());
+
+        // Attempt made to update batch leader-epoch to an older leader-epoch(100).
+        // Check batch leader-epoch remains unchanged as 101.
+        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch - 1));
+        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
+        assertEquals(1, batch.attemptsWhenLeaderLastChanged());
+
+        // Attempt made to update batch leader-epoch to an unknown leader(optional.empty())
+        // Check batch leader-epoch remains unchanged as 101.
+        batch.maybeUpdateLeaderEpoch(OptionalInt.empty());
+        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
+        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
         assertEquals(1, batch.attemptsWhenLeaderLastChanged());
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java
index a046efe2cb2ac..9ed0e70380969 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java
@@ -17,10 +17,11 @@
 package org.apache.kafka.clients.producer.internals;
 
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.function.Function;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.Partitioner;
@@ -33,6 +34,7 @@
 import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.CompressionRatioEstimator;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.DefaultRecord;
@@ -42,6 +44,8 @@
 import org.apache.kafka.common.record.MutableRecordBatch;
 import org.apache.kafka.common.record.Record;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.ProducerIdAndEpoch;
@@ -89,29 +93,38 @@ public class RecordAccumulatorTest {
     private TopicPartition tp1 = new TopicPartition(topic, partition1);
     private TopicPartition tp2 = new TopicPartition(topic, partition2);
     private TopicPartition tp3 = new TopicPartition(topic, partition3);
-    private PartitionInfo part1 = new PartitionInfo(topic, partition1, node1, null, null);
-    private PartitionInfo part2 = new PartitionInfo(topic, partition2, node1, null, null);
-    private PartitionInfo part3 = new PartitionInfo(topic, partition3, node2, null, null);
+
+    private PartitionMetadata partMetadata1 = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
+    private PartitionMetadata partMetadata2 = new PartitionMetadata(Errors.NONE, tp2, Optional.of(node1.id()), Optional.empty(), null, null, null);
+    private PartitionMetadata partMetadata3 = new PartitionMetadata(Errors.NONE, tp3, Optional.of(node2.id()), Optional.empty(), null, null, null);
+    private List<PartitionMetadata> partMetadatas = new ArrayList<>(Arrays.asList(partMetadata1, partMetadata2, partMetadata3));
+
+    private Map<Integer, Node> nodes = Arrays.asList(node1, node2).stream().collect(Collectors.toMap(Node::id, Function.identity()));
+    private MetadataSnapshot metadataCache = new MetadataSnapshot(null,
+        nodes,
+        partMetadatas,
+        Collections.emptySet(),
+        Collections.emptySet(),
+        Collections.emptySet(),
+        null,
+        Collections.emptyMap());
+
+    private Cluster cluster = metadataCache.cluster();
+
     private MockTime time = new MockTime();
     private byte[] key = "key".getBytes();
     private byte[] value = "value".getBytes();
     private int msgSize = DefaultRecord.sizeInBytes(0, 0, key.length, value.length, Record.EMPTY_HEADERS);
-    Metadata metadataMock;
-    private Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3),
-        Collections.emptySet(), Collections.emptySet());
+
     private Metrics metrics = new Metrics(time);
     private final long maxBlockTimeMs = 1000;
     private final LogContext logContext = new LogContext();
 
-    @BeforeEach
-    public void setup() {
-        metadataMock = setupMetadata(cluster);
-    }
+    @BeforeEach void setup() {}
 
     @AfterEach
     public void teardown() {
         this.metrics.close();
-        Mockito.reset(metadataMock);
     }
 
     @Test
@@ -120,13 +133,31 @@ public void testDrainBatches() throws Exception {
         // add tp-4
         int partition4 = 3;
         TopicPartition tp4 = new TopicPartition(topic, partition4);
-        PartitionInfo part4 = new PartitionInfo(topic, partition4, node2, null, null);
+        PartitionMetadata partMetadata4 = new PartitionMetadata(Errors.NONE, tp4, Optional.of(node2.id()), Optional.empty(), null, null, null);
+        partMetadatas.add(partMetadata4);
+
+        // This test requires that partitions to be drained in order for each node i.e.
+        // node1 -> tp1, tp3, and node2 -> tp2, tp4.
+        // So setup cluster with this order, and pass this cluster to MetadataCache to preserve this order.
+        PartitionInfo part1 = MetadataResponse.toPartitionInfo(partMetadata1, nodes);
+        PartitionInfo part2 = MetadataResponse.toPartitionInfo(partMetadata2, nodes);
+        PartitionInfo part3 = MetadataResponse.toPartitionInfo(partMetadata3, nodes);
+        PartitionInfo part4 = MetadataResponse.toPartitionInfo(partMetadata4, nodes);
+        Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3, part4),
+            Collections.emptySet(), Collections.emptySet());
 
+        metadataCache = new MetadataSnapshot(null,
+            nodes,
+            partMetadatas,
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            null,
+            Collections.emptyMap(),
+            cluster);
         long batchSize = value.length + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
         RecordAccumulator accum = createTestRecordAccumulator((int) batchSize, Integer.MAX_VALUE, CompressionType.NONE, 10);
-        Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3, part4),
-                Collections.emptySet(), Collections.emptySet());
-        metadataMock = setupMetadata(cluster);
+
 
         //  initial data
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
@@ -135,7 +166,7 @@ public void testDrainBatches() throws Exception {
         accum.append(topic, partition4, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
 
         // drain batches from 2 nodes: node1 => tp1, node2 => tp3, because the max request size is full after the first batch drained
-        Map<Integer, List<ProducerBatch>> batches1 = accum.drain(metadataMock, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
+        Map<Integer, List<ProducerBatch>> batches1 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
         verifyTopicPartitionInBatches(batches1, tp1, tp3);
 
         // add record for tp1, tp3
@@ -144,11 +175,11 @@ public void testDrainBatches() throws Exception {
 
         // drain batches from 2 nodes: node1 => tp2, node2 => tp4, because the max request size is full after the first batch drained
         // The drain index should start from next topic partition, that is, node1 => tp2, node2 => tp4
-        Map<Integer, List<ProducerBatch>> batches2 = accum.drain(metadataMock, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
+        Map<Integer, List<ProducerBatch>> batches2 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
         verifyTopicPartitionInBatches(batches2, tp2, tp4);
 
         // make sure in next run, the drain index will start from the beginning
-        Map<Integer, List<ProducerBatch>> batches3 = accum.drain(metadataMock, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
+        Map<Integer, List<ProducerBatch>> batches3 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
         verifyTopicPartitionInBatches(batches3, tp1, tp3);
 
         // add record for tp2, tp3, tp4 and mute the tp4
@@ -157,7 +188,7 @@ public void testDrainBatches() throws Exception {
         accum.append(topic, partition4, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
         accum.mutePartition(tp4);
         // drain batches from 2 nodes: node1 => tp2, node2 => tp3 (because tp4 is muted)
-        Map<Integer, List<ProducerBatch>> batches4 = accum.drain(metadataMock, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
+        Map<Integer, List<ProducerBatch>> batches4 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), (int) batchSize, 0);
         verifyTopicPartitionInBatches(batches4, tp2, tp3);
 
         // add record for tp1, tp2, tp3, and unmute tp4
@@ -166,7 +197,7 @@ public void testDrainBatches() throws Exception {
         accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
         accum.unmutePartition(tp4);
         // set maxSize as a max value, so that the all partitions in 2 nodes should be drained: node1 => [tp1, tp2], node2 => [tp3, tp4]
-        Map<Integer, List<ProducerBatch>> batches5 = accum.drain(metadataMock, new HashSet<>(Arrays.asList(node1, node2)), Integer.MAX_VALUE, 0);
+        Map<Integer, List<ProducerBatch>> batches5 = accum.drain(metadataCache, new HashSet<>(Arrays.asList(node1, node2)), Integer.MAX_VALUE, 0);
         verifyTopicPartitionInBatches(batches5, tp1, tp2, tp3, tp4);
     }
 
@@ -197,25 +228,25 @@ public void testFull() throws Exception {
         int appends = expectedNumAppends(batchSize);
         for (int i = 0; i < appends; i++) {
             // append to the first batch
-            accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
+            accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), metadataCache.cluster());
             Deque<ProducerBatch> partitionBatches = accum.getDeque(tp1);
             assertEquals(1, partitionBatches.size());
 
             ProducerBatch batch = partitionBatches.peekFirst();
             assertTrue(batch.isWritable());
-            assertEquals(0, accum.ready(metadataMock, now).readyNodes.size(), "No partitions should be ready.");
+            assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready.");
         }
 
         // this append doesn't fit in the first batch, so a new batch is created and the first batch is closed
 
-        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
+        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), metadataCache.cluster());
         Deque<ProducerBatch> partitionBatches = accum.getDeque(tp1);
         assertEquals(2, partitionBatches.size());
         Iterator<ProducerBatch> partitionBatchesIterator = partitionBatches.iterator();
         assertTrue(partitionBatchesIterator.next().isWritable());
-        assertEquals(Collections.singleton(node1), accum.ready(metadataMock, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
+        assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
 
-        List<ProducerBatch> batches = accum.drain(metadataMock, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
+        List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
         assertEquals(1, batches.size());
         ProducerBatch batch = batches.get(0);
 
@@ -243,8 +274,8 @@ private void testAppendLarge(CompressionType compressionType) throws Exception {
         byte[] value = new byte[2 * batchSize];
         RecordAccumulator accum = createTestRecordAccumulator(
                 batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0);
-        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        assertEquals(Collections.singleton(node1), accum.ready(metadataMock, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
+        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), metadataCache.cluster());
+        assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
 
         Deque<ProducerBatch> batches = accum.getDeque(tp1);
         assertEquals(1, batches.size());
@@ -281,8 +312,8 @@ private void testAppendLargeOldMessageFormat(CompressionType compressionType) th
 
         RecordAccumulator accum = createTestRecordAccumulator(
                 batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, compressionType, 0);
-        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        assertEquals(Collections.singleton(node1), accum.ready(metadataMock, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
+        accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), metadataCache.cluster());
+        assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
 
         Deque<ProducerBatch> batches = accum.getDeque(tp1);
         assertEquals(1, batches.size());
@@ -306,10 +337,10 @@ public void testLinger() throws Exception {
         RecordAccumulator accum = createTestRecordAccumulator(
                 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024, CompressionType.NONE, lingerMs);
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready");
+        assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready");
         time.sleep(10);
-        assertEquals(Collections.singleton(node1), accum.ready(metadataMock, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
-        List<ProducerBatch> batches = accum.drain(metadataMock, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
+        assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready");
+        List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id());
         assertEquals(1, batches.size());
         ProducerBatch batch = batches.get(0);
 
@@ -330,9 +361,9 @@ public void testPartialDrain() throws Exception {
             for (int i = 0; i < appends; i++)
                 accum.append(tp.topic(), tp.partition(), 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
         }
-        assertEquals(Collections.singleton(node1), accum.ready(metadataMock, time.milliseconds()).readyNodes, "Partition's leader should be ready");
+        assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Partition's leader should be ready");
 
-        List<ProducerBatch> batches = accum.drain(metadataMock, Collections.singleton(node1), 1024, 0).get(node1.id());
+        List<ProducerBatch> batches = accum.drain(metadataCache, Collections.singleton(node1), 1024, 0).get(node1.id());
         assertEquals(1, batches.size(), "But due to size bound only one partition should have been retrieved");
     }
 
@@ -361,8 +392,8 @@ public void testStressfulSituation() throws Exception {
         int read = 0;
         long now = time.milliseconds();
         while (read < numThreads * msgs) {
-            Set<Node> nodes = accum.ready(metadataMock, now).readyNodes;
-            List<ProducerBatch> batches = accum.drain(metadataMock, nodes, 5 * 1024, 0).get(node1.id());
+            Set<Node> nodes = accum.ready(metadataCache, now).readyNodes;
+            List<ProducerBatch> batches = accum.drain(metadataCache, nodes, 5 * 1024, 0).get(node1.id());
             if (batches != null) {
                 for (ProducerBatch batch : batches) {
                     for (Record record : batch.records().records())
@@ -393,7 +424,7 @@ public void testNextReadyCheckDelay() throws Exception {
         // Partition on node1 only
         for (int i = 0; i < appends; i++)
             accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
         assertEquals(lingerMs, result.nextReadyCheckDelayMs, "Next check time should be the linger time");
 
@@ -402,14 +433,14 @@ public void testNextReadyCheckDelay() throws Exception {
         // Add partition on node2 only
         for (int i = 0; i < appends; i++)
             accum.append(topic, partition3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        result = accum.ready(metadataMock, time.milliseconds());
+        result = accum.ready(metadataCache, time.milliseconds());
         assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
         assertEquals(lingerMs / 2, result.nextReadyCheckDelayMs, "Next check time should be defined by node1, half remaining linger time");
 
         // Add data for another partition on node1, enough to make data sendable immediately
         for (int i = 0; i < appends + 1; i++)
             accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        result = accum.ready(metadataMock, time.milliseconds());
+        result = accum.ready(metadataCache, time.milliseconds());
         assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
         // Note this can actually be < linger time because it may use delays from partitions that aren't sendable
         // but have leaders with other sendable data.
@@ -433,9 +464,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
 
         long now = time.milliseconds();
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now + lingerMs + 1);
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now + lingerMs + 1);
         assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
-        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
+        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
         assertEquals(1, batches.size(), "Node1 should be the only ready node.");
         assertEquals(1, batches.get(0).size(), "Partition 0 should only have one batch drained.");
 
@@ -445,37 +476,37 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
 
         // Put message for partition 1 into accumulator
         accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        result = accum.ready(metadataMock, now + lingerMs + 1);
+        result = accum.ready(metadataCache, now + lingerMs + 1);
         assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
 
         // tp1 should backoff while tp2 should not
-        batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
+        batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + lingerMs + 1);
         assertEquals(1, batches.size(), "Node1 should be the only ready node.");
         assertEquals(1, batches.get(0).size(), "Node1 should only have one batch drained.");
         assertEquals(tp2, batches.get(0).get(0).topicPartition, "Node1 should only have one batch for partition 1.");
 
         // Partition 0 can be drained after retry backoff
         long upperBoundBackoffMs = (long) (retryBackoffMs * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
-        result = accum.ready(metadataMock, now + upperBoundBackoffMs + 1);
+        result = accum.ready(metadataCache, now + upperBoundBackoffMs + 1);
         assertEquals(Collections.singleton(node1), result.readyNodes, "Node1 should be ready");
-        batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, now + upperBoundBackoffMs + 1);
+        batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now + upperBoundBackoffMs + 1);
         assertEquals(1, batches.size(), "Node1 should be the only ready node.");
         assertEquals(1, batches.get(0).size(), "Node1 should only have one batch drained.");
         assertEquals(tp1, batches.get(0).get(0).topicPartition, "Node1 should only have one batch for partition 0.");
     }
 
-    private Map<Integer, List<ProducerBatch>> drainAndCheckBatchAmount(Cluster cluster, Node leader, RecordAccumulator accum, long now, int expected) {
-        metadataMock = setupMetadata(cluster);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+    private Map<Integer, List<ProducerBatch>> drainAndCheckBatchAmount(
+        MetadataSnapshot metadataCache, Node leader, RecordAccumulator accum, long now, int expected) {
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
         if (expected > 0) {
             assertEquals(Collections.singleton(leader), result.readyNodes, "Leader should be ready");
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, now);
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now);
             assertEquals(expected, batches.size(), "Leader should be the only ready node.");
             assertEquals(expected, batches.get(leader.id()).size(), "Partition should only have " + expected + " batch drained.");
             return batches;
         } else {
             assertEquals(0, result.readyNodes.size(), "Leader should not be ready");
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, now);
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, now);
             assertEquals(0, batches.size(), "Leader should not be drained.");
             return null;
         }
@@ -501,7 +532,7 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
 
         // No backoff for initial attempt
-        Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(cluster, node1, accum, now + lingerMs + 1, 1);
+        Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(metadataCache, node1, accum, now + lingerMs + 1, 1);
         ProducerBatch batch = batches.get(0).get(0);
         long currentRetryBackoffMs = 0;
 
@@ -513,9 +544,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
             long upperBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, i) * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
             currentRetryBackoffMs = upperBoundBackoffMs;
             // Should back off
-            drainAndCheckBatchAmount(cluster, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
+            drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
             // Should not back off
-            drainAndCheckBatchAmount(cluster, node1, accum, initial + upperBoundBackoffMs + 1, 1);
+            drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
         }
     }
 
@@ -529,14 +560,28 @@ public void testExponentialRetryBackoffLeaderChange() throws Exception {
         int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
         String metricGrpName = "producer-metrics";
 
-        PartitionInfo part1 = new PartitionInfo(topic, partition1, node1, null, null);
-        PartitionInfo part1Change = new PartitionInfo(topic, partition1, node2, null, null);
-        PartitionInfo part2 = new PartitionInfo(topic, partition2, node1, null, null);
-        PartitionInfo part3 = new PartitionInfo(topic, partition3, node2, null, null);
-        Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3),
-                Collections.emptySet(), Collections.emptySet());
-        Cluster clusterChange = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1Change, part2, part3),
-                Collections.emptySet(), Collections.emptySet());
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        PartitionMetadata part1MetadataChange = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node2.id()), Optional.empty(), null, null, null);
+        PartitionMetadata part2Metadata = new PartitionMetadata(Errors.NONE, tp2, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        PartitionMetadata part3Metadata = new PartitionMetadata(Errors.NONE, tp3, Optional.of(node2.id()), Optional.empty(), null, null, null);
+
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null,
+            nodes,
+            Arrays.asList(part1Metadata, part2Metadata, part3Metadata),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            null,
+            Collections.emptyMap());
+
+        MetadataSnapshot metadataCacheChange = new MetadataSnapshot(null,
+            nodes,
+            Arrays.asList(part1MetadataChange, part2Metadata, part3Metadata),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            null,
+            Collections.emptyMap());
 
         final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
                 CompressionType.NONE, lingerMs, retryBackoffMs, retryBackoffMaxMs,
@@ -548,7 +593,7 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
 
         // No backoff for initial attempt
-        Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(cluster, node1, accum, now + lingerMs + 1, 1);
+        Map<Integer, List<ProducerBatch>> batches = drainAndCheckBatchAmount(metadataCache, node1, accum, now + lingerMs + 1, 1);
         ProducerBatch batch = batches.get(0).get(0);
 
         long lowerBoundBackoffMs;
@@ -560,9 +605,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         lowerBoundBackoffMs = (long) (retryBackoffMs * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
         upperBoundBackoffMs = (long) (retryBackoffMs * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
         // Should back off
-        drainAndCheckBatchAmount(cluster, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
+        drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
         // Should not back off
-        drainAndCheckBatchAmount(cluster, node1, accum, initial + upperBoundBackoffMs + 1, 1);
+        drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
 
         // Retry 2 - delay by retryBackoffMs * 2 +/- jitter
         now = time.milliseconds();
@@ -570,9 +615,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         lowerBoundBackoffMs = (long) (retryBackoffMs * CommonClientConfigs.RETRY_BACKOFF_EXP_BASE * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
         upperBoundBackoffMs = (long) (retryBackoffMs * CommonClientConfigs.RETRY_BACKOFF_EXP_BASE * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
         // Should back off
-        drainAndCheckBatchAmount(cluster, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
+        drainAndCheckBatchAmount(metadataCache, node1, accum, initial + lowerBoundBackoffMs - 1, 0);
         // Should not back off
-        drainAndCheckBatchAmount(cluster, node1, accum, initial + upperBoundBackoffMs + 1, 1);
+        drainAndCheckBatchAmount(metadataCache, node1, accum, initial + upperBoundBackoffMs + 1, 1);
 
         // Retry 3 - after a leader change, delay by retryBackoffMs * 2^2 +/- jitter (could optimise to do not delay at all)
         now = time.milliseconds();
@@ -580,9 +625,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         lowerBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 2) * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
         upperBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 2) * (1 + CommonClientConfigs.RETRY_BACKOFF_JITTER));
         // Should back off
-        drainAndCheckBatchAmount(clusterChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
+        drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
         // Should not back off
-        drainAndCheckBatchAmount(clusterChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
+        drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
 
         // Retry 4 - delay by retryBackoffMs * 2^3 +/- jitter (capped to retryBackoffMaxMs)
         now = time.milliseconds();
@@ -590,9 +635,9 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         lowerBoundBackoffMs = (long) (retryBackoffMs * Math.pow(CommonClientConfigs.RETRY_BACKOFF_EXP_BASE, 3) * (1 - CommonClientConfigs.RETRY_BACKOFF_JITTER));
         upperBoundBackoffMs = retryBackoffMaxMs;
         // Should back off
-        drainAndCheckBatchAmount(clusterChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
+        drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + lowerBoundBackoffMs - 1, 0);
         // Should not back off
-        drainAndCheckBatchAmount(clusterChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
+        drainAndCheckBatchAmount(metadataCacheChange, node2, accum, initial + upperBoundBackoffMs + 1, 1);
     }
 
     @Test
@@ -605,14 +650,14 @@ public void testFlush() throws Exception {
             accum.append(topic, i % 3, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
             assertTrue(accum.hasIncomplete());
         }
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertEquals(0, result.readyNodes.size(), "No nodes should be ready.");
 
         accum.beginFlush();
-        result = accum.ready(metadataMock, time.milliseconds());
+        result = accum.ready(metadataCache, time.milliseconds());
 
         // drain and deallocate all batches
-        Map<Integer, List<ProducerBatch>> results = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> results = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertTrue(accum.hasIncomplete());
 
         for (List<ProducerBatch> batches: results.values())
@@ -672,9 +717,9 @@ public void setPartition(int partition) {
         }
         for (int i = 0; i < numRecords; i++)
             accum.append(topic, i % 3, 0L, key, value, null, new TestCallback(), maxBlockTimeMs, false, time.milliseconds(), cluster);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertTrue(accum.hasUndrained());
         assertTrue(accum.hasIncomplete());
 
@@ -717,9 +762,9 @@ public void setPartition(int partition) {
         }
         for (int i = 0; i < numRecords; i++)
             accum.append(topic, i % 3, 0L, key, value, null, new TestCallback(), maxBlockTimeMs, false, time.milliseconds(), cluster);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE,
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE,
                 time.milliseconds());
         assertTrue(accum.hasUndrained());
         assertTrue(accum.hasIncomplete());
@@ -756,10 +801,10 @@ private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedExcept
             if (time.milliseconds() < System.currentTimeMillis())
                 time.setCurrentTimeMs(System.currentTimeMillis());
             accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-            assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partition should be ready.");
+            assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partition should be ready.");
 
             time.sleep(lingerMs);
-            readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+            readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
             assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
 
             expiredBatches = accum.expiredBatches(time.milliseconds());
@@ -774,7 +819,7 @@ private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedExcept
             time.sleep(deliveryTimeoutMs - lingerMs);
             expiredBatches = accum.expiredBatches(time.milliseconds());
             assertEquals(1, expiredBatches.size(), "The batch may expire when the partition is muted");
-            assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
+            assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
         }
     }
 
@@ -805,11 +850,11 @@ public void testExpiredBatches() throws InterruptedException {
         // Test batches not in retry
         for (int i = 0; i < appends; i++) {
             accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-            assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
+            assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
         }
         // Make the batches ready due to batch full
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false, time.milliseconds(), cluster);
-        Set<Node> readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+        Set<Node> readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
         assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
         // Advance the clock to expire the batch.
         time.sleep(deliveryTimeoutMs + 1);
@@ -820,7 +865,7 @@ public void testExpiredBatches() throws InterruptedException {
         accum.unmutePartition(tp1);
         expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals(0, expiredBatches.size(), "All batches should have been expired earlier");
-        assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
+        assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
 
         // Advance the clock to make the next batch ready due to linger.ms
         time.sleep(lingerMs);
@@ -834,15 +879,15 @@ public void testExpiredBatches() throws InterruptedException {
         accum.unmutePartition(tp1);
         expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals(0, expiredBatches.size(), "All batches should have been expired");
-        assertEquals(0, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
+        assertEquals(0, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
 
         // Test batches in retry.
         // Create a retried batch
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false, time.milliseconds(), cluster);
         time.sleep(lingerMs);
-        readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+        readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
         assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(drained.get(node1.id()).size(), 1, "There should be only one batch.");
         time.sleep(1000L);
         accum.reenqueue(drained.get(node1.id()).get(0), time.milliseconds());
@@ -864,7 +909,7 @@ public void testExpiredBatches() throws InterruptedException {
         // Test that when being throttled muted batches are expired before the throttle time is over.
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false, time.milliseconds(), cluster);
         time.sleep(lingerMs);
-        readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+        readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
         assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
         // Advance the clock to expire the batch.
         time.sleep(requestTimeout + 1);
@@ -882,7 +927,7 @@ public void testExpiredBatches() throws InterruptedException {
         time.sleep(throttleTimeMs);
         expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals(0, expiredBatches.size(), "All batches should have been expired earlier");
-        assertEquals(1, accum.ready(metadataMock, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
+        assertEquals(1, accum.ready(metadataCache, time.milliseconds()).readyNodes.size(), "No partitions should be ready.");
     }
 
     @Test
@@ -896,28 +941,28 @@ public void testMutedPartitions() throws InterruptedException {
         int appends = expectedNumAppends(batchSize);
         for (int i = 0; i < appends; i++) {
             accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-            assertEquals(0, accum.ready(metadataMock, now).readyNodes.size(), "No partitions should be ready.");
+            assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready.");
         }
         time.sleep(2000);
 
         // Test ready with muted partition
         accum.mutePartition(tp1);
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertEquals(0, result.readyNodes.size(), "No node should be ready");
 
         // Test ready without muted partition
         accum.unmutePartition(tp1);
-        result = accum.ready(metadataMock, time.milliseconds());
+        result = accum.ready(metadataCache, time.milliseconds());
         assertTrue(result.readyNodes.size() > 0, "The batch should be ready");
 
         // Test drain with muted partition
         accum.mutePartition(tp1);
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(0, drained.get(node1.id()).size(), "No batch should have been drained");
 
         // Test drain without muted partition.
         accum.unmutePartition(tp1);
-        drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertTrue(drained.get(node1.id()).size() > 0, "The batch should have been drained.");
     }
 
@@ -967,20 +1012,20 @@ public void testRecordsDrainedWhenTransactionCompleting() throws Exception {
             false, time.milliseconds(), cluster);
         assertTrue(accumulator.hasUndrained());
 
-        RecordAccumulator.ReadyCheckResult firstResult = accumulator.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult firstResult = accumulator.ready(metadataCache, time.milliseconds());
         assertEquals(0, firstResult.readyNodes.size());
-        Map<Integer, List<ProducerBatch>> firstDrained = accumulator.drain(metadataMock, firstResult.readyNodes,
+        Map<Integer, List<ProducerBatch>> firstDrained = accumulator.drain(metadataCache, firstResult.readyNodes,
             Integer.MAX_VALUE, time.milliseconds());
         assertEquals(0, firstDrained.size());
 
         // Once the transaction begins completion, then the batch should be drained immediately.
         Mockito.when(transactionManager.isCompleting()).thenReturn(true);
 
-        RecordAccumulator.ReadyCheckResult secondResult = accumulator.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult secondResult = accumulator.ready(metadataCache, time.milliseconds());
         assertEquals(1, secondResult.readyNodes.size());
         Node readyNode = secondResult.readyNodes.iterator().next();
 
-        Map<Integer, List<ProducerBatch>> secondDrained = accumulator.drain(metadataMock, secondResult.readyNodes,
+        Map<Integer, List<ProducerBatch>> secondDrained = accumulator.drain(metadataCache, secondResult.readyNodes,
             Integer.MAX_VALUE, time.milliseconds());
         assertEquals(Collections.singleton(readyNode.id()), secondDrained.keySet());
         List<ProducerBatch> batches = secondDrained.get(readyNode.id());
@@ -1011,16 +1056,16 @@ public void testSplitAndReenqueue() throws ExecutionException, InterruptedExcept
         // Re-enqueuing counts as a second attempt, so the delay with jitter is 100 * (1 + 0.2) + 1
         time.sleep(121L);
         // Drain the batch.
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertTrue(result.readyNodes.size() > 0, "The batch should be ready");
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(1, drained.size(), "Only node1 should be drained");
         assertEquals(1, drained.get(node1.id()).size(), "Only one batch should be drained");
         // Split and reenqueue the batch.
         accum.splitAndReenqueue(drained.get(node1.id()).get(0));
         time.sleep(101L);
 
-        drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertFalse(drained.isEmpty());
         assertFalse(drained.get(node1.id()).isEmpty());
         drained.get(node1.id()).get(0).complete(acked.get(), 100L);
@@ -1028,7 +1073,7 @@ public void testSplitAndReenqueue() throws ExecutionException, InterruptedExcept
         assertTrue(future1.isDone());
         assertEquals(0, future1.get().offset());
 
-        drained = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        drained = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertFalse(drained.isEmpty());
         assertFalse(drained.get(node1.id()).isEmpty());
         drained.get(node1.id()).get(0).complete(acked.get(), 100L);
@@ -1049,14 +1094,14 @@ public void testSplitBatchOffAccumulator() throws InterruptedException {
         int numSplitBatches = prepareSplitBatches(accum, seed, 100, 20);
         assertTrue(numSplitBatches > 0, "There should be some split batches");
         // Drain all the split batches.
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         for (int i = 0; i < numSplitBatches; i++) {
             Map<Integer, List<ProducerBatch>> drained =
-                accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+                accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
             assertFalse(drained.isEmpty());
             assertFalse(drained.get(node1.id()).isEmpty());
         }
-        assertTrue(accum.ready(metadataMock, time.milliseconds()).readyNodes.isEmpty(), "All the batches should have been drained.");
+        assertTrue(accum.ready(metadataCache, time.milliseconds()).readyNodes.isEmpty(), "All the batches should have been drained.");
         assertEquals(bufferCapacity, accum.bufferPoolAvailableMemory(),
             "The split batches should be allocated off the accumulator");
     }
@@ -1103,16 +1148,16 @@ public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedExce
             batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
 
         accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
-        Set<Node> readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
-        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Set<Node> readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
+        Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertTrue(drained.isEmpty());
         //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
 
         // advanced clock and send one batch out but it should not be included in soon to expire inflight
         // batches because batch's expiry is quite far.
         time.sleep(lingerMs + 1);
-        readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
-        drained = accum.drain(metadataMock, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
+        drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(1, drained.size(), "A batch did not drain after linger");
         //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
 
@@ -1121,8 +1166,8 @@ public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedExce
         time.sleep(lingerMs * 4);
 
         // Now drain and check that accumulator picked up the drained batch because its expiry is soon.
-        readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
-        drained = accum.drain(metadataMock, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
+        drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(1, drained.size(), "A batch did not drain after linger");
     }
 
@@ -1144,9 +1189,9 @@ public void testExpiredBatchesRetry() throws InterruptedException {
         for (Boolean mute : muteStates) {
             accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false, time.milliseconds(), cluster);
             time.sleep(lingerMs);
-            readyNodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+            readyNodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
             assertEquals(Collections.singleton(node1), readyNodes, "Our partition's leader should be ready");
-            Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataMock, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+            Map<Integer, List<ProducerBatch>> drained = accum.drain(metadataCache, readyNodes, Integer.MAX_VALUE, time.milliseconds());
             assertEquals(1, drained.get(node1.id()).size(), "There should be only one batch.");
             time.sleep(rtt);
             accum.reenqueue(drained.get(node1.id()).get(0), time.milliseconds());
@@ -1158,7 +1203,7 @@ public void testExpiredBatchesRetry() throws InterruptedException {
 
             // test expiration
             time.sleep(deliveryTimeoutMs - rtt);
-            accum.drain(metadataMock, Collections.singleton(node1), Integer.MAX_VALUE, time.milliseconds());
+            accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, time.milliseconds());
             expiredBatches = accum.expiredBatches(time.milliseconds());
             assertEquals(mute ? 1 : 0, expiredBatches.size(), "RecordAccumulator has expired batches if the partition is not muted");
         }
@@ -1199,12 +1244,12 @@ public void testStickyBatches() throws Exception {
             // We only appended if we do not retry.
             if (!switchPartition) {
                 appends++;
-                assertEquals(0, accum.ready(metadataMock, now).readyNodes.size(), "No partitions should be ready.");
+                assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready.");
             }
         }
 
         // Batch should be full.
-        assertEquals(1, accum.ready(metadataMock, time.milliseconds()).readyNodes.size());
+        assertEquals(1, accum.ready(metadataCache, time.milliseconds()).readyNodes.size());
         assertEquals(appends, expectedAppends);
         switchPartition = false;
 
@@ -1263,6 +1308,12 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
                 }
             };
 
+            PartitionInfo part1 = MetadataResponse.toPartitionInfo(partMetadata1, nodes);
+            PartitionInfo part2 = MetadataResponse.toPartitionInfo(partMetadata2, nodes);
+            PartitionInfo part3 = MetadataResponse.toPartitionInfo(partMetadata3, nodes);
+            Cluster cluster = new Cluster(null, asList(node1, node2), asList(part1, part2, part3),
+                Collections.emptySet(), Collections.emptySet());
+
             // Produce small record, we should switch to first partition.
             accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, value, Record.EMPTY_HEADERS,
                 callbacks, maxBlockTimeMs, false, time.milliseconds(), cluster);
@@ -1332,7 +1383,7 @@ public void testAdaptiveBuiltInPartitioner() throws Exception {
             }
 
             // Let the accumulator generate the probability tables.
-            accum.ready(metadataMock, time.milliseconds());
+            accum.ready(metadataCache, time.milliseconds());
 
             // Set up callbacks so that we know what partition is chosen.
             final AtomicInteger partition = new AtomicInteger(RecordMetadata.UNKNOWN_PARTITION);
@@ -1376,7 +1427,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
             // Test that partitions residing on high-latency nodes don't get switched to.
             accum.updateNodeLatencyStats(0, time.milliseconds() - 200, true);
             accum.updateNodeLatencyStats(0, time.milliseconds(), false);
-            accum.ready(metadataMock, time.milliseconds());
+            accum.ready(metadataCache, time.milliseconds());
 
             // Do one append, because partition gets switched after append.
             accum.append(topic, RecordMetadata.UNKNOWN_PARTITION, 0L, null, largeValue, Record.EMPTY_HEADERS,
@@ -1413,9 +1464,9 @@ public void testBuiltInPartitionerFractionalBatches() throws Exception {
             time.sleep(10);
 
             // We should have one batch ready.
-            Set<Node> nodes = accum.ready(metadataMock, time.milliseconds()).readyNodes;
+            Set<Node> nodes = accum.ready(metadataCache, time.milliseconds()).readyNodes;
             assertEquals(1, nodes.size(), "Should have 1 leader ready");
-            List<ProducerBatch> batches = accum.drain(metadataMock, nodes, Integer.MAX_VALUE, 0).entrySet().iterator().next().getValue();
+            List<ProducerBatch> batches = accum.drain(metadataCache, nodes, Integer.MAX_VALUE, 0).entrySet().iterator().next().getValue();
             assertEquals(1, batches.size(), "Should have 1 batch ready");
             int actualBatchSize = batches.get(0).records().sizeInBytes();
             assertTrue(actualBatchSize > batchSize / 2, "Batch must be greater than half batch.size");
@@ -1431,12 +1482,9 @@ public void testBuiltInPartitionerFractionalBatches() throws Exception {
     @Test
     public void testReadyAndDrainWhenABatchIsBeingRetried() throws InterruptedException {
         int part1LeaderEpoch = 100;
-        // Create cluster metadata, partition1 being hosted by node1.
-        part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-        cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1),
-            Collections.emptySet(), Collections.emptySet());
-        final int finalEpoch = part1LeaderEpoch;
-        metadataMock = setupMetadata(cluster, tp -> finalEpoch);
+        // Create cluster metadata, partition1 being hosted by node1
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
 
         int batchSize = 10;
         int lingerMs = 10;
@@ -1457,14 +1505,14 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         // 1st attempt(not a retry) to produce batchA, it should be ready & drained to be produced.
         {
             now += lingerMs + 1;
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
             assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
 
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                 result.readyNodes, 999999 /* maxSize */, now);
             assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
             ProducerBatch batch = batches.get(node1.id()).get(0);
-            assertEquals(Optional.of(part1LeaderEpoch), batch.currentLeaderEpoch());
+            assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
             assertEquals(0, batch.attemptsWhenLeaderLastChanged());
             // Re-enqueue batch for subsequent retries & test-cases
             accum.reenqueue(batch, now);
@@ -1473,11 +1521,11 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         // In this retry of batchA, wait-time between retries is less than configured and no leader change, so should backoff.
         {
             now += 1;
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
             assertFalse(result.readyNodes.contains(node1), "Node1 is not ready");
 
             // Try to drain from node1, it should return no batches.
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                 new HashSet<>(Arrays.asList(node1)), 999999 /* maxSize */, now);
             assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).isEmpty(),
                 "No batches ready to be drained on Node1");
@@ -1488,19 +1536,16 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
             now += 1;
             part1LeaderEpoch++;
             // Create cluster metadata, with new leader epoch.
-            part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-            cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1),
-                Collections.emptySet(), Collections.emptySet());
-            final int finalPart1LeaderEpoch = part1LeaderEpoch;
-            metadataMock = setupMetadata(cluster, tp -> finalPart1LeaderEpoch);
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+            part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
+            metadataCache = new MetadataSnapshot(null, nodes, Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
             assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
 
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                 result.readyNodes, 999999 /* maxSize */, now);
             assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
             ProducerBatch batch = batches.get(node1.id()).get(0);
-            assertEquals(Optional.of(part1LeaderEpoch), batch.currentLeaderEpoch());
+            assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
             assertEquals(1, batch.attemptsWhenLeaderLastChanged());
 
             // Re-enqueue batch for subsequent retries/test-cases.
@@ -1511,19 +1556,16 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
         {
             now += 2 * retryBackoffMaxMs;
             // Create cluster metadata, with new leader epoch.
-            part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-            cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1),
-                Collections.emptySet(), Collections.emptySet());
-            final int finalPart1LeaderEpoch = part1LeaderEpoch;
-            metadataMock = setupMetadata(cluster, tp -> finalPart1LeaderEpoch);
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+            part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
+            metadataCache = new MetadataSnapshot(null, nodes, Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
             assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
 
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                 result.readyNodes, 999999 /* maxSize */, now);
             assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
             ProducerBatch batch = batches.get(node1.id()).get(0);
-            assertEquals(Optional.of(part1LeaderEpoch), batch.currentLeaderEpoch());
+            assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
             assertEquals(1, batch.attemptsWhenLeaderLastChanged());
 
             // Re-enqueue batch for subsequent retries/test-cases.
@@ -1535,19 +1577,16 @@ deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
             now += 2 * retryBackoffMaxMs;
             part1LeaderEpoch++;
             // Create cluster metadata, with new leader epoch.
-            part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-            cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1),
-                Collections.emptySet(), Collections.emptySet());
-            final int finalPart1LeaderEpoch = part1LeaderEpoch;
-            metadataMock = setupMetadata(cluster, tp -> finalPart1LeaderEpoch);
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, now);
+            part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()),  Optional.of(part1LeaderEpoch), null, null, null);
+            metadataCache = new MetadataSnapshot(null, nodes, Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, now);
             assertTrue(result.readyNodes.contains(node1), "Node1 is ready");
 
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
                 result.readyNodes, 999999 /* maxSize */, now);
             assertTrue(batches.containsKey(node1.id()) && batches.get(node1.id()).size() == 1, "Node1 has 1 batch ready & drained");
             ProducerBatch batch = batches.get(node1.id()).get(0);
-            assertEquals(Optional.of(part1LeaderEpoch), batch.currentLeaderEpoch());
+            assertEquals(OptionalInt.of(part1LeaderEpoch), batch.currentLeaderEpoch());
             assertEquals(3, batch.attemptsWhenLeaderLastChanged());
 
             // Re-enqueue batch for subsequent retries/test-cases.
@@ -1564,97 +1603,15 @@ public void testDrainWithANodeThatDoesntHostAnyPartitions() {
             CompressionType.NONE, lingerMs);
 
         // Create cluster metadata, node2 doesn't host any partitions.
-        part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-        cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1),
-            Collections.emptySet(), Collections.emptySet());
-        metadataMock = Mockito.mock(Metadata.class);
-        Mockito.when(metadataMock.fetch()).thenReturn(cluster);
-        Mockito.when(metadataMock.currentLeader(tp1)).thenReturn(
-            new Metadata.LeaderAndEpoch(Optional.of(node1),
-                Optional.of(999 /* dummy value */)));
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
 
         // Drain for node2, it should return 0 batches,
-        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
+        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache,
             new HashSet<>(Arrays.asList(node2)), 999999 /* maxSize */, time.milliseconds());
         assertTrue(batches.get(node2.id()).isEmpty());
     }
 
-    @Test
-    public void testDrainOnANodeWhenItCeasesToBeALeader() throws InterruptedException {
-        int batchSize = 10;
-        int lingerMs = 10;
-        long totalSize = 10 * 1024;
-        RecordAccumulator accum = createTestRecordAccumulator(batchSize, totalSize,
-            CompressionType.NONE, lingerMs);
-
-        // While node1 is being drained, leader changes from node1 -> node2 for a partition.
-        {
-            // Create cluster metadata, partition1&2 being hosted by node1&2 resp.
-            part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-            part2 = new PartitionInfo(topic, partition2, node2, null, null, null);
-            cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2),
-                Collections.emptySet(), Collections.emptySet());
-            metadataMock = Mockito.mock(Metadata.class);
-            Mockito.when(metadataMock.fetch()).thenReturn(cluster);
-            // But metadata has a newer leader for partition1 i.e node2.
-            Mockito.when(metadataMock.currentLeader(tp1)).thenReturn(
-                new Metadata.LeaderAndEpoch(Optional.of(node2),
-                    Optional.of(999 /* dummy value */)));
-            Mockito.when(metadataMock.currentLeader(tp2)).thenReturn(
-                new Metadata.LeaderAndEpoch(Optional.of(node2),
-                    Optional.of(999 /* dummy value */)));
-
-            // Create 1 batch each for partition1 & partition2.
-            long now = time.milliseconds();
-            accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null,
-                maxBlockTimeMs, false, now, cluster);
-            accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null,
-                maxBlockTimeMs, false, now, cluster);
-
-            // Drain for node1, it should return 0 batches, as partition1's leader in metadata changed.
-            // Drain for node2, it should return 1 batch, for partition2.
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
-                new HashSet<>(Arrays.asList(node1, node2)), 999999 /* maxSize */, now);
-            assertTrue(batches.get(node1.id()).isEmpty());
-            assertEquals(1, batches.get(node2.id()).size());
-        }
-
-        // Cleanup un-drained batches to have an empty accum before next test.
-        accum.abortUndrainedBatches(new RuntimeException());
-
-        // While node1 is being drained, leader changes from node1 -> "no-leader" for partition.
-        {
-            // Create cluster metadata, partition1&2 being hosted by node1&2 resp.
-            part1 = new PartitionInfo(topic, partition1, node1, null, null, null);
-            part2 = new PartitionInfo(topic, partition2, node2, null, null, null);
-            cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2),
-                Collections.emptySet(), Collections.emptySet());
-            metadataMock = Mockito.mock(Metadata.class);
-            Mockito.when(metadataMock.fetch()).thenReturn(cluster);
-            // But metadata no longer has a leader for partition1.
-            Mockito.when(metadataMock.currentLeader(tp1)).thenReturn(
-                new Metadata.LeaderAndEpoch(Optional.empty(),
-                    Optional.of(999 /* dummy value */)));
-            Mockito.when(metadataMock.currentLeader(tp2)).thenReturn(
-                new Metadata.LeaderAndEpoch(Optional.of(node2),
-                    Optional.of(999 /* dummy value */)));
-
-            // Create 1 batch each for partition1 & partition2.
-            long now = time.milliseconds();
-            accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null,
-                maxBlockTimeMs, false, now, cluster);
-            accum.append(topic, partition2, 0L, key, value, Record.EMPTY_HEADERS, null,
-                maxBlockTimeMs, false, now, cluster);
-
-            // Drain for node1, it should return 0 batches, as partition1's leader in metadata changed.
-            // Drain for node2, it should return 1 batch, for partition2.
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock,
-                new HashSet<>(Arrays.asList(node1, node2)), 999999 /* maxSize */, now);
-            assertTrue(batches.get(node1.id()).isEmpty());
-            assertEquals(1, batches.get(node2.id()).size());
-        }
-    }
-
     private int prepareSplitBatches(RecordAccumulator accum, long seed, int recordSize, int numRecords)
         throws InterruptedException {
         Random random = new Random();
@@ -1667,9 +1624,9 @@ private int prepareSplitBatches(RecordAccumulator accum, long seed, int recordSi
             accum.append(topic, partition1, 0L, null, bytesWithPoorCompression(random, recordSize), Record.EMPTY_HEADERS, null, 0, false, time.milliseconds(), cluster);
         }
 
-        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
+        RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
-        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
         assertEquals(1, batches.size());
         assertEquals(1, batches.values().iterator().next().size());
         ProducerBatch batch = batches.values().iterator().next().get(0);
@@ -1685,8 +1642,8 @@ private BatchDrainedResult completeOrSplitBatches(RecordAccumulator accum, int b
         boolean batchDrained;
         do {
             batchDrained = false;
-            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataMock, time.milliseconds());
-            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataMock, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
+            RecordAccumulator.ReadyCheckResult result = accum.ready(metadataCache, time.milliseconds());
+            Map<Integer, List<ProducerBatch>> batches = accum.drain(metadataCache, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
             for (List<ProducerBatch> batchList : batches.values()) {
                 for (ProducerBatch batch : batchList) {
                     batchDrained = true;
@@ -1805,27 +1762,4 @@ private RecordAccumulator createTestRecordAccumulator(
             txnManager,
             new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
     }
-
-    /**
-     * Setup a mocked metadata object.
-     */
-    private Metadata setupMetadata(Cluster cluster) {
-        return setupMetadata(cluster, tp -> 999 /* dummy epoch */);
-    }
-
-    /**
-     * Setup a mocked metadata object.
-     */
-    private Metadata setupMetadata(Cluster cluster, final Function<TopicPartition, Integer> epochSupplier) {
-        Metadata metadataMock = Mockito.mock(Metadata.class);
-        Mockito.when(metadataMock.fetch()).thenReturn(cluster);
-        for (String topic: cluster.topics()) {
-            for (PartitionInfo partInfo: cluster.partitionsForTopic(topic)) {
-                TopicPartition tp = new TopicPartition(partInfo.topic(), partInfo.partition());
-                Integer partLeaderEpoch = epochSupplier.apply(tp);
-                Mockito.when(metadataMock.currentLeader(tp)).thenReturn(new Metadata.LeaderAndEpoch(Optional.of(partInfo.leader()), Optional.of(partLeaderEpoch)));
-            }
-        }
-        return metadataMock;
-    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
index a524f2f78dc91..9af3aff8bbf3f 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
@@ -20,6 +20,7 @@
 import org.apache.kafka.clients.ClientRequest;
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.NodeApiVersions;
@@ -471,7 +472,7 @@ public void testAppendInExpiryCallback() throws InterruptedException {
         final byte[] key = "key".getBytes();
         final byte[] value = "value".getBytes();
         final long maxBlockTimeMs = 1000;
-        Cluster cluster = TestUtils.singletonCluster();
+        MetadataSnapshot metadataCache = TestUtils.metadataSnapshotWith(1);
         RecordAccumulator.AppendCallbacks callbacks = new RecordAccumulator.AppendCallbacks() {
             @Override
             public void setPartition(int partition) {
@@ -483,7 +484,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
                     expiryCallbackCount.incrementAndGet();
                     try {
                         accumulator.append(tp1.topic(), tp1.partition(), 0L, key, value,
-                            Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), cluster);
+                            Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds(), metadataCache.cluster());
                     } catch (InterruptedException e) {
                         throw new RuntimeException("Unexpected interruption", e);
                     }
@@ -494,14 +495,14 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
 
         final long nowMs = time.milliseconds();
         for (int i = 0; i < messagesPerBatch; i++)
-            accumulator.append(tp1.topic(), tp1.partition(), 0L, key, value, null, callbacks, maxBlockTimeMs, false, nowMs, cluster);
+            accumulator.append(tp1.topic(), tp1.partition(), 0L, key, value, null, callbacks, maxBlockTimeMs, false, nowMs, metadataCache.cluster());
 
         // Advance the clock to expire the first batch.
         time.sleep(10000);
 
         Node clusterNode = metadata.fetch().nodes().get(0);
         Map<Integer, List<ProducerBatch>> drainedBatches =
-            accumulator.drain(metadata, Collections.singleton(clusterNode), Integer.MAX_VALUE, time.milliseconds());
+            accumulator.drain(metadataCache, Collections.singleton(clusterNode), Integer.MAX_VALUE, time.milliseconds());
         sender.addToInflightBatches(drainedBatches);
 
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index 0e9bd834c616c..5a2527deee99e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.consumer.CommitFailedException;
@@ -69,6 +70,7 @@
 import org.apache.kafka.common.requests.InitProducerIdRequest;
 import org.apache.kafka.common.requests.InitProducerIdResponse;
 import org.apache.kafka.common.requests.JoinGroupRequest;
+import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata;
 import org.apache.kafka.common.requests.ProduceRequest;
 import org.apache.kafka.common.requests.ProduceResponse;
 import org.apache.kafka.common.requests.RequestTestUtils;
@@ -2476,16 +2478,16 @@ public void testNoDrainWhenPartitionsPending() throws InterruptedException {
 
         Node node1 = new Node(0, "localhost", 1111);
         Node node2 = new Node(1, "localhost", 1112);
-        PartitionInfo part1 = new PartitionInfo(topic, 0, node1, null, null);
-        PartitionInfo part2 = new PartitionInfo(topic, 1, node2, null, null);
-
-        Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2),
-                Collections.emptySet(), Collections.emptySet());
-        Metadata metadataMock = setupMetadata(cluster);
+        Map<Integer, Node> nodesById = new HashMap<>();
+        nodesById.put(node1.id(), node1);
+        nodesById.put(node2.id(), node2);
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp0, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        PartitionMetadata part2Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node2.id()), Optional.empty(), null, null, null);
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodesById, Arrays.asList(part1Metadata, part2Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
         Set<Node> nodes = new HashSet<>();
         nodes.add(node1);
         nodes.add(node2);
-        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataMock, nodes, Integer.MAX_VALUE,
+        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataCache, nodes, Integer.MAX_VALUE,
                 time.milliseconds());
 
         // We shouldn't drain batches which haven't been added to the transaction yet.
@@ -2511,12 +2513,10 @@ public void testAllowDrainInAbortableErrorState() throws InterruptedException {
 
         // Try to drain a message destined for tp1, it should get drained.
         Node node1 = new Node(1, "localhost", 1112);
-        PartitionInfo part1 = new PartitionInfo(topic, 1, node1, null, null);
-        Cluster cluster = new Cluster(null, Collections.singletonList(node1), Collections.singletonList(part1),
-                Collections.emptySet(), Collections.emptySet());
-        Metadata metadataMock = setupMetadata(cluster);
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null, Collections.singletonMap(node1.id(), node1), Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
         appendToAccumulator(tp1);
-        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataMock, Collections.singleton(node1),
+        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataCache, Collections.singleton(node1),
                 Integer.MAX_VALUE,
                 time.milliseconds());
 
@@ -2534,15 +2534,12 @@ public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedExc
         // Don't execute transactionManager.maybeAddPartitionToTransaction(tp0). This should result in an error on drain.
         appendToAccumulator(tp0);
         Node node1 = new Node(0, "localhost", 1111);
-        PartitionInfo part1 = new PartitionInfo(topic, 0, node1, null, null);
-
-        Cluster cluster = new Cluster(null, Collections.singletonList(node1), Collections.singletonList(part1),
-                Collections.emptySet(), Collections.emptySet());
-        Metadata metadataMock = setupMetadata(cluster);
+        PartitionMetadata part1Metadata = new PartitionMetadata(Errors.NONE, tp0, Optional.of(node1.id()), Optional.empty(), null, null, null);
+        MetadataSnapshot metadataCache = new MetadataSnapshot(null, Collections.singletonMap(node1.id(), node1), Arrays.asList(part1Metadata), Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
 
         Set<Node> nodes = new HashSet<>();
         nodes.add(node1);
-        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataMock, nodes, Integer.MAX_VALUE,
+        Map<Integer, List<ProducerBatch>> drainedBatches = accumulator.drain(metadataCache, nodes, Integer.MAX_VALUE,
                 time.milliseconds());
 
         // We shouldn't drain batches which haven't been added to the transaction yet.
diff --git a/clients/src/test/java/org/apache/kafka/test/TestUtils.java b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
index c8a6db6f6ca32..db86558f7d0b3 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.test;
 
+import org.apache.kafka.clients.MetadataSnapshot;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.Cluster;
@@ -29,10 +30,12 @@
 import org.apache.kafka.common.network.NetworkReceive;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.RecordVersion;
 import org.apache.kafka.common.record.UnalignedRecords;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
 import org.apache.kafka.common.requests.ByteBufferChannel;
+import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
@@ -123,6 +126,42 @@ public static Cluster clusterWith(final int nodes, final String topic, final int
         return clusterWith(nodes, Collections.singletonMap(topic, partitions));
     }
 
+    /**
+     * Test utility function to get MetadataSnapshot with configured nodes and partitions.
+     * @param nodes number of nodes in the cluster
+     * @param topicPartitionCounts map of topic -> # of partitions
+     * @return a MetadataSnapshot with number of nodes, partitions as per the input.
+     */
+
+    public static MetadataSnapshot metadataSnapshotWith(final int nodes, final Map<String, Integer> topicPartitionCounts) {
+        final Node[] ns = new Node[nodes];
+        Map<Integer, Node> nodesById = new HashMap<>();
+        for (int i = 0; i < nodes; i++) {
+            ns[i] = new Node(i, "localhost", 1969);
+            nodesById.put(ns[i].id(), ns[i]);
+        }
+        final List<PartitionMetadata> partsMetadatas = new ArrayList<>();
+        for (final Map.Entry<String, Integer> topicPartition : topicPartitionCounts.entrySet()) {
+            final String topic = topicPartition.getKey();
+            final int partitions = topicPartition.getValue();
+            for (int i = 0; i < partitions; i++) {
+                TopicPartition tp = new TopicPartition(topic, partitions);
+                Node node = ns[i % ns.length];
+                partsMetadatas.add(new PartitionMetadata(Errors.NONE, tp, Optional.of(node.id()), Optional.empty(), null, null, null));
+            }
+        }
+        return new MetadataSnapshot("kafka-cluster", nodesById, partsMetadatas, Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap());
+    }
+
+    /**
+     * Test utility function to get MetadataSnapshot of cluster with configured, and 0 partitions.
+     * @param nodes number of nodes in the cluster.
+     * @return a MetadataSnapshot of cluster with number of nodes in the input.
+     */
+    public static MetadataSnapshot metadataSnapshotWith(int nodes) {
+        return metadataSnapshotWith(nodes, new HashMap<>());
+    }
+
     /**
      * Generate an array of random bytes
      *

From 553f45bca8a936d971af03e1552f2fe3f9c54fac Mon Sep 17 00:00:00 2001
From: "Minha, Jeong" <vwjdalsgkv@gmail.com>
Date: Fri, 16 Feb 2024 12:07:18 +0900
Subject: [PATCH 041/258] MINOR: Fix toString method of IsolationLevel (#14782)

Reviewers: Matthias J. Sax <matthias@confluent.io>, Ashwin Pankaj <apankaj@confluent.io>
---
 .../apache/kafka/clients/consumer/ConsumerConfig.java |  4 ++--
 .../java/org/apache/kafka/common/IsolationLevel.java  |  7 +++++++
 .../kafka/clients/consumer/KafkaConsumerTest.java     |  3 +--
 .../kafka/connect/mirror/MirrorSourceConnector.java   |  3 +--
 .../java/org/apache/kafka/connect/runtime/Worker.java |  5 ++---
 .../connect/storage/KafkaConfigBackingStore.java      |  3 +--
 .../connect/storage/KafkaOffsetBackingStore.java      |  5 ++---
 .../org/apache/kafka/connect/util/KafkaBasedLog.java  |  5 ++---
 .../connect/storage/KafkaConfigBackingStoreTest.java  | 11 +++++------
 .../connect/storage/KafkaOffsetBackingStoreTest.java  | 11 +++++------
 .../java/org/apache/kafka/streams/StreamsConfig.java  |  3 +--
 .../org/apache/kafka/streams/StreamsConfigTest.java   |  8 ++++----
 .../kafka/streams/integration/EosIntegrationTest.java |  5 ++---
 .../integration/EosV2UpgradeIntegrationTest.java      |  3 +--
 .../kafka/streams/tests/BrokerCompatibilityTest.java  |  3 +--
 .../org/apache/kafka/streams/tests/EosTestDriver.java |  5 ++---
 16 files changed, 39 insertions(+), 45 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerConfig.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerConfig.java
index 512f8c4f06ca3..7ec147e9e3ced 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerConfig.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerConfig.java
@@ -363,7 +363,7 @@ public class ConsumerConfig extends AbstractConfig {
             " consumers will not be able to read up to the high watermark when there are in flight transactions.</p><p> Further, when in <code>read_committed</code> the seekToEnd method will" +
             " return the LSO</p>";
 
-    public static final String DEFAULT_ISOLATION_LEVEL = IsolationLevel.READ_UNCOMMITTED.toString().toLowerCase(Locale.ROOT);
+    public static final String DEFAULT_ISOLATION_LEVEL = IsolationLevel.READ_UNCOMMITTED.toString();
 
     /** <code>allow.auto.create.topics</code> */
     public static final String ALLOW_AUTO_CREATE_TOPICS_CONFIG = "allow.auto.create.topics";
@@ -620,7 +620,7 @@ public class ConsumerConfig extends AbstractConfig {
                                 .define(ISOLATION_LEVEL_CONFIG,
                                         Type.STRING,
                                         DEFAULT_ISOLATION_LEVEL,
-                                        in(IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT), IsolationLevel.READ_UNCOMMITTED.toString().toLowerCase(Locale.ROOT)),
+                                        in(IsolationLevel.READ_COMMITTED.toString(), IsolationLevel.READ_UNCOMMITTED.toString()),
                                         Importance.MEDIUM,
                                         ISOLATION_LEVEL_DOC)
                                 .define(ALLOW_AUTO_CREATE_TOPICS_CONFIG,
diff --git a/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java b/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
index 79f0a92954bf1..fd4f45f39642c 100644
--- a/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
+++ b/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common;
 
+import java.util.Locale;
+
 public enum IsolationLevel {
     READ_UNCOMMITTED((byte) 0), READ_COMMITTED((byte) 1);
 
@@ -39,4 +41,9 @@ public static IsolationLevel forId(byte id) {
                 throw new IllegalArgumentException("Unknown isolation level " + id);
         }
     }
+
+    @Override
+    public String toString() {
+        return super.toString().toLowerCase(Locale.ROOT);
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
index dc5b5ff2f9069..7dec7305d4db4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
@@ -116,7 +116,6 @@
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -2959,7 +2958,7 @@ private ConsumerConfig newConsumerConfig(GroupProtocol groupProtocol,
         configs.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
         configs.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
         configs.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, heartbeatIntervalMs);
-        configs.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        configs.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
         configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
         configs.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, fetchSize);
         configs.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, rebalanceTimeoutMs);
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java
index 700ba3d139fa6..45c01527a4d36 100644
--- a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.connect.mirror;
 
 import java.util.HashMap;
-import java.util.Locale;
 import java.util.Map.Entry;
 
 import java.util.concurrent.atomic.AtomicReference;
@@ -88,7 +87,7 @@ public class MirrorSourceConnector extends SourceConnector {
     private static final ResourcePatternFilter ANY_TOPIC = new ResourcePatternFilter(ResourceType.TOPIC,
         null, PatternType.ANY);
     private static final AclBindingFilter ANY_TOPIC_ACL = new AclBindingFilter(ANY_TOPIC, AccessControlEntryFilter.ANY);
-    private static final String READ_COMMITTED = IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT);
+    private static final String READ_COMMITTED = IsolationLevel.READ_COMMITTED.toString();
     private static final String EXACTLY_ONCE_SUPPORT_CONFIG = "exactly.once.support";
 
     private final AtomicBoolean noAclAuthorizer = new AtomicBoolean(false);
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java
index 81ab861724715..2ce09ee28b6df 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java
@@ -107,7 +107,6 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -855,7 +854,7 @@ static Map<String, Object> exactlyOnceSourceOffsetsConsumerConfigs(String connNa
                 connName, defaultClientId, config, connConfig, connectorClass,
                 connectorClientConfigOverridePolicy, clusterId, ConnectorType.SOURCE);
         ConnectUtils.ensureProperty(
-                result, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                result, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString(),
                 "for source connectors' offset consumers when exactly-once source support is enabled",
                 false
         );
@@ -875,7 +874,7 @@ static Map<String, Object> regularSourceOffsetsConsumerConfigs(String connName,
         // Users can disable this if they want to since the task isn't exactly-once anyways
         result.putIfAbsent(
                 ConsumerConfig.ISOLATION_LEVEL_CONFIG,
-                IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT));
+                IsolationLevel.READ_COMMITTED.toString());
         return result;
     }
 
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
index 35d43ea3ccaed..735020173996c 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
@@ -67,7 +67,6 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -781,7 +780,7 @@ KafkaBasedLog<String, byte[]> setupAndCreateKafkaBasedLog(String topic, final Wo
         ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
         if (config.exactlyOnceSourceEnabled()) {
             ConnectUtils.ensureProperty(
-                    consumerProps, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                    consumerProps, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString(),
                     "for the worker's config topic consumer when exactly-once source support is enabled",
                     true
             );
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java
index b5f8c51d83760..786bd55351fcb 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java
@@ -47,7 +47,6 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -203,7 +202,7 @@ public void configure(final WorkerConfig config) {
         ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
         if (config.exactlyOnceSourceEnabled()) {
             ConnectUtils.ensureProperty(
-                    consumerProps, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                    consumerProps, ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString(),
                     "for the worker offsets topic consumer when exactly-once source support is enabled",
                     false
             );
@@ -250,7 +249,7 @@ public void start() {
                         + "support for source connectors, or upgrade to a newer Kafka broker version.";
             } else {
                 message = "When " + ConsumerConfig.ISOLATION_LEVEL_CONFIG + "is set to "
-                        + IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT)
+                        + IsolationLevel.READ_COMMITTED.toString()
                         + ", a Kafka broker version that allows admin clients to read consumer offsets is required. "
                         + "Please either reconfigure the worker or connector, or upgrade to a newer Kafka broker version.";
             }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/KafkaBasedLog.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/KafkaBasedLog.java
index 68e117e71f865..a66ac4e0546d9 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/KafkaBasedLog.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/KafkaBasedLog.java
@@ -46,7 +46,6 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -179,7 +178,7 @@ public KafkaBasedLog(String topic,
         // as it will not take records from currently-open transactions into account. We want to err on the side of caution in that
         // case: when users request a read to the end of the log, we will read up to the point where the latest offsets visible to the
         // consumer are at least as high as the (possibly-part-of-a-transaction) end offsets of the topic.
-        this.requireAdminForOffsets = IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT)
+        this.requireAdminForOffsets = IsolationLevel.READ_COMMITTED.toString()
                 .equals(consumerConfigs.get(ConsumerConfig.ISOLATION_LEVEL_CONFIG));
     }
 
@@ -252,7 +251,7 @@ public void start() {
             throw new ConnectException(
                     "Must provide a TopicAdmin to KafkaBasedLog when consumer is configured with "
                             + ConsumerConfig.ISOLATION_LEVEL_CONFIG + " set to "
-                            + IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT)
+                            + IsolationLevel.READ_COMMITTED.toString()
             );
         }
         initializer.accept(admin);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
index a03224149511f..39068a990459f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
@@ -65,7 +65,6 @@
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -1643,7 +1642,7 @@ public void testConsumerPropertiesInsertedByDefaultWithExactlyOnceSourceEnabled(
         configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
 
         assertEquals(
-                IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_COMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
 
@@ -1653,7 +1652,7 @@ public void testConsumerPropertiesInsertedByDefaultWithExactlyOnceSourceEnabled(
     @Test
     public void testConsumerPropertiesOverrideUserSuppliedValuesWithExactlyOnceSourceEnabled() throws Exception {
         props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "enabled");
-        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
         createStore();
 
         expectConfigure();
@@ -1662,7 +1661,7 @@ public void testConsumerPropertiesOverrideUserSuppliedValuesWithExactlyOnceSourc
         configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
 
         assertEquals(
-                IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_COMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
 
@@ -1688,7 +1687,7 @@ public void testConsumerPropertiesNotInsertedByDefaultWithoutExactlyOnceSourceEn
     @Test
     public void testConsumerPropertiesDoNotOverrideUserSuppliedValuesWithoutExactlyOnceSourceEnabled() throws Exception {
         props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "preparing");
-        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
         createStore();
 
         expectConfigure();
@@ -1697,7 +1696,7 @@ public void testConsumerPropertiesDoNotOverrideUserSuppliedValuesWithoutExactlyO
         configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
 
         assertEquals(
-                IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_UNCOMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
index 2aaf97d326318..78f995253b337 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
@@ -46,7 +46,6 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -420,7 +419,7 @@ public void testConsumerPropertiesInsertedByDefaultWithExactlyOnceSourceEnabled(
         store.configure(mockConfig(props));
 
         assertEquals(
-                IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_COMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
     }
@@ -428,12 +427,12 @@ public void testConsumerPropertiesInsertedByDefaultWithExactlyOnceSourceEnabled(
     @Test
     public void testConsumerPropertiesOverrideUserSuppliedValuesWithExactlyOnceSourceEnabled() {
         props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "enabled");
-        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
 
         store.configure(mockConfig(props));
 
         assertEquals(
-                IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_COMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
     }
@@ -451,12 +450,12 @@ public void testConsumerPropertiesNotInsertedByDefaultWithoutExactlyOnceSourceEn
     @Test
     public void testConsumerPropertiesDoNotOverrideUserSuppliedValuesWithoutExactlyOnceSourceEnabled() {
         props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "disabled");
-        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
 
         store.configure(mockConfig(props));
 
         assertEquals(
-                IsolationLevel.READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT),
+                IsolationLevel.READ_UNCOMMITTED.toString(),
                 capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
         );
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 70520c1eee26e..2f4841269b5dd 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -57,7 +57,6 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Properties;
@@ -1208,7 +1207,7 @@ public class StreamsConfig extends AbstractConfig {
     private static final Map<String, Object> CONSUMER_EOS_OVERRIDES;
     static {
         final Map<String, Object> tempConsumerDefaultOverrides = new HashMap<>(CONSUMER_DEFAULT_OVERRIDES);
-        tempConsumerDefaultOverrides.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, READ_COMMITTED.name().toLowerCase(Locale.ROOT));
+        tempConsumerDefaultOverrides.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, READ_COMMITTED.toString());
         CONSUMER_EOS_OVERRIDES = Collections.unmodifiableMap(tempConsumerDefaultOverrides);
     }
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java b/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
index 33c8f1c2eb3bc..3f1bfef63194e 100644
--- a/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
@@ -674,18 +674,18 @@ private void shouldResetToDefaultIfConsumerIsolationLevelIsOverriddenIfEosEnable
         final Map<String, Object> consumerConfigs = streamsConfig.getMainConsumerConfigs(groupId, clientId, threadIdx);
         assertThat(
             consumerConfigs.get(ConsumerConfig.ISOLATION_LEVEL_CONFIG),
-            equalTo(READ_COMMITTED.name().toLowerCase(Locale.ROOT))
+            equalTo(READ_COMMITTED.toString())
         );
     }
 
     @Test
     public void shouldAllowSettingConsumerIsolationLevelIfEosDisabled() {
-        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT));
+        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, READ_UNCOMMITTED.toString());
         final StreamsConfig streamsConfig = new StreamsConfig(props);
         final Map<String, Object> consumerConfigs = streamsConfig.getMainConsumerConfigs(groupId, clientId, threadIdx);
         assertThat(
             consumerConfigs.get(ConsumerConfig.ISOLATION_LEVEL_CONFIG),
-            equalTo(READ_UNCOMMITTED.name().toLowerCase(Locale.ROOT))
+            equalTo(READ_UNCOMMITTED.toString())
         );
     }
 
@@ -752,7 +752,7 @@ private void shouldSetDifferentDefaultsIfEosEnabled() {
 
         assertThat(
             consumerConfigs.get(ConsumerConfig.ISOLATION_LEVEL_CONFIG),
-            equalTo(READ_COMMITTED.name().toLowerCase(Locale.ROOT))
+            equalTo(READ_COMMITTED.toString())
         );
         assertTrue((Boolean) producerConfigs.get(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG));
         assertThat(producerConfigs.get(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG), equalTo(Integer.MAX_VALUE));
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
index d06613dd9a87b..d79631f8f7a0c 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
@@ -88,7 +88,6 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
@@ -1171,7 +1170,7 @@ private <K, V> List<KeyValue<K, V>> readResult(final String topic,
                     valueDeserializer,
                     Utils.mkProperties(Collections.singletonMap(
                         ConsumerConfig.ISOLATION_LEVEL_CONFIG,
-                        IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT)))),
+                        IsolationLevel.READ_COMMITTED.toString()))),
                 topic,
                 numberOfRecords
             );
@@ -1203,7 +1202,7 @@ private <K, V> void ensureCommittedRecordsInTopicPartition(final String topic,
                     valueDeserializer,
                     Utils.mkProperties(Collections.singletonMap(
                         ConsumerConfig.ISOLATION_LEVEL_CONFIG,
-                        IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT))
+                        IsolationLevel.READ_COMMITTED.toString())
                     )
                 ),
                 topic,
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
index 5d8e4bf319036..011090c152667 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
@@ -75,7 +75,6 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
@@ -1067,7 +1066,7 @@ private List<KeyValue<Long, Long>> readResult(final int numberOfRecords,
                     LongDeserializer.class,
                     Utils.mkProperties(Collections.singletonMap(
                         ConsumerConfig.ISOLATION_LEVEL_CONFIG,
-                        IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT))
+                        IsolationLevel.READ_COMMITTED.toString())
                     )
                 ),
                 MULTI_PARTITION_OUTPUT_TOPIC,
diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/BrokerCompatibilityTest.java b/streams/src/test/java/org/apache/kafka/streams/tests/BrokerCompatibilityTest.java
index 1e06b4ab31415..7f8d144be45fb 100644
--- a/streams/src/test/java/org/apache/kafka/streams/tests/BrokerCompatibilityTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/tests/BrokerCompatibilityTest.java
@@ -40,7 +40,6 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Collections;
-import java.util.Locale;
 import java.util.Properties;
 
 public class BrokerCompatibilityTest {
@@ -146,7 +145,7 @@ private static void loopUntilRecordReceived(final String kafka, final boolean eo
         consumerProperties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
         consumerProperties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
         if (eosEnabled) {
-            consumerProperties.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.name().toLowerCase(Locale.ROOT));
+            consumerProperties.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString());
         }
 
         try (final KafkaConsumer<String, String> consumer = new KafkaConsumer<>(consumerProperties)) {
diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/EosTestDriver.java b/streams/src/test/java/org/apache/kafka/streams/tests/EosTestDriver.java
index 18822d3b08ba6..77d0fc3ad6c12 100644
--- a/streams/src/test/java/org/apache/kafka/streams/tests/EosTestDriver.java
+++ b/streams/src/test/java/org/apache/kafka/streams/tests/EosTestDriver.java
@@ -47,7 +47,6 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Random;
@@ -152,7 +151,7 @@ static void generate(final String kafka) {
             props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka);
             props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
             props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
-            props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT));
+            props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString());
 
             try (final KafkaConsumer<byte[], byte[]> consumer = new KafkaConsumer<>(props)) {
                 final List<TopicPartition> partitions = getAllPartitions(consumer, "data");
@@ -178,7 +177,7 @@ public static void verify(final String kafka, final boolean withRepartitioning)
         props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka);
         props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
         props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
-        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString().toLowerCase(Locale.ROOT));
+        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_COMMITTED.toString());
 
         try (final KafkaConsumer<byte[], byte[]> consumer = new KafkaConsumer<>(props)) {
             verifyAllTransactionFinished(consumer, kafka, withRepartitioning);

From 051d4274da6b07f1edc1d4d7556b12c9bb2a7f79 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 16 Feb 2024 00:05:02 -0800
Subject: [PATCH 042/258] KAFKA-16167: re-enable
 PlaintextConsumerTest.testAutoCommitOnCloseAfterWakeup (#15358)

This integration test is now passing, presumably based on recent related changes. Re-enabling to ensure it is included in the test suite to catch any regressions.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../scala/integration/kafka/api/PlaintextConsumerTest.scala   | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 74009fb6e0ac3..fbbd7c42a9c5e 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -289,10 +289,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
-  // TODO: Enable this test for both protocols when the Jira tracking its failure (KAFKA-16167) is fixed. This
-  //       is done by setting the @MethodSource value to "getTestQuorumAndGroupProtocolParametersAll"
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testAutoCommitOnCloseAfterWakeup(quorum: String, groupProtocol: String): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
     val consumer = createConsumer()

From 14d5e170706be0b843d08147a1046fafb442f9be Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Fri, 16 Feb 2024 05:38:46 -0500
Subject: [PATCH 043/258] KAFKA-16165: Fix invalid transition on poll timer
 expiration (#15375)

This fixes an invalid transition (leaving->stale) that was discovered in the system tests. The underlying issue was that the poll timer expiration logic was blindly forcing a transition to stale and sending a leave group, without considering that the member could be already leaving.
The fix included in this PR ensures that the poll timer expiration logic, whose purpose is to leave the group, is only applied if the member is not already leaving. Note that it also fixes the transition out of the STALE state, that should only happen when the poll timer is reset.

As a result of this changes:

If the poll timer expires while the member is not leaving, the poll timer expiration logic is applied: it will transition to stale, send a leave group, and remain in STALE state until the timer is reset. At that point the member will transition to JOINING to rejoin the group.
If the poll timer expires while the member is already leaving, the poll timer expiration logic does not apply, and just lets the HB continue. Not that this would be the case of member in PREPARE_LEAVING waiting for callbacks to complete (needs to continue sending HB), or LEAVING (needs to send the last HB to leave).

Reviewers: Kirk True <ktrue@confluent.io>, Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/HeartbeatRequestManager.java    | 24 +++---
 .../consumer/internals/MemberState.java       |  9 ++-
 .../consumer/internals/MembershipManager.java | 11 +--
 .../internals/MembershipManagerImpl.java      | 35 +++++----
 .../HeartbeatRequestManagerTest.java          | 31 +++++++-
 .../internals/MembershipManagerImplTest.java  | 77 +++++++++++++++++--
 6 files changed, 146 insertions(+), 41 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 246ea05b220d8..a6b0b62c4ebe2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -188,18 +188,18 @@ public HeartbeatRequestManager(
     @Override
     public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
         if (!coordinatorRequestManager.coordinator().isPresent() ||
-            membershipManager.shouldSkipHeartbeat() ||
-            pollTimer.isExpired()) {
+            membershipManager.shouldSkipHeartbeat()) {
             membershipManager.onHeartbeatRequestSkipped();
             return NetworkClientDelegate.PollResult.EMPTY;
         }
         pollTimer.update(currentTimeMs);
-        if (pollTimer.isExpired()) {
-            logger.warn("consumer poll timeout has expired. This means the time between subsequent calls to poll() " +
-                "was longer than the configured max.poll.interval.ms, which typically implies that " +
-                "the poll loop is spending too much time processing messages. You can address this " +
-                "either by increasing max.poll.interval.ms or by reducing the maximum size of batches " +
-                "returned in poll() with max.poll.records.");
+        if (pollTimer.isExpired() && !membershipManager.isLeavingGroup()) {
+            logger.warn("Consumer poll timeout has expired. This means the time between " +
+                "subsequent calls to poll() was longer than the configured max.poll.interval.ms, " +
+                "which typically implies that the poll loop is spending too much time processing " +
+                "messages. You can address this either by increasing max.poll.interval.ms or by " +
+                "reducing the maximum size of batches returned in poll() with max.poll.records.");
+
             // This should trigger a heartbeat with leave group epoch
             membershipManager.transitionToStale();
             NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(currentTimeMs, true);
@@ -242,12 +242,16 @@ public long maximumTimeToWait(long currentTimeMs) {
     }
 
     /**
-     * When consumer polls, we need to reset the pollTimer.  If the poll timer has expired, we rejoin when the user
-     * repoll the consumer.
+     * Reset the poll timer, indicating that the user has called consumer.poll(). If the member
+     * is in {@link MemberState#STALE} state due to expired poll timer, this will transition the
+     * member to {@link MemberState#JOINING}, so that it rejoins the group.
      */
     public void resetPollTimer(final long pollMs) {
         pollTimer.update(pollMs);
         pollTimer.reset(maxPollIntervalMs);
+        if (membershipManager.state() == MemberState.STALE) {
+            membershipManager.transitionToJoining();
+        }
     }
 
     private NetworkClientDelegate.UnsentRequest makeHeartbeatRequest(final long currentTimeMs,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
index 100795d55f937..1df4d30e5945d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
@@ -103,9 +103,12 @@ public enum MemberState {
     FATAL,
 
     /**
-     * An intermediate state indicating the consumer is staled because the user has not polled the consumer
-     * within the <code>max.poll.interval.ms</code> time bound; therefore causing the member to leave the
-     * group. The member rejoins on the next poll.
+     * The member transitions to this state when the poll timer expires, indicating that there
+     * hasn't been a call to consumer.poll within the <code>max.poll.interval.ms</code>. While in
+     * this state, the member will send a heartbeat to leave the group, invoke the
+     * onPartitionsLost callback, and clear its assignments. The member will only transition
+     * out of this state on the next application poll event. The member will then transition
+     * to JOINING, to rejoin the group.
      */
     STALE;
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
index 7f0975b5aa968..f95552cc3a353 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
@@ -63,11 +63,6 @@ public interface MembershipManager extends RequestManager {
      */
     MemberState state();
 
-    /**
-     * @return True if the member is staled due to expired poll timer.
-     */
-    boolean isStaled();
-
     /**
      * Update member info and transition member state based on a successful heartbeat response.
      *
@@ -174,4 +169,10 @@ public interface MembershipManager extends RequestManager {
      * transitions of new data received from the server, as defined in {@link MemberStateListener}.
      */
     void registerStateListener(MemberStateListener listener);
+
+    /**
+     * @return True if the member is preparing to leave the group (waiting for callbacks), or
+     * leaving (sending last heartbeat).
+     */
+    boolean isLeavingGroup();
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 7132b43a869b6..6556495e55b34 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -335,11 +335,6 @@ public int memberEpoch() {
         return memberEpoch;
     }
 
-    @Override
-    public boolean isStaled() {
-        return state == MemberState.STALE;
-    }
-
     /**
      * {@inheritDoc}
      */
@@ -685,13 +680,6 @@ public boolean shouldHeartbeatNow() {
     @Override
     public void onHeartbeatRequestSent() {
         MemberState state = state();
-        if (isStaled()) {
-            log.debug("Member {} is staled and is therefore leaving the group.  It will rejoin upon the next poll.", memberEpoch);
-            // TODO: Integrate partition revocation/loss callback
-            transitionToJoining();
-            return;
-        }
-
         if (state == MemberState.ACKNOWLEDGING) {
             if (targetAssignmentReconciled()) {
                 transitionTo(MemberState.STABLE);
@@ -731,15 +719,34 @@ private boolean targetAssignmentReconciled() {
         return currentAssignment.equals(currentTargetAssignment);
     }
 
+    /**
+     * @return True if the member should not send heartbeats, which would be one of the following
+     * cases:
+     * <ul>
+     * <li>Member is not subscribed to any topics</li>
+     * <li>Member has received a fatal error in a previous heartbeat response</li>
+     * <li>Member is stale, meaning that it has left the group due to expired poll timer</li>
+     * </ul>
+     */
     @Override
     public boolean shouldSkipHeartbeat() {
         MemberState state = state();
-        return state == MemberState.UNSUBSCRIBED || state == MemberState.FATAL;
+        return state == MemberState.UNSUBSCRIBED || state == MemberState.FATAL || state == MemberState.STALE;
+    }
+
+    /**
+     * @return True if the member is preparing to leave the group (waiting for callbacks), or
+     * leaving (sending last heartbeat). This is used to skip proactively leaving the group when
+     * the consumer poll timer expires.
+     */
+    public boolean isLeavingGroup() {
+        MemberState state = state();
+        return state == MemberState.PREPARE_LEAVING || state == MemberState.LEAVING;
     }
 
     /**
      * Sets the epoch to the leave group epoch and clears the assignments. The member will rejoin with
-     * the existing subscriptions on the next time user polls.
+     * the existing subscriptions after the next application poll event.
      */
     @Override
     public void transitionToStale() {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 4040d6823314d..4016b74b27bd6 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -72,6 +72,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
@@ -573,20 +574,48 @@ public void testPollTimerExpiration() {
                 backgroundEventHandler);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(new Node(1, "localhost", 9999)));
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
-        when(membershipManager.state()).thenReturn(MemberState.STABLE);
 
+        // On poll timer expiration, the member should transition to stale and a last heartbeat 
+        // should be sent to leave the group
         time.sleep(maxPollIntervalMs);
         assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
         verify(heartbeatState).reset();
         verify(heartbeatRequestState).reset();
         verify(membershipManager).transitionToStale();
 
+        when(membershipManager.state()).thenReturn(MemberState.STALE);
+        when(membershipManager.shouldSkipHeartbeat()).thenReturn(true);
         assertNoHeartbeat(heartbeatRequestManager);
         heartbeatRequestManager.resetPollTimer(time.milliseconds());
         assertTrue(pollTimer.notExpired());
+        verify(membershipManager).transitionToJoining();
+        when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
         assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
     }
 
+    /**
+     * This is expected to be the case where a member is already leaving the group and the poll
+     * timer expires. The poll timer expiration should not transition the member to STALE, and
+     * the member should continue to send heartbeats while the ongoing leaving operation
+     * completes (send heartbeats while waiting for callbacks before leaving, or send last
+     * heartbeat to leave).
+     */
+    @Test
+    public void testPollTimerExpirationShouldNotMarkMemberStaleIfMemberAlreadyLeaving() {
+        when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
+        when(membershipManager.isLeavingGroup()).thenReturn(true);
+        doNothing().when(membershipManager).transitionToStale();
+
+        time.sleep(maxPollIntervalMs);
+        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
+
+        // No transition to STALE should be triggered, because the member is already leaving the group
+        verify(membershipManager, never()).transitionToStale();
+
+        assertEquals(1, result.unsentRequests.size(), "A heartbeat request should be generated to" +
+            " complete the ongoing leaving operation that was triggered before the poll timer expired.");
+    }
+
     @Test
     public void testHeartbeatMetrics() {
         // setup
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 50fede2e9ef49..2d8ef9397f35e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -1611,6 +1611,61 @@ public void testOnPartitionsLostError() {
         testOnPartitionsLost(Optional.of(new KafkaException("Intentional error for test")));
     }
 
+    @Test
+    public void testTransitionToStaleWhileReconciling() {
+        MembershipManagerImpl membershipManager = memberJoinWithAssignment();
+        clearInvocations(subscriptionState);
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+
+        membershipManager.transitionToStale();
+        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+    }
+
+    @Test
+    public void testTransitionToStaleWhileJoining() {
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        doNothing().when(subscriptionState).assignFromSubscribed(any());
+        assertEquals(MemberState.JOINING, membershipManager.state());
+
+        membershipManager.transitionToStale();
+        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+    }
+
+    @Test
+    public void testTransitionToStaleWhileStable() {
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        doNothing().when(subscriptionState).assignFromSubscribed(any());
+        assertEquals(MemberState.STABLE, membershipManager.state());
+
+        membershipManager.transitionToStale();
+        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+    }
+
+    @Test
+    public void testTransitionToStaleWhileAcknowledging() {
+        MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true);
+        doNothing().when(subscriptionState).assignFromSubscribed(any());
+        clearInvocations(subscriptionState);
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+
+        membershipManager.transitionToStale();
+        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+    }
+
+    @Test
+    public void testStaleMemberDoesNotSendHeartbeatAndAllowsTransitionToJoiningToRecover() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        doNothing().when(subscriptionState).assignFromSubscribed(any());
+        membershipManager.transitionToStale();
+        assertTrue(membershipManager.shouldSkipHeartbeat(), "Stale member should not send " +
+            "heartbeats");
+        // Check that a transition to joining is allowed, which is what is expected to happen
+        // when the poll timer is reset.
+        membershipManager.transitionToJoining();
+    }
+
     private void mockPartitionOwnedAndNewPartitionAdded(String topicName,
                                                         int partitionOwned,
                                                         int partitionAdded,
@@ -1756,21 +1811,26 @@ private void testFenceIsNoOp(MembershipManagerImpl membershipManager) {
         verify(subscriptionState, never()).rebalanceListener();
     }
 
-    @Test
-    public void testTransitionToStaled() {
-        MembershipManager membershipManager = memberJoinWithAssignment("topic", Uuid.randomUuid());
-        membershipManager.transitionToStale();
+    private void assertStaleMemberClearsAssignmentsAndLeaves(MembershipManagerImpl membershipManager) {
+        assertEquals(MemberState.STALE, membershipManager.state());
+
+        // Should clear subscriptions, current assignments, and reset epoch to leave the group
+        verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
+        assertTrue(membershipManager.currentAssignment().isEmpty());
+        assertTrue(membershipManager.topicsAwaitingReconciliation().isEmpty());
         assertEquals(LEAVE_GROUP_MEMBER_EPOCH, membershipManager.memberEpoch());
     }
 
     @Test
-    public void testHeartbeatSentOnStaledMember() {
+    public void testHeartbeatSentOnStaleMember() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         subscriptionState.subscribe(Collections.singleton("topic"), Optional.empty());
         subscriptionState.assignFromSubscribed(Collections.singleton(new TopicPartition("topic", 0)));
         membershipManager.transitionToStale();
         membershipManager.onHeartbeatRequestSent();
-        assertEquals(MemberState.JOINING, membershipManager.state());
+        // Member should remain in STALE state. Only when the poll timer is reset the member will
+        // transition to JOINING.
+        assertEquals(MemberState.STALE, membershipManager.state());
         assertTrue(membershipManager.currentAssignment().isEmpty());
         assertTrue(subscriptionState.assignedPartitions().isEmpty());
     }
@@ -2157,10 +2217,11 @@ private ConsumerGroupHeartbeatResponseData.Assignment createAssignment(boolean m
                 ));
     }
 
-    private MembershipManager memberJoinWithAssignment(String topicName, Uuid topicId) {
+    private MembershipManagerImpl memberJoinWithAssignment() {
+        Uuid topicId = Uuid.randomUuid();
         MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true);
         membershipManager.onHeartbeatRequestSent();
-        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, "topic"));
         receiveAssignment(topicId, Collections.singletonList(0), membershipManager);
         membershipManager.onHeartbeatRequestSent();
         assertFalse(membershipManager.currentAssignment().isEmpty());

From 7a07aefd86f32824ade4821deb4bd81400b99d91 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 16 Feb 2024 02:53:53 -0800
Subject: [PATCH 044/258] =?UTF-8?q?KAFKA-16230:=20Update=20verifiable=5Fco?=
 =?UTF-8?q?nsumer.py=20to=20support=20KIP-848=E2=80=99s=20group=20protocol?=
 =?UTF-8?q?=20config=20(#15328)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Python VerifiableConsumer now passes in the --group-protocol and --group-remote-assignor command line arguments to VerifiableConsumer if the node is running 3.7.0+ and using the new consumer group.protocol.

Reviewers: Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../kafkatest/services/verifiable_consumer.py | 25 +++++++++++++---
 .../kafka/tools/VerifiableConsumer.java       | 30 ++++++++++++++++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/tests/kafkatest/services/verifiable_consumer.py b/tests/kafkatest/services/verifiable_consumer.py
index 93d9446fb9bba..e1155c16aaef4 100644
--- a/tests/kafkatest/services/verifiable_consumer.py
+++ b/tests/kafkatest/services/verifiable_consumer.py
@@ -21,7 +21,7 @@
 from kafkatest.directory_layout.kafka_path import KafkaPathResolverMixin
 from kafkatest.services.kafka import TopicPartition
 from kafkatest.services.verifiable_client import VerifiableClientMixin
-from kafkatest.version import DEV_BRANCH, V_2_3_0, V_2_3_1, V_0_10_0_0
+from kafkatest.version import DEV_BRANCH, V_2_3_0, V_2_3_1, V_3_7_0, V_0_10_0_0
 
 
 class ConsumerState:
@@ -167,7 +167,7 @@ class VerifiableConsumer(KafkaPathResolverMixin, VerifiableClientMixin, Backgrou
 
     def __init__(self, context, num_nodes, kafka, topic, group_id,
                  static_membership=False, max_messages=-1, session_timeout_sec=30, enable_autocommit=False,
-                 assignment_strategy=None,
+                 assignment_strategy=None, group_protocol=None, group_remote_assignor=None,
                  version=DEV_BRANCH, stop_timeout_sec=30, log_level="INFO", jaas_override_variables=None,
                  on_record_consumed=None, reset_policy="earliest", verify_offsets=True):
         """
@@ -184,6 +184,8 @@ def __init__(self, context, num_nodes, kafka, topic, group_id,
         self.session_timeout_sec = session_timeout_sec
         self.enable_autocommit = enable_autocommit
         self.assignment_strategy = assignment_strategy
+        self.group_protocol = group_protocol
+        self.group_remote_assignor = group_remote_assignor
         self.prop_file = ""
         self.stop_timeout_sec = stop_timeout_sec
         self.on_record_consumed = on_record_consumed
@@ -306,8 +308,20 @@ def start_cmd(self, node):
             # if `None` is passed as the argument value
             cmd += " --group-instance-id None"
 
-        if self.assignment_strategy:
-            cmd += " --assignment-strategy %s" % self.assignment_strategy
+        # 3.7.0 includes support for KIP-848 which introduced a new implementation of the consumer group protocol.
+        # The two implementations use slightly different configuration, hence these arguments are conditional.
+        #
+        # See the Java class/method VerifiableConsumer.createFromArgs() for how the command line arguments are
+        # parsed and used as configuration in the runner.
+        if node.version >= V_3_7_0 and self.is_consumer_group_protocol_enabled():
+            cmd += " --group-protocol %s" % self.group_protocol
+
+            if self.group_remote_assignor:
+                cmd += " --group-remote-assignor %s" % self.group_remote_assignor
+        else:
+            # Either we're an older consumer version or we're using the old consumer group protocol.
+            if self.assignment_strategy:
+                cmd += " --assignment-strategy %s" % self.assignment_strategy
 
         if self.enable_autocommit:
             cmd += " --enable-autocommit "
@@ -416,3 +430,6 @@ def alive_nodes(self):
         with self.lock:
             return [handler.node for handler in self.event_handlers.values()
                     if handler.state != ConsumerState.Dead]
+
+    def is_consumer_group_protocol_enabled(self):
+        return self.group_protocol and self.group_protocol.upper() == "CONSUMER"
diff --git a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
index 5fe66a5998834..afff2ce893c0f 100644
--- a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
+++ b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
@@ -533,11 +533,20 @@ private static ArgumentParser argParser() {
                 .action(store())
                 .required(false)
                 .type(String.class)
-                .setDefault(GroupProtocol.CLASSIC.name)
+                .setDefault(ConsumerConfig.DEFAULT_GROUP_PROTOCOL)
                 .metavar("GROUP_PROTOCOL")
                 .dest("groupProtocol")
                 .help(String.format("Group protocol (must be one of %s)", Utils.join(GroupProtocol.values(), ", ")));
 
+        parser.addArgument("--group-remote-assignor")
+                .action(store())
+                .required(false)
+                .type(String.class)
+                .setDefault(ConsumerConfig.DEFAULT_GROUP_REMOTE_ASSIGNOR)
+                .metavar("GROUP_REMOTE_ASSIGNOR")
+                .dest("groupRemoteAssignor")
+                .help(String.format("Group remote assignor; only used if the group protocol is %s", GroupProtocol.CONSUMER.name()));
+
         parser.addArgument("--group-id")
                 .action(store())
                 .required(true)
@@ -599,7 +608,7 @@ private static ArgumentParser argParser() {
                 .setDefault(RangeAssignor.class.getName())
                 .type(String.class)
                 .dest("assignmentStrategy")
-                .help("Set assignment strategy (e.g. " + RoundRobinAssignor.class.getName() + ")");
+                .help(String.format("Set assignment strategy (e.g. %s); only used if the group protocol is %s", RoundRobinAssignor.class.getName(), GroupProtocol.CLASSIC.name()));
 
         parser.addArgument("--consumer.config")
                 .action(store())
@@ -627,7 +636,21 @@ public static VerifiableConsumer createFromArgs(ArgumentParser parser, String[]
             }
         }
 
-        consumerProps.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, res.getString("groupProtocol"));
+        String groupProtocol = res.getString("groupProtocol");
+
+        // 3.7.0 includes support for KIP-848 which introduced a new implementation of the consumer group protocol.
+        // The two implementations use slightly different configuration, hence these arguments are conditional.
+        //
+        // See the Python class/method VerifiableConsumer.start_cmd() in verifiable_consumer.py for how the
+        // command line arguments are passed in by the system test framework.
+        if (groupProtocol.equalsIgnoreCase(GroupProtocol.CONSUMER.name())) {
+            consumerProps.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol);
+            consumerProps.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, res.getString("groupRemoteAssignor"));
+        } else {
+            // This means we're using the old consumer group protocol.
+            consumerProps.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, res.getString("assignmentStrategy"));
+        }
+
         consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, res.getString("groupId"));
 
         String groupInstanceId = res.getString("groupInstanceId");
@@ -650,7 +673,6 @@ public static VerifiableConsumer createFromArgs(ArgumentParser parser, String[]
         consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, useAutoCommit);
         consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, res.getString("resetPolicy"));
         consumerProps.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, Integer.toString(res.getInt("sessionTimeout")));
-        consumerProps.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, res.getString("assignmentStrategy"));
 
         StringDeserializer deserializer = new StringDeserializer();
         KafkaConsumer<String, String> consumer = new KafkaConsumer<>(consumerProps, deserializer, deserializer);

From 501f82b91c01b1bd54ced1248baaf0ff2a30cc04 Mon Sep 17 00:00:00 2001
From: Luke Chen <showuon@gmail.com>
Date: Fri, 16 Feb 2024 18:56:06 +0800
Subject: [PATCH 045/258] KAFKA-15670: add "inter.broker.listener.name" config
 in KRaft controller config (#14631)

During ZK migrating to KRaft, before entering dual-write mode, the KRaft controller will send RPCs (i.e. UpdateMetadataRequest, LeaderAndIsrRequest, and StopReplicaRequest) to the brokers. Currently, we use the inter broker listener to send the RPC to brokers from the controller. But in the doc, we didn't provide this info to users because the normal KRaft controller won't use inter.broker.listener.names.

This PR adds the missing config in the ZK migrating to KRaft doc.

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Paolo Patierno <ppatierno@live.com>
---
 docs/ops.html | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/docs/ops.html b/docs/ops.html
index 8cb65549db6ec..563c2a65000e1 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -3869,6 +3869,9 @@ <h3>Provisioning the KRaft controller quorum</h3>
 # ZooKeeper client configuration
 zookeeper.connect=localhost:2181
 
+# The inter broker listener in brokers to allow KRaft controller send RPCs to brokers
+inter.broker.listener.name=PLAINTEXT
+
 # Other configs ...</pre>
 
   <p><em>Note: The KRaft cluster <code>node.id</code> values must be different from any existing ZK broker <code>broker.id</code>.

From ddc5d1dc349d98dc430fb1f67a6f4ca108dc00e4 Mon Sep 17 00:00:00 2001
From: Paolo Patierno <ppatierno@live.com>
Date: Fri, 16 Feb 2024 11:57:22 +0100
Subject: [PATCH 046/258] MINOR: Added ACLs authorizer change during migration
 (#15333)

This trivial PR makes clear when it's the right time to switch from AclAuthorizer to StandardAuthorizer during the migration process.

Reviewers: Luke Chen <showuon@gmail.com>
---
 docs/ops.html | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docs/ops.html b/docs/ops.html
index 563c2a65000e1..cf039c752207c 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -3935,6 +3935,12 @@ <h3>Migrating brokers to KRaft</h3>
     The zookeeper configurations should be removed at this point.
   </p>
 
+  <p>
+    If your broker has authorization configured via the <code>authorizer.class.name</code> property
+    using <code>kafka.security.authorizer.AclAuthorizer</code>, this is also the time to change it
+    to use <code>org.apache.kafka.metadata.authorizer.StandardAuthorizer</code> instead.
+  </p>
+
   <pre>
 # Sample KRaft broker server.properties listening on 9092
 process.roles=broker

From 0789952b68963eb8dd1d20b6bd8db1729026fa88 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Fri, 16 Feb 2024 08:59:48 -0800
Subject: [PATCH 047/258] MINOR: add note about Kafka Streams feature for 3.7
 release (#15380)

Reviewers: Walker Carlson <wcarlson@confluent.io>
---
 docs/upgrade.html | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/docs/upgrade.html b/docs/upgrade.html
index bf66a3d3ad470..628396f361f54 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -31,6 +31,9 @@ <h5><a id="upgrade_370_notable" href="#upgrade_370_notable">Notable changes in 3
             of running Tiered Storage in production.
             For more detailed information, please refer to <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-963%3A+Additional+metrics+in+Tiered+Storage">KIP-963</a>.
         </li>
+	<li>Kafka Streams ships multiple KIPs for IQv2 support.
+            See the <a href="/{{version}}/documentation/streams/upgrade-guide#streams_api_changes_370">Kafka Streams upgrade section</a> for more details.
+        </li>
     </ul>
 
 <h4><a id="upgrade_3_6_0" href="#upgrade_3_6_0">Upgrading to 3.6.0 from any version 0.8.x through 3.5.x</a></h4>

From 756f44a3e582902827ae45daf382556c7fba99a0 Mon Sep 17 00:00:00 2001
From: Calvin Liu <83986057+CalvinConfluent@users.noreply.github.com>
Date: Fri, 16 Feb 2024 10:27:43 -0800
Subject: [PATCH 048/258] KAFKA-15665: Enforce partition reassignment should
 complete when all target replicas are in ISR (#15359)

When completing the partition reassignment, the new ISR should have all the target replicas.

Reviewers: Justine Olshan <jolshan@confluent.io>, David Mao <dmao@confluent.io>
---
 .../controller/PartitionReassignmentReplicas.java      |  4 +---
 .../controller/PartitionReassignmentReplicasTest.java  | 10 ++++++++++
 .../controller/ReplicationControlManagerTest.java      |  6 +++---
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/metadata/src/main/java/org/apache/kafka/controller/PartitionReassignmentReplicas.java b/metadata/src/main/java/org/apache/kafka/controller/PartitionReassignmentReplicas.java
index 62f84adea4cc7..56c3188f74184 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/PartitionReassignmentReplicas.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/PartitionReassignmentReplicas.java
@@ -120,9 +120,7 @@ Optional<CompletedReassignment> maybeCompleteReassignment(List<Integer> targetIs
             }
             if (newTargetReplicas.isEmpty()) return Optional.empty();
         }
-        for (int replica : adding) {
-            if (!newTargetIsr.contains(replica)) return Optional.empty();
-        }
+        if (!newTargetIsr.containsAll(newTargetReplicas)) return Optional.empty();
 
         return Optional.of(
             new CompletedReassignment(
diff --git a/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java b/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
index dd1d567e58715..b2bc540bda5f6 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
@@ -196,6 +196,16 @@ public void testIsReassignmentInProgress() {
                 build()));
     }
 
+    @Test
+    public void testDoesNotCompleteReassignmentIfIsrDoesNotHaveAllTargetReplicas() {
+        PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
+            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(0, 1, 3)));
+        assertTrue(replicas.isReassignmentInProgress());
+        Optional<PartitionReassignmentReplicas.CompletedReassignment> reassignmentOptional =
+            replicas.maybeCompleteReassignment(Arrays.asList(3));
+        assertFalse(reassignmentOptional.isPresent());
+    }
+
     @Test
     public void testOriginalReplicas() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
index bf7f6c82e051b..3d54720be921f 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
@@ -1989,7 +1989,7 @@ public void testCancelReassignPartitions() throws Exception {
                 new AlterPartitionReassignmentsRequestData().setTopics(asList(
                     new ReassignableTopic().setName("foo").setPartitions(asList(
                         new ReassignablePartition().setPartitionIndex(0).
-                            setReplicas(asList(1, 2, 3)),
+                            setReplicas(asList(1, 2, 4)),
                         new ReassignablePartition().setPartitionIndex(1).
                             setReplicas(asList(1, 2, 3, 0)),
                         new ReassignablePartition().setPartitionIndex(2).
@@ -2019,11 +2019,11 @@ public void testCancelReassignPartitions() throws Exception {
                     setErrorMessage(null))))),
             alterResult.response());
         ctx.replay(alterResult.records());
-        assertEquals(new PartitionRegistration.Builder().setReplicas(new int[] {1, 2, 3}).setIsr(new int[] {1, 2}).
+        assertEquals(new PartitionRegistration.Builder().setReplicas(new int[] {1, 2, 4}).setIsr(new int[] {1, 2, 4}).
             setDirectories(new Uuid[] {
                     Uuid.fromString("TESTBROKER00001DIRAAAA"),
                     Uuid.fromString("TESTBROKER00002DIRAAAA"),
-                    Uuid.fromString("TESTBROKER00003DIRAAAA")
+                    Uuid.fromString("TESTBROKER00004DIRAAAA")
             }).
             setLeader(1).setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).setLeaderEpoch(1).setPartitionEpoch(2).build(), replication.getPartition(fooId, 0));
         assertEquals(new PartitionRegistration.Builder().setReplicas(new int[] {1, 2, 3, 0}).setIsr(new int[] {0, 1, 2}).

From 98fb3bd304ccf6eb58e4a84bf2827669775259ad Mon Sep 17 00:00:00 2001
From: Luke Chen <showuon@gmail.com>
Date: Sat, 17 Feb 2024 13:58:11 +0800
Subject: [PATCH 049/258] MINOR: log error when initialLoadFuture is not done
 in authorizer (#14953)

Currently, when initializing StandardAuthorizer, it'll wait until all ACL loaded and complete the initialLoadFuture. So, checking logs, we'll see:

2023-12-06 14:07:50,325 INFO [StandardAuthorizer 1] Initialized with 6 acl(s). (org.apache.kafka.metadata.authorizer.StandardAuthorizerData) [kafka-1-metadata-loader-event-handler]
2023-12-06 14:07:50,325 INFO [StandardAuthorizer 1] Completed initial ACL load process. (org.apache.kafka.metadata.authorizer.StandardAuthorizerData) [kafka-1-metadata-loader-event-handler]

But then, when shutting down the node, we will also see this error:

2023-12-06 14:12:32,752 ERROR [StandardAuthorizer 1] Failed to complete initial ACL load process. (org.apache.kafka.metadata.authorizer.StandardAuthorizerData) [kafka-1-metadata-loader-event-handler]
java.util.concurrent.TimeoutException
	at kafka.server.metadata.AclPublisher.close(AclPublisher.scala:98)
	at org.apache.kafka.image.loader.MetadataLoader.closePublisher(MetadataLoader.java:568)
	at org.apache.kafka.image.loader.MetadataLoader.lambda$removeAndClosePublisher$7(MetadataLoader.java:528)
	at org.apache.kafka.queue.KafkaEventQueue$EventContext.run(KafkaEventQueue.java:127)
	at org.apache.kafka.queue.KafkaEventQueue$EventHandler.handleEvents(KafkaEventQueue.java:210)
	at org.apache.kafka.queue.KafkaEventQueue$EventHandler.run(KafkaEventQueue.java:181)
	at java.base/java.lang.Thread.run(Thread.java:840)

It's confusing. And it's because we'll try to complete authorizer initialLoad, and complete the initialLoadFuture if not done. But we'll log the error no matter it's completed or not. This patch improves the logging.

Reviewers: Josep Prat <josep.prat@aiven.io>
---
 .../kafka/metadata/authorizer/StandardAuthorizer.java       | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/metadata/src/main/java/org/apache/kafka/metadata/authorizer/StandardAuthorizer.java b/metadata/src/main/java/org/apache/kafka/metadata/authorizer/StandardAuthorizer.java
index 7e9f779093ecb..106f6aba94990 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/authorizer/StandardAuthorizer.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/authorizer/StandardAuthorizer.java
@@ -93,8 +93,10 @@ public CompletableFuture<Void> initialLoadFuture() {
 
     @Override
     public void completeInitialLoad(Exception e) {
-        data.log.error("Failed to complete initial ACL load process.", e);
-        initialLoadFuture.completeExceptionally(e);
+        if (!initialLoadFuture.isDone()) {
+            data.log.error("Failed to complete initial ACL load process.", e);
+            initialLoadFuture.completeExceptionally(e);
+        }
     }
 
     @Override

From e247bd03afe66d61426a9029220d06438dede3dc Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Sat, 17 Feb 2024 00:07:50 -0800
Subject: [PATCH 050/258] MINOR: Improve
 ListConsumerGroupTest.testListGroupCommand (#15382)

While reviewing https://github.com/apache/kafka/pull/15150, I found that our tests verifying the console output are really hard to read. Here is my proposal to make it better.

Reviewers: Ritika Reddy <rreddy@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../consumer/group/ListConsumerGroupTest.java | 121 +++++++++++++-----
 1 file changed, 86 insertions(+), 35 deletions(-)

diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
index 6648cfe2a2414..894f00df5e7f6 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
@@ -26,10 +26,15 @@
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 
+import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -123,42 +128,88 @@ public void testListGroupCommand(String quorum) throws Exception {
         String simpleGroup = "simple-group";
         addSimpleGroupExecutor(simpleGroup);
         addConsumerGroupExecutor(1);
-        final AtomicReference<String> out = new AtomicReference<>("");
 
-        String[] cgcArgs1 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
-        TestUtils.waitForCondition(() -> {
-            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
-                ConsumerGroupCommand.main(cgcArgs1);
-                return null;
-            }));
-            return !out.get().contains("STATE") && out.get().contains(simpleGroup) && out.get().contains(GROUP);
-        }, "Expected to find " + simpleGroup + ", " + GROUP + " and no header, but found " + out.get());
-
-        String[] cgcArgs2 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"};
-        TestUtils.waitForCondition(() -> {
-            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
-                ConsumerGroupCommand.main(cgcArgs2);
-                return null;
-            }));
-            return out.get().contains("STATE") && out.get().contains(simpleGroup) && out.get().contains(GROUP);
-        }, "Expected to find " + simpleGroup + ", " + GROUP + " and the header, but found " + out.get());
-
-        String[] cgcArgs3 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "Stable"};
-        TestUtils.waitForCondition(() -> {
-            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
-                ConsumerGroupCommand.main(cgcArgs3);
-                return null;
-            }));
-            return out.get().contains("STATE") && out.get().contains(GROUP) && out.get().contains("Stable");
-        }, "Expected to find " + GROUP + " in state Stable and the header, but found " + out.get());
-
-        String[] cgcArgs4 = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "stable"};
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list"),
+            Collections.emptyList(),
+            mkSet(
+                Collections.singletonList(GROUP),
+                Collections.singletonList(simpleGroup)
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"),
+            Arrays.asList("GROUP", "STATE"),
+            mkSet(
+                Arrays.asList(GROUP, "Stable"),
+                Arrays.asList(simpleGroup, "Empty")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "Stable"),
+            Arrays.asList("GROUP", "STATE"),
+            mkSet(
+                Arrays.asList(GROUP, "Stable")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "stable"),
+            Arrays.asList("GROUP", "STATE"),
+            mkSet(
+                Arrays.asList(GROUP, "Stable")
+            )
+        );
+    }
+
+    /**
+     * Validates that the output of the list command corresponds to the expected values.
+     *
+     * @param args              The arguments for the command line tool.
+     * @param expectedHeader    The expected header as a list of strings; or an empty list
+     *                          if a header is not expected.
+     * @param expectedRows      The expected rows as a set of list of columns.
+     * @throws InterruptedException
+     */
+    private static void validateListOutput(
+        List<String> args,
+        List<String> expectedHeader,
+        Set<List<String>> expectedRows
+    ) throws InterruptedException {
+        final AtomicReference<String> out = new AtomicReference<>("");
         TestUtils.waitForCondition(() -> {
-            out.set(kafka.utils.TestUtils.grabConsoleOutput(() -> {
-                ConsumerGroupCommand.main(cgcArgs4);
-                return null;
-            }));
-            return out.get().contains("STATE") && out.get().contains(GROUP) && out.get().contains("Stable");
-        }, "Expected to find " + GROUP + " in state Stable and the header, but found " + out.get());
+            String output = runAndGrabConsoleOutput(args);
+            out.set(output);
+
+            int index = 0;
+            String[] lines = output.split("\n");
+
+            // Parse the header if one is expected.
+            if (!expectedHeader.isEmpty()) {
+                if (lines.length == 0) return false;
+                List<String> header = Arrays.stream(lines[index++].split("\\s+")).collect(Collectors.toList());
+                if (!expectedHeader.equals(header)) {
+                    return false;
+                }
+            }
+
+            // Parse the groups.
+            Set<List<String>> groups = new HashSet<>();
+            for (; index < lines.length; index++) {
+                groups.add(Arrays.stream(lines[index].split("\\s+")).collect(Collectors.toList()));
+            }
+            return expectedRows.equals(groups);
+        }, () -> String.format("Expected header=%s and groups=%s, but found:%n%s", expectedHeader, expectedRows, out.get()));
+    }
+
+    private static String runAndGrabConsoleOutput(
+        List<String> args
+    ) {
+        return kafka.utils.TestUtils.grabConsoleOutput(() -> {
+            ConsumerGroupCommand.main(args.toArray(new String[0]));
+            return null;
+        });
     }
 }

From 71a4e6fc0ce43e907c320ed5afca95b84620b0e8 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Mon, 19 Feb 2024 19:37:31 +0800
Subject: [PATCH 051/258] KAFKA-15140: improve TopicCommandIntegrationTest to
 be less flaky (#14891)

This PR improves TopicCommandIntegrationTest by :
    - using TestUtils.createTopicWithAdmin
    - replacing \n with lineSeperator
    - using waitForAllReassignmentsToComplete
    - adding more log when assertion fails

Reviewers: Luke Chen <showuon@gmail.com>, Justine Olshan <jolshan@confluent.io>
---
 .../tools/TopicCommandIntegrationTest.java    | 526 ++++++++++--------
 1 file changed, 298 insertions(+), 228 deletions(-)

diff --git a/tools/src/test/java/org/apache/kafka/tools/TopicCommandIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/TopicCommandIntegrationTest.java
index b934e04012cc8..fa1d8ea8c511a 100644
--- a/tools/src/test/java/org/apache/kafka/tools/TopicCommandIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/TopicCommandIntegrationTest.java
@@ -18,6 +18,7 @@
 package org.apache.kafka.tools;
 
 import kafka.admin.RackAwareTest;
+import kafka.server.ControllerServer;
 import kafka.server.KafkaBroker;
 import kafka.server.KafkaConfig;
 import kafka.utils.Logging;
@@ -28,7 +29,6 @@
 import org.apache.kafka.clients.admin.Config;
 import org.apache.kafka.clients.admin.ListPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.NewPartitionReassignment;
-import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.admin.PartitionReassignment;
 import org.apache.kafka.clients.admin.TopicDescription;
 import org.apache.kafka.common.Node;
@@ -38,7 +38,6 @@
 import org.apache.kafka.common.config.ConfigResource;
 import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
-import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.TopicExistsException;
 import org.apache.kafka.common.internals.Topic;
 import org.apache.kafka.common.message.MetadataResponseData;
@@ -53,6 +52,8 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import scala.collection.JavaConverters;
+import scala.collection.mutable.Buffer;
+import scala.collection.Seq;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -66,6 +67,7 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -82,13 +84,16 @@
 @SuppressWarnings("deprecation") // Added for Scala 2.12 compatibility for usages of JavaConverters
 public class TopicCommandIntegrationTest extends kafka.integration.KafkaServerTestHarness implements Logging, RackAwareTest {
     private short defaultReplicationFactor = 1;
-    private int numPartitions = 1;
+    private int defaultNumPartitions = 1;
     private TopicCommand.TopicService topicService;
     private Admin adminClient;
     private String bootstrapServer;
     private String testTopicName;
     private long defaultTimeout = 10000;
 
+    private Buffer<KafkaBroker> scalaBrokers;
+    private Seq<ControllerServer> scalaControllers;
+
     /**
      * Implementations must override this method to return a set of KafkaConfigs. This method will be invoked for every
      * test and should not reuse previous configurations unless they select their ports randomly when servers are started.
@@ -107,7 +112,7 @@ public scala.collection.Seq<KafkaConfig> generateConfigs() {
         rackInfo.put(5, "rack3");
 
         List<Properties> brokerConfigs = ToolsTestUtils
-            .createBrokerProperties(6, zkConnectOrNull(), rackInfo, numPartitions, defaultReplicationFactor);
+            .createBrokerProperties(6, zkConnectOrNull(), rackInfo, defaultNumPartitions, defaultReplicationFactor);
 
         List<KafkaConfig> configs = new ArrayList<>();
         for (Properties props : brokerConfigs) {
@@ -124,19 +129,6 @@ private TopicCommand.TopicCommandOptions buildTopicCommandOptionsWithBootstrap(S
         return new TopicCommand.TopicCommandOptions(finalOptions);
     }
 
-    private void createAndWaitTopic(TopicCommand.TopicCommandOptions opts) throws Exception {
-        topicService.createTopic(opts);
-        waitForTopicCreated(opts.topic().get());
-    }
-
-    private void waitForTopicCreated(String topicName) {
-        waitForTopicCreated(topicName, defaultTimeout);
-    }
-
-    private void waitForTopicCreated(String topicName, long timeout) {
-        TestUtils.waitForPartitionMetadata(brokers(), topicName, 0, timeout);
-    }
-
     @BeforeEach
     public void setUp(TestInfo info) {
         super.setUp(info);
@@ -147,7 +139,10 @@ public void setUp(TestInfo info) {
         adminClient = Admin.create(props);
         topicService = new TopicCommand.TopicService(props, Optional.of(bootstrapServer));
         testTopicName = String.format("%s-%s", info.getTestMethod().get().getName(), TestUtils.randomString(10));
+        scalaBrokers = brokers();
+        scalaControllers = controllerServers();
     }
+
     @AfterEach
     public void close() throws Exception {
         if (topicService != null)
@@ -159,47 +154,51 @@ public void close() throws Exception {
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreate(String quorum) throws Exception {
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap(
-            "--create", "--partitions", "2", "--replication-factor", "1", "--topic", testTopicName));
-
-        assertTrue(adminClient.listTopics().names().get().contains(testTopicName));
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, 2, 1,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        assertTrue(adminClient.listTopics().names().get().contains(testTopicName),
+                "Admin client didn't see the created topic. It saw: " + adminClient.listTopics().names().get());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithDefaults(String quorum) throws Exception {
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap("--create", "--topic", testTopicName));
-
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         List<TopicPartitionInfo> partitions = adminClient
             .describeTopics(Collections.singletonList(testTopicName))
             .allTopicNames()
             .get()
             .get(testTopicName)
             .partitions();
-        assertEquals(numPartitions, partitions.size());
-        assertEquals(defaultReplicationFactor, (short) partitions.get(0).replicas().size());
+        assertEquals(defaultNumPartitions, partitions.size(), "Unequal partition size: " + partitions.size());
+        assertEquals(defaultReplicationFactor, (short) partitions.get(0).replicas().size(), "Unequal replication factor: " + partitions.get(0).replicas().size());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithDefaultReplication(String quorum) throws Exception {
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap("--create", "--topic", testTopicName, "--partitions", "2"));
-
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, 2, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         List<TopicPartitionInfo>  partitions = adminClient
             .describeTopics(Collections.singletonList(testTopicName))
             .allTopicNames()
             .get()
             .get(testTopicName)
             .partitions();
-        assertEquals(2, partitions.size());
-        assertEquals(defaultReplicationFactor, (short) partitions.get(0).replicas().size());
+        assertEquals(2, partitions.size(), "Unequal partition size: " + partitions.size());
+        assertEquals(defaultReplicationFactor, (short) partitions.get(0).replicas().size(), "Unequal replication factor: " + partitions.get(0).replicas().size());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithDefaultPartitions(String quorum) throws Exception {
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap("--create", "--topic", testTopicName, "--replication-factor", "2"));
-
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, 2,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         List<TopicPartitionInfo> partitions = adminClient
             .describeTopics(Collections.singletonList(testTopicName))
             .allTopicNames()
@@ -207,52 +206,65 @@ public void testCreateWithDefaultPartitions(String quorum) throws Exception {
             .get(testTopicName)
             .partitions();
 
-        assertEquals(numPartitions, partitions.size());
-        assertEquals(2, (short) partitions.get(0).replicas().size());
+        assertEquals(defaultNumPartitions, partitions.size(), "Unequal partition size: " + partitions.size());
+        assertEquals(2, (short) partitions.get(0).replicas().size(), "Partitions not replicated: " + partitions.get(0).replicas().size());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithConfigs(String quorum) throws Exception {
         ConfigResource configResource = new ConfigResource(ConfigResource.Type.TOPIC, testTopicName);
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "2", "--replication-factor", "2", "--topic", testTopicName, "--config",
-                "delete.retention.ms=1000"));
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.DELETE_RETENTION_MS_CONFIG, "1000");
+
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, 2, 2,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
 
         Config configs = adminClient.describeConfigs(Collections.singleton(configResource)).all().get().get(configResource);
-        assertEquals(1000, Integer.valueOf(configs.get("delete.retention.ms").value()));
+        assertEquals(1000, Integer.valueOf(configs.get("delete.retention.ms").value()),
+                "Config not set correctly: " + configs.get("delete.retention.ms").value());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWhenAlreadyExists(String quorum) throws Exception {
-        int numPartitions = 1;
-
         // create the topic
         TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap(
-            "--create", "--partitions", Integer.toString(numPartitions), "--replication-factor", "1",
+            "--create", "--partitions", Integer.toString(defaultNumPartitions), "--replication-factor", "1",
                 "--topic", testTopicName);
-        createAndWaitTopic(createOpts);
 
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         // try to re-create the topic
-        assertThrows(TopicExistsException.class, () -> topicService.createTopic(createOpts));
+        assertThrows(TopicExistsException.class, () -> topicService.createTopic(createOpts),
+                "Expected TopicExistsException to throw");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWhenAlreadyExistsWithIfNotExists(String quorum) throws Exception {
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         TopicCommand.TopicCommandOptions createOpts =
                 buildTopicCommandOptionsWithBootstrap("--create", "--topic", testTopicName, "--if-not-exists");
-        createAndWaitTopic(createOpts);
         topicService.createTopic(createOpts);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithReplicaAssignment(String quorum) throws Exception {
-        // create the topic
-        TopicCommand.TopicCommandOptions createOpts =
-                buildTopicCommandOptionsWithBootstrap("--create", "--replica-assignment", "5:4,3:2,1:0", "--topic", testTopicName);
-        createAndWaitTopic(createOpts);
+        scala.collection.mutable.HashMap<Object, Seq<Object>> replicaAssignmentMap = new scala.collection.mutable.HashMap<>();
+
+        replicaAssignmentMap.put(0, JavaConverters.asScalaBufferConverter(Arrays.asList((Object) 5, (Object) 4)).asScala().toSeq());
+        replicaAssignmentMap.put(1, JavaConverters.asScalaBufferConverter(Arrays.asList((Object) 3, (Object) 2)).asScala().toSeq());
+        replicaAssignmentMap.put(2, JavaConverters.asScalaBufferConverter(Arrays.asList((Object) 1, (Object) 0)).asScala().toSeq());
+
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions,
+                defaultReplicationFactor, replicaAssignmentMap, new Properties()
+        );
 
         List<TopicPartitionInfo> partitions = adminClient
             .describeTopics(Collections.singletonList(testTopicName))
@@ -261,10 +273,14 @@ public void testCreateWithReplicaAssignment(String quorum) throws Exception {
             .get(testTopicName)
             .partitions();
         
-        assertEquals(3, partitions.size());
-        assertEquals(Arrays.asList(5, 4), getPartitionReplicas(partitions, 0));
-        assertEquals(Arrays.asList(3, 2), getPartitionReplicas(partitions, 1));
-        assertEquals(Arrays.asList(1, 0), getPartitionReplicas(partitions, 2));
+        assertEquals(3, partitions.size(),
+                "Unequal partition size: " + partitions.size());
+        assertEquals(Arrays.asList(5, 4), getPartitionReplicas(partitions, 0),
+                "Unexpected replica assignment: " + getPartitionReplicas(partitions, 0));
+        assertEquals(Arrays.asList(3, 2), getPartitionReplicas(partitions, 1),
+                "Unexpected replica assignment: " + getPartitionReplicas(partitions, 1));
+        assertEquals(Arrays.asList(1, 0), getPartitionReplicas(partitions, 2),
+                "Unexpected replica assignment: " + getPartitionReplicas(partitions, 2));
     }
 
     private List<Integer> getPartitionReplicas(List<TopicPartitionInfo> partitions, int partitionNumber) {
@@ -276,7 +292,7 @@ private List<Integer> getPartitionReplicas(List<TopicPartitionInfo> partitions,
     public void testCreateWithInvalidReplicationFactor(String quorum) {
         TopicCommand.TopicCommandOptions opts = buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "2", "--replication-factor", Integer.toString(Short.MAX_VALUE + 1),
             "--topic", testTopicName);
-        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts));
+        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts), "Expected IllegalArgumentException to throw");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -284,14 +300,14 @@ public void testCreateWithInvalidReplicationFactor(String quorum) {
     public void testCreateWithNegativeReplicationFactor(String quorum) {
         TopicCommand.TopicCommandOptions opts = buildTopicCommandOptionsWithBootstrap("--create",
             "--partitions", "2", "--replication-factor", "-1", "--topic", testTopicName);
-        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts));
+        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts), "Expected IllegalArgumentException to throw");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithNegativePartitionCount(String quorum) {
         TopicCommand.TopicCommandOptions opts = buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "-1", "--replication-factor", "1", "--topic", testTopicName);
-        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts));
+        assertThrows(IllegalArgumentException.class, () -> topicService.createTopic(opts), "Expected IllegalArgumentException to throw");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -300,17 +316,18 @@ public void testInvalidTopicLevelConfig(String quorum) {
         TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap("--create",
             "--partitions", "1", "--replication-factor", "1", "--topic", testTopicName,
             "--config", "message.timestamp.type=boom");
-        assertThrows(ConfigException.class, () -> topicService.createTopic(createOpts));
+        assertThrows(ConfigException.class, () -> topicService.createTopic(createOpts), "Expected ConfigException to throw");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testListTopics(String quorum) throws Exception {
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap(
-            "--create", "--partitions", "1", "--replication-factor", "1", "--topic", testTopicName));
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         String output = captureListTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--list"));
-        assertTrue(output.contains(testTopicName));
+        assertTrue(output.contains(testTopicName), "Expected topic name to be present in output: " + output);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -319,107 +336,124 @@ public void testListTopicsWithIncludeList(String quorum) throws ExecutionExcepti
         String topic1 = "kafka.testTopic1";
         String topic2 = "kafka.testTopic2";
         String topic3 = "oooof.testTopic1";
-        adminClient.createTopics(
-                Arrays.asList(new NewTopic(topic1, 2, (short) 2),
-                    new NewTopic(topic2, 2, (short) 2),
-                    new NewTopic(topic3, 2, (short) 2)))
-            .all().get();
-        waitForTopicCreated(topic1);
-        waitForTopicCreated(topic2);
-        waitForTopicCreated(topic3);
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, topic1, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, topic2, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, topic3, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         String output = captureListTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--list", "--topic", "kafka.*"));
 
-        assertTrue(output.contains(topic1));
-        assertTrue(output.contains(topic2));
-        assertFalse(output.contains(topic3));
+        assertTrue(output.contains(topic1), "Expected topic name " + topic1 + " to be present in output: " + output);
+        assertTrue(output.contains(topic2), "Expected topic name " + topic2 + " to be present in output: " + output);
+        assertFalse(output.contains(topic3), "Do not expect topic name " + topic3 + " to be present in output: " + output);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testListTopicsWithExcludeInternal(String quorum) throws ExecutionException, InterruptedException {
         String topic1 = "kafka.testTopic1";
-        adminClient.createTopics(
-                Arrays.asList(new NewTopic(topic1, 2, (short) 2),
-                    new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, 2, (short) 2)))
-            .all().get();
-        waitForTopicCreated(topic1);
+        String hiddenConsumerTopic = Topic.GROUP_METADATA_TOPIC_NAME;
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, topic1, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, hiddenConsumerTopic, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         String output = captureListTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--list", "--exclude-internal"));
 
-        assertTrue(output.contains(topic1));
-        assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME));
+        assertTrue(output.contains(topic1), "Expected topic name " + topic1 + " to be present in output: " + output);
+        assertFalse(output.contains(hiddenConsumerTopic), "Do not expect topic name " + hiddenConsumerTopic + " to be present in output: " + output);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testAlterPartitionCount(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Arrays.asList(new NewTopic(testTopicName, 2, (short) 2))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter", "--topic", testTopicName, "--partitions", "3"));
 
+        TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
         kafka.utils.TestUtils.waitUntilTrue(
             () -> brokers().forall(b -> b.metadataCache().getTopicPartitions(testTopicName).size() == 3),
             () -> "Timeout waiting for new assignment propagating to broker", org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS, 100L);
         TopicDescription topicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName)).topicNameValues().get(testTopicName).get();
-        assertEquals(3, topicDescription.partitions().size());
+        assertEquals(3, topicDescription.partitions().size(), "Expected partition count to be 3. Got: " + topicDescription.partitions().size());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testAlterAssignment(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 2, (short) 2))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter",
             "--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2", "--partitions", "3"));
+
+        TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
         kafka.utils.TestUtils.waitUntilTrue(
             () -> brokers().forall(b -> b.metadataCache().getTopicPartitions(testTopicName).size() == 3),
             () -> "Timeout waiting for new assignment propagating to broker",
             org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS, 100L);
 
         TopicDescription topicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName)).topicNameValues().get(testTopicName).get();
-        assertTrue(topicDescription.partitions().size() == 3);
+        assertTrue(topicDescription.partitions().size() == 3, "Expected partition count to be 3. Got: " + topicDescription.partitions().size());
         List<Integer> partitionReplicas = getPartitionReplicas(topicDescription.partitions(), 2);
-        assertEquals(Arrays.asList(4, 2), partitionReplicas);
+        assertEquals(Arrays.asList(4, 2), partitionReplicas, "Expected to have replicas 4,2. Got: " + partitionReplicas);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testAlterAssignmentWithMoreAssignmentThanPartitions(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Arrays.asList(new NewTopic(testTopicName, 2, (short) 2))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         assertThrows(ExecutionException.class,
             () -> topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter",
-                "--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2,3:2", "--partitions", "3")));
+                "--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2,3:2", "--partitions", "3")),
+                "Expected to fail with ExecutionException");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testAlterAssignmentWithMorePartitionsThanAssignment(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Arrays.asList(new NewTopic(testTopicName, 2, (short) 2))).all().get();
-        waitForTopicCreated(testTopicName);
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         assertThrows(ExecutionException.class,
             () -> topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter", "--topic", testTopicName,
-                "--replica-assignment", "5:3,3:1,4:2", "--partitions", "6")));
+                "--replica-assignment", "5:3,3:1,4:2", "--partitions", "6")),
+                "Expected to fail with ExecutionException");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testAlterWithInvalidPartitionCount(String quorum) throws Exception {
-        createAndWaitTopic(
-                buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
         );
-
         assertThrows(ExecutionException.class,
-            () -> topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter", "--partitions", "-1", "--topic", testTopicName)));
+            () -> topicService.alterTopic(buildTopicCommandOptionsWithBootstrap("--alter", "--partitions", "-1", "--topic", testTopicName)),
+                "Expected to fail with ExecutionException");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -428,7 +462,7 @@ public void testAlterWhenTopicDoesntExist(String quorum) {
         // alter a topic that does not exist without --if-exists
         TopicCommand.TopicCommandOptions alterOpts = buildTopicCommandOptionsWithBootstrap("--alter", "--topic", testTopicName, "--partitions", "1");
         TopicCommand.TopicService topicService = new TopicCommand.TopicService(adminClient);
-        assertThrows(IllegalArgumentException.class, () -> topicService.alterTopic(alterOpts));
+        assertThrows(IllegalArgumentException.class, () -> topicService.alterTopic(alterOpts), "Expected to fail with IllegalArgumentException");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -450,11 +484,9 @@ public void testCreateAlterTopicWithRackAware(String quorum) throws Exception {
 
         int numPartitions = 18;
         int replicationFactor = 3;
-        TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap("--create",
-            "--partitions", Integer.toString(numPartitions),
-            "--replication-factor", Integer.toString(replicationFactor),
-            "--topic", testTopicName);
-        createAndWaitTopic(createOpts);
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, numPartitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         Map<Integer, List<Integer>> assignment = adminClient.describeTopics(Collections.singletonList(testTopicName))
             .allTopicNames().get().get(testTopicName).partitions()
@@ -472,6 +504,7 @@ public void testCreateAlterTopicWithRackAware(String quorum) throws Exception {
             "--topic", testTopicName);
         topicService.alterTopic(alterOpts);
 
+        TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
         kafka.utils.TestUtils.waitUntilTrue(
             () -> brokers().forall(p -> p.metadataCache().getTopicPartitions(testTopicName).size() == alteredNumPartitions),
             () -> "Timeout waiting for new assignment propagating to broker",
@@ -487,22 +520,18 @@ public void testCreateAlterTopicWithRackAware(String quorum) throws Exception {
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testConfigPreservationAcrossPartitionAlteration(String quorum) throws Exception {
-        int numPartitionsOriginal = 1;
-        String cleanupKey = "cleanup.policy";
-        String cleanupVal = "compact";
+        String cleanUpPolicy = "compact";
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.CLEANUP_POLICY_CONFIG, cleanUpPolicy);
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
 
-        // create the topic
-        TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap("--create",
-            "--partitions", Integer.toString(numPartitionsOriginal),
-            "--replication-factor", "1",
-            "--config", cleanupKey + "=" + cleanupVal,
-            "--topic", testTopicName);
-        createAndWaitTopic(createOpts);
         ConfigResource configResource = new ConfigResource(ConfigResource.Type.TOPIC, testTopicName);
         Config props = adminClient.describeConfigs(Collections.singleton(configResource)).all().get().get(configResource);
         // val props = adminZkClient.fetchEntityConfig(ConfigType.Topic, testTopicName)
-        assertNotNull(props.get(cleanupKey), "Properties after creation don't contain " + cleanupKey);
-        assertEquals(cleanupVal, props.get(cleanupKey).value(), "Properties after creation have incorrect value");
+        assertNotNull(props.get(TopicConfig.CLEANUP_POLICY_CONFIG), "Properties after creation don't contain " + cleanUpPolicy);
+        assertEquals(cleanUpPolicy, props.get(TopicConfig.CLEANUP_POLICY_CONFIG).value(), "Properties after creation have incorrect value");
 
         // pre-create the topic config changes path to avoid a NoNodeException
         if (!isKRaftTest()) {
@@ -514,21 +543,20 @@ public void testConfigPreservationAcrossPartitionAlteration(String quorum) throw
         TopicCommand.TopicCommandOptions alterOpts = buildTopicCommandOptionsWithBootstrap("--alter",
             "--partitions", Integer.toString(numPartitionsModified), "--topic", testTopicName);
         topicService.alterTopic(alterOpts);
+
+        TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
         Config newProps = adminClient.describeConfigs(Collections.singleton(configResource)).all().get().get(configResource);
-        assertNotNull(newProps.get(cleanupKey), "Updated properties do not contain " + cleanupKey);
-        assertEquals(cleanupVal, newProps.get(cleanupKey).value(), "Updated properties have incorrect value");
+        assertNotNull(newProps.get(TopicConfig.CLEANUP_POLICY_CONFIG), "Updated properties do not contain " + TopicConfig.CLEANUP_POLICY_CONFIG);
+        assertEquals(cleanUpPolicy, newProps.get(TopicConfig.CLEANUP_POLICY_CONFIG).value(), "Updated properties have incorrect value");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testTopicDeletion(String quorum) throws Exception {
         // create the NormalTopic
-        TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap("--create",
-            "--partitions", "1",
-            "--replication-factor", "1",
-            "--topic", testTopicName);
-        createAndWaitTopic(createOpts);
-
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         // delete the NormalTopic
         TopicCommand.TopicCommandOptions deleteOpts = buildTopicCommandOptionsWithBootstrap("--delete", "--topic", testTopicName);
 
@@ -545,12 +573,9 @@ public void testTopicDeletion(String quorum) throws Exception {
     public void testTopicWithCollidingCharDeletionAndCreateAgain(String quorum) throws Exception {
         // create the topic with colliding chars
         String topicWithCollidingChar = "test.a";
-        TopicCommand.TopicCommandOptions createOpts = buildTopicCommandOptionsWithBootstrap("--create",
-            "--partitions", "1",
-            "--replication-factor", "1",
-            "--topic", topicWithCollidingChar);
-        createAndWaitTopic(createOpts);
-
+        TestUtils.createTopicWithAdmin(adminClient, topicWithCollidingChar, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         // delete the topic
         TopicCommand.TopicCommandOptions deleteOpts = buildTopicCommandOptionsWithBootstrap("--delete", "--topic", topicWithCollidingChar);
 
@@ -560,19 +585,18 @@ public void testTopicWithCollidingCharDeletionAndCreateAgain(String quorum) thro
         }
         topicService.deleteTopic(deleteOpts);
         TestUtils.verifyTopicDeletion(zkClientOrNull(), topicWithCollidingChar, 1, brokers());
-        assertDoesNotThrow(() -> createAndWaitTopic(createOpts));
+        assertDoesNotThrow(() -> TestUtils.createTopicWithAdmin(adminClient, topicWithCollidingChar, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        ), "Should be able to create a topic with colliding chars after deletion.");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDeleteInternalTopic(String quorum) throws Exception {
         // create the offset topic
-        TopicCommand.TopicCommandOptions createOffsetTopicOpts = buildTopicCommandOptionsWithBootstrap("--create",
-            "--partitions", "1",
-            "--replication-factor", "1",
-            "--topic", Topic.GROUP_METADATA_TOPIC_NAME);
-        createAndWaitTopic(createOffsetTopicOpts);
-
+        TestUtils.createTopicWithAdmin(adminClient, Topic.GROUP_METADATA_TOPIC_NAME, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         // Try to delete the Topic.GROUP_METADATA_TOPIC_NAME which is allowed by default.
         // This is a difference between the new and the old command as the old one didn't allow internal topic deletion.
         // If deleting internal topics is not desired, ACLS should be used to control it.
@@ -583,7 +607,7 @@ public void testDeleteInternalTopic(String quorum) throws Exception {
             assertFalse(zkClient().pathExists(deleteOffsetTopicPath), "Delete path for topic shouldn't exist before deletion.");
         }
         topicService.deleteTopic(deleteOffsetTopicOpts);
-        TestUtils.verifyTopicDeletion(zkClientOrNull(), Topic.GROUP_METADATA_TOPIC_NAME, 1, brokers());
+        TestUtils.verifyTopicDeletion(zkClientOrNull(), Topic.GROUP_METADATA_TOPIC_NAME, defaultNumPartitions, brokers());
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -591,7 +615,8 @@ public void testDeleteInternalTopic(String quorum) throws Exception {
     public void testDeleteWhenTopicDoesntExist(String quorum) {
         // delete a topic that does not exist
         TopicCommand.TopicCommandOptions deleteOpts = buildTopicCommandOptionsWithBootstrap("--delete", "--topic", testTopicName);
-        assertThrows(IllegalArgumentException.class, () -> topicService.deleteTopic(deleteOpts));
+        assertThrows(IllegalArgumentException.class, () -> topicService.deleteTopic(deleteOpts),
+                "Expected an exception when trying to delete a topic that does not exist.");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -603,21 +628,23 @@ public void testDeleteWhenTopicDoesntExistWithIfExists(String quorum) throws Exe
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribe(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 2, (short) 2))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partition = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partition, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName));
-        String[] rows = output.split("\n");
-        assertEquals(3, rows.length);
-        assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)));
+        String[] rows = output.split(System.lineSeparator());
+        assertEquals(3, rows.length, "Expected 3 rows in output, got " + rows.length);
+        assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)), "Row does not start with " + testTopicName + ". Row is: " + rows[0]);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeWhenTopicDoesntExist(String quorum) {
         assertThrows(IllegalArgumentException.class,
-            () -> topicService.describeTopic(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName)));
+            () -> topicService.describeTopic(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName)),
+                "Expected an exception when trying to describe a topic that does not exist.");
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -629,10 +656,11 @@ public void testDescribeWhenTopicDoesntExistWithIfExists(String quorum) throws E
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeUnavailablePartitions(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 6, (short) 1))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partitions = 6;
+        short replicationFactor = 1;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         try {
             // check which partition is on broker 0 which we'll kill
             TopicDescription testTopicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName))
@@ -678,9 +706,11 @@ public void testDescribeUnavailablePartitions(String quorum) throws ExecutionExc
 
             // grab the console output and assert
             String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName, "--unavailable-partitions"));
-            String[] rows = output.split("\n");
-            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)));
-            assertTrue(rows[0].contains("Leader: none\tReplicas: 0\tIsr:"));
+            String[] rows = output.split(System.lineSeparator());
+            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)),
+                    "Unexpected Topic " + rows[0] + " received. Expect " + String.format("Topic: %s", testTopicName));
+            assertTrue(rows[0].contains("Leader: none\tReplicas: 0\tIsr:"),
+                    "Rows did not contain 'Leader: none\tReplicas: 0\tIsr:'");
         } finally {
             restartDeadBrokers(false);
         }
@@ -689,10 +719,11 @@ public void testDescribeUnavailablePartitions(String quorum) throws ExecutionExc
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeUnderReplicatedPartitions(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 1, (short) 6))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        int partitions = 1;
+        short replicationFactor = 6;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         try {
             killBroker(0);
             if (isKRaftTest()) {
@@ -701,7 +732,7 @@ public void testDescribeUnderReplicatedPartitions(String quorum) throws Executio
                 TestUtils.waitForPartitionMetadata(aliveBrokers(), testTopicName, 0, org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS);
             }
             String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--under-replicated-partitions"));
-            String[] rows = output.split("\n");
+            String[] rows = output.split(System.lineSeparator());
             assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)), String.format("Unexpected output: %s", rows[0]));
         } finally {
             restartDeadBrokers(false);
@@ -711,13 +742,13 @@ public void testDescribeUnderReplicatedPartitions(String quorum) throws Executio
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeUnderMinIsrPartitions(String quorum) throws ExecutionException, InterruptedException {
-        Map<String, String> configMap = new HashMap<>();
-        configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6");
-
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 1, (short) 6).configs(configMap))).all().get();
-        waitForTopicCreated(testTopicName);
-
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6");
+        int partitions = 1;
+        short replicationFactor = 6;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
         try {
             killBroker(0);
             if (isKRaftTest()) {
@@ -730,8 +761,9 @@ public void testDescribeUnderMinIsrPartitions(String quorum) throws ExecutionExc
                 );
             }
             String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--under-min-isr-partitions"));
-            String[] rows = output.split("\n");
-            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)));
+            String[] rows = output.split(System.lineSeparator());
+            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)),
+                    "Unexpected topic: " + rows[0]);
         } finally {
             restartDeadBrokers(false);
         }
@@ -740,15 +772,11 @@ public void testDescribeUnderMinIsrPartitions(String quorum) throws ExecutionExc
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeUnderReplicatedPartitionsWhenReassignmentIsInProgress(String quorum) throws ExecutionException, InterruptedException {
-        Map<String, String> configMap = new HashMap<>();
-        short replicationFactor = 1;
-        int partitions = 1;
         TopicPartition tp = new TopicPartition(testTopicName, 0);
 
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, partitions, replicationFactor).configs(configMap))
-        ).all().get();
-        waitForTopicCreated(testTopicName);
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         // Produce multiple batches.
         TestUtils.generateAndProduceMessages(brokers(), testTopicName, 10, -1);
@@ -785,19 +813,36 @@ public void testDescribeUnderReplicatedPartitionsWhenReassignmentIsInProgress(St
             () -> "Reassignment didn't add the second node",
             org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS, 100L);
 
+        ensureConsistentKRaftMetadata();
+
         // describe the topic and test if it's under-replicated
         String simpleDescribeOutput = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName));
-        String[] simpleDescribeOutputRows = simpleDescribeOutput.split("\n");
-        assertTrue(simpleDescribeOutputRows[0].startsWith(String.format("Topic: %s", testTopicName)));
-        assertEquals(2, simpleDescribeOutputRows.length);
+        String[] simpleDescribeOutputRows = simpleDescribeOutput.split(System.lineSeparator());
+        assertTrue(simpleDescribeOutputRows[0].startsWith(String.format("Topic: %s", testTopicName)),
+                "Unexpected describe output: " + simpleDescribeOutputRows[0]);
+        assertEquals(2, simpleDescribeOutputRows.length,
+                "Unexpected describe output length: " + simpleDescribeOutputRows.length);
 
         String underReplicatedOutput = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--under-replicated-partitions"));
         assertEquals("", underReplicatedOutput,
             String.format("--under-replicated-partitions shouldn't return anything: '%s'", underReplicatedOutput));
 
-        // Verify reassignment is still ongoing.
-        PartitionReassignment reassignments = adminClient.listPartitionReassignments(Collections.singleton(tp)).reassignments().get().get(tp);
-        assertFalse(reassignments.addingReplicas().isEmpty());
+        int maxRetries = 20;
+        long pause = 100L;
+        long waitTimeMs = maxRetries * pause;
+        AtomicReference<PartitionReassignment> reassignmentsRef = new AtomicReference<>();
+
+        TestUtils.waitUntilTrue(() -> {
+            try {
+                PartitionReassignment tempReassignments = adminClient.listPartitionReassignments(Collections.singleton(tp)).reassignments().get().get(tp);
+                reassignmentsRef.set(tempReassignments);
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException("Error while fetching reassignments", e);
+            }
+            return reassignmentsRef.get() != null;
+        }, () -> "Reassignments did not become non-null within the specified time", waitTimeMs, pause);
+
+        assertFalse(reassignmentsRef.get().addingReplicas().isEmpty());
 
         ToolsTestUtils.removeReplicationThrottleForPartitions(adminClient, brokerIds, Collections.singleton(tp));
         TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
@@ -806,13 +851,14 @@ public void testDescribeUnderReplicatedPartitionsWhenReassignmentIsInProgress(St
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeAtMinIsrPartitions(String quorum) throws ExecutionException, InterruptedException {
-        Map<String, String> configMap = new HashMap<>();
-        configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "4");
-
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic(testTopicName, 1, (short) 6).configs(configMap))).all().get();
-        waitForTopicCreated(testTopicName);
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "4");
 
+        int partitions = 1;
+        short replicationFactor = 6;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
         try {
             killBroker(0);
             killBroker(1);
@@ -828,8 +874,9 @@ public void testDescribeAtMinIsrPartitions(String quorum) throws ExecutionExcept
             }
 
             String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--at-min-isr-partitions"));
-            String[] rows = output.split("\n");
-            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)));
+            String[] rows = output.split(System.lineSeparator());
+            assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)),
+                    "Unexpected output: " + rows[0]);
             assertEquals(1, rows.length);
         } finally {
             restartDeadBrokers(false);
@@ -853,21 +900,32 @@ public void testDescribeUnderMinIsrPartitionsMixed(String quorum) throws Executi
         String offlineTopic = "offline-topic";
         String fullyReplicatedTopic = "fully-replicated-topic";
 
-        Map<String, String> configMap = new HashMap<>();
-        configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6");
+        scala.collection.mutable.HashMap<Object, Seq<Object>> fullyReplicatedReplicaAssignmentMap = new scala.collection.mutable.HashMap<>();
+        fullyReplicatedReplicaAssignmentMap.put(0, JavaConverters.asScalaBufferConverter(Arrays.asList((Object) 1, (Object) 2, (Object) 3)).asScala().toSeq());
 
-        adminClient.createTopics(
-            java.util.Arrays.asList(
-                new NewTopic(underMinIsrTopic, 1, (short) 6).configs(configMap),
-                new NewTopic(notUnderMinIsrTopic, 1, (short) 6),
-                new NewTopic(offlineTopic, Collections.singletonMap(0, Collections.singletonList(0))),
-                new NewTopic(fullyReplicatedTopic, Collections.singletonMap(0, java.util.Arrays.asList(1, 2, 3))))
-        ).all().get();
+        scala.collection.mutable.HashMap<Object, Seq<Object>> offlineReplicaAssignmentMap = new scala.collection.mutable.HashMap<>();
+        offlineReplicaAssignmentMap.put(0, JavaConverters.asScalaBufferConverter(Arrays.asList((Object) 0)).asScala().toSeq());
+
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6");
+
+        int partitions = 1;
+        short replicationFactor = 6;
+        int negativePartition = -1;
+        short negativeReplicationFactor = -1;
+        TestUtils.createTopicWithAdmin(adminClient, underMinIsrTopic, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
+        TestUtils.createTopicWithAdmin(adminClient, notUnderMinIsrTopic, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, offlineTopic, scalaBrokers, scalaControllers, negativePartition, negativeReplicationFactor,
+                offlineReplicaAssignmentMap, new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, fullyReplicatedTopic, scalaBrokers, scalaControllers, negativePartition, negativeReplicationFactor,
+                fullyReplicatedReplicaAssignmentMap, new Properties()
+        );
 
-        waitForTopicCreated(underMinIsrTopic);
-        waitForTopicCreated(notUnderMinIsrTopic);
-        waitForTopicCreated(offlineTopic);
-        waitForTopicCreated(fullyReplicatedTopic);
 
         try {
             killBroker(0);
@@ -881,10 +939,13 @@ public void testDescribeUnderMinIsrPartitionsMixed(String quorum) throws Executi
                     () -> "Timeout waiting for partition metadata propagating to brokers for underMinIsrTopic topic",
                     org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS, 100L);
             }
+            TestUtils.waitForAllReassignmentsToComplete(adminClient, 100L);
             String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--under-min-isr-partitions"));
-            String[] rows = output.split("\n");
-            assertTrue(rows[0].startsWith(String.format("Topic: %s", underMinIsrTopic)));
-            assertTrue(rows[1].startsWith(String.format("\tTopic: %s", offlineTopic)));
+            String[] rows = output.split(System.lineSeparator());
+            assertTrue(rows[0].startsWith(String.format("Topic: %s", underMinIsrTopic)),
+                    "Unexpected output: " + rows[0]);
+            assertTrue(rows[1].startsWith(String.format("\tTopic: %s", offlineTopic)),
+                    "Unexpected output: " + rows[1]);
             assertEquals(2, rows.length);
         } finally {
             restartDeadBrokers(false);
@@ -895,8 +956,14 @@ public void testDescribeUnderMinIsrPartitionsMixed(String quorum) throws Executi
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeReportOverriddenConfigs(String quorum) throws Exception {
         String config = "file.delete.delay.ms=1000";
-        createAndWaitTopic(buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "2",
-            "--replication-factor", "2", "--topic", testTopicName, "--config", config));
+        Properties topicConfig = new Properties();
+        topicConfig.put(TopicConfig.FILE_DELETE_DELAY_MS_CONFIG, "1000");
+
+        int partitions = 2;
+        short replicationFactor = 2;
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), topicConfig
+        );
         String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe"));
         assertTrue(output.contains(config), String.format("Describe output should have contained %s", config));
     }
@@ -904,22 +971,24 @@ public void testDescribeReportOverriddenConfigs(String quorum) throws Exception
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeAndListTopicsWithoutInternalTopics(String quorum) throws Exception {
-        createAndWaitTopic(
-            buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "1", "--replication-factor", "1", "--topic", testTopicName));
-        // create a internal topic
-        createAndWaitTopic(
-            buildTopicCommandOptionsWithBootstrap("--create", "--partitions", "1", "--replication-factor", "1", "--topic", Topic.GROUP_METADATA_TOPIC_NAME));
-
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        TestUtils.createTopicWithAdmin(adminClient, Topic.GROUP_METADATA_TOPIC_NAME, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
         // test describe
         String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--describe", "--exclude-internal"));
         assertTrue(output.contains(testTopicName),
             String.format("Output should have contained %s", testTopicName));
-        assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME));
+        assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME),
+                "Output should not have contained " + Topic.GROUP_METADATA_TOPIC_NAME);
 
         // test list
         output = captureListTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--list", "--exclude-internal"));
-        assertTrue(output.contains(testTopicName));
-        assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME));
+        assertTrue(output.contains(testTopicName), String.format("Output should have contained %s", testTopicName));
+        assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME),
+                "Output should not have contained " + Topic.GROUP_METADATA_TOPIC_NAME);
     }
 
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -936,13 +1005,12 @@ public void testDescribeDoesNotFailWhenListingReassignmentIsUnauthorized(String
 
         topicService = new TopicCommand.TopicService(adminClient);
 
-        adminClient.createTopics(
-                Collections.singletonList(new NewTopic(testTopicName, 1, (short) 1))
-        ).all().get();
-        waitForTopicCreated(testTopicName);
+        TestUtils.createTopicWithAdmin(adminClient, testTopicName, scalaBrokers, scalaControllers, defaultNumPartitions, defaultReplicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
 
         String output = captureDescribeTopicStandardOut(buildTopicCommandOptionsWithBootstrap("--describe", "--topic", testTopicName));
-        String[] rows = output.split("\n");
+        String[] rows = output.split(System.lineSeparator());
         assertEquals(2, rows.length, "Unexpected output: " + output);
         assertTrue(rows[0].startsWith(String.format("Topic: %s", testTopicName)), "Unexpected output: " + rows[0]);
     }
@@ -950,12 +1018,14 @@ public void testDescribeDoesNotFailWhenListingReassignmentIsUnauthorized(String
     @ParameterizedTest(name = ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testCreateWithTopicNameCollision(String quorum) throws ExecutionException, InterruptedException {
-        adminClient.createTopics(
-            Collections.singletonList(new NewTopic("foo_bar", 1, (short) 6))).all().get();
-        waitForTopicCreated("foo_bar");
-
-        assertThrows(InvalidTopicException.class,
-            () -> topicService.createTopic(buildTopicCommandOptionsWithBootstrap("--create", "--topic", "foo.bar")));
+        String topic = "foo_bar";
+        int partitions = 1;
+        short replicationFactor = 6;
+        TestUtils.createTopicWithAdmin(adminClient, topic, scalaBrokers, scalaControllers, partitions, replicationFactor,
+                scala.collection.immutable.Map$.MODULE$.empty(), new Properties()
+        );
+        assertThrows(TopicExistsException.class,
+            () -> topicService.createTopic(buildTopicCommandOptionsWithBootstrap("--create", "--topic", topic)));
     }
 
     private void checkReplicaDistribution(Map<Integer, List<Integer>> assignment,

From 1442862bbd7195b4dde76f9076cf94fcd500d3b9 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Mon, 19 Feb 2024 15:33:37 +0100
Subject: [PATCH 052/258] KAFKA-16009: Fix PlaintextConsumerTest.
 testMaxPollIntervalMsDelayInRevocation (#15383)

The wake-up mechanism in the new consumer is preventing from committing within a rebalance listener callback. The reason is that we are trying to register two wake-uppable actions at the same time.

The fix is to register the wake-uppable action more closely to where we are in fact blocking on it, so that the action is not registered when we execute rebalance listeneners and callback listeners.

Reviewers: Bruno Cadonna <cadonna@apache.org>
---
 .../internals/AsyncKafkaConsumer.java         |  4 +--
 .../internals/AsyncKafkaConsumerTest.java     | 34 ++++++++++++++++++-
 .../kafka/api/PlaintextConsumerTest.scala     |  4 +--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index 31481079cc9e4..28d26a83be6b6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -692,7 +692,6 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
 
         acquireAndEnsureOpen();
         try {
-            wakeupTrigger.setFetchAction(fetchBuffer);
             kafkaConsumerMetrics.recordPollStart(timer.currentTimeMs());
 
             if (subscriptions.hasNoSubscriptionOrUserAssignment()) {
@@ -724,7 +723,6 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
             return ConsumerRecords.empty();
         } finally {
             kafkaConsumerMetrics.recordPollEnd(timer.currentTimeMs());
-            wakeupTrigger.clearTask();
             release();
         }
     }
@@ -1511,6 +1509,7 @@ private Fetch<K, V> pollForFetches(Timer timer) {
         log.trace("Polling for fetches with timeout {}", pollTimeout);
 
         Timer pollTimer = time.timer(pollTimeout);
+        wakeupTrigger.setFetchAction(fetchBuffer);
 
         // Wait a bit for some fetched data to arrive, as there may not be anything immediately available. Note the
         // use of a shorter, dedicated "pollTimer" here which updates "timer" so that calling method (poll) will
@@ -1521,6 +1520,7 @@ private Fetch<K, V> pollForFetches(Timer timer) {
             log.trace("Timeout during fetch", e);
         } finally {
             timer.update(pollTimer.currentTimeMs());
+            wakeupTrigger.clearTask();
         }
 
         return collectFetch();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 691467d29fbd5..575469519823b 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -81,6 +81,7 @@
 
 import java.time.Duration;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -90,6 +91,7 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.SortedSet;
+import java.util.TreeSet;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Future;
@@ -107,6 +109,7 @@
 import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST;
 import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED;
 import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED;
+import static org.apache.kafka.clients.consumer.internals.MembershipManagerImpl.TOPIC_PARTITION_COMPARATOR;
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
@@ -384,7 +387,7 @@ public void testWakeupAfterEmptyFetch() {
         doAnswer(invocation -> {
             consumer.wakeup();
             return Fetch.empty();
-        }).when(fetchCollector).collectFetch(any(FetchBuffer.class));
+        }).doAnswer(invocation -> Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
         Map<TopicPartition, OffsetAndMetadata> offsets = mkMap(mkEntry(tp, new OffsetAndMetadata(1)));
         completeFetchedCommittedOffsetApplicationEventSuccessfully(offsets);
         doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
@@ -419,6 +422,35 @@ public void testWakeupAfterNonEmptyFetch() {
         assertThrows(WakeupException.class, () -> consumer.poll(Duration.ZERO));
     }
 
+    @Test
+    public void testCommitInRebalanceCallback() {
+        consumer = newConsumer();
+        final String topicName = "foo";
+        final int partition = 3;
+        final TopicPartition tp = new TopicPartition(topicName, partition);
+        doAnswer(invocation -> Fetch.empty()).when(fetchCollector).collectFetch(Mockito.any(FetchBuffer.class));
+        SortedSet<TopicPartition> sortedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+        sortedPartitions.add(tp);
+        CompletableBackgroundEvent<Void> e = new ConsumerRebalanceListenerCallbackNeededEvent(ON_PARTITIONS_REVOKED, sortedPartitions);
+        backgroundEventQueue.add(e);
+        completeCommitApplicationEventSuccessfully();
+
+        ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
+            @Override
+            public void onPartitionsRevoked(final Collection<TopicPartition> partitions) {
+                assertDoesNotThrow(() -> consumer.commitSync(mkMap(mkEntry(tp, new OffsetAndMetadata(0)))));
+            }
+
+            @Override
+            public void onPartitionsAssigned(final Collection<TopicPartition> partitions) {
+                // no-op
+            }
+        };
+
+        consumer.subscribe(Collections.singletonList(topicName), listener);
+        consumer.poll(Duration.ZERO);
+    }
+
     @Test
     public void testClearWakeupTriggerAfterPoll() {
         consumer = newConsumer();
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index fbbd7c42a9c5e..553188900cdcf 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -197,10 +197,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertEquals(1, listener.callsToRevoked)
   }
 
-  // TODO: Enable this test for both protocols when the Jira tracking its failure (KAFKA-16009) is fixed. This
-  //       is done by setting the @MethodSource value to "getTestQuorumAndGroupProtocolParametersAll"
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testMaxPollIntervalMsDelayInRevocation(quorum: String, groupProtocol: String): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 5000.toString)
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)

From b71999be95325f6ea54e925cbe5b426425781014 Mon Sep 17 00:00:00 2001
From: Josep Prat <josep.prat@aiven.io>
Date: Mon, 19 Feb 2024 16:54:50 +0100
Subject: [PATCH 053/258] MINOR: Clean up core modules (#15279)

This PR cleans up: metrics, migration, network, raft, security, serializer, tools, utils, and zookeeper package classes

Mark methods and fields private where possible
Annotate public methods and fields
Remove unused classes and methods
Make sure Arrays are not printed with .toString
Optimize minor warnings

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 core/src/main/scala/kafka/Kafka.scala         |  2 +-
 .../kafka/metrics/KafkaMetricsReporter.scala  |  2 +-
 .../metrics/LinuxIoMetricsCollector.scala     | 19 ++--
 .../kafka/migration/MigrationPropagator.scala | 10 +-
 .../scala/kafka/network/RequestChannel.scala  | 94 +++++++++----------
 .../scala/kafka/network/SocketServer.scala    | 44 ++++-----
 .../raft/TimingWheelExpirationService.scala   |  2 +-
 .../kafka/security/CredentialProvider.scala   | 11 ---
 .../security/authorizer/AclAuthorizer.scala   | 30 +++---
 .../kafka/security/authorizer/AclEntry.scala  | 16 ++--
 .../main/scala/kafka/serializer/Decoder.scala |  2 +-
 .../scala/kafka/tools/ConsoleProducer.scala   | 89 +++++++++---------
 .../scala/kafka/tools/DumpLogSegments.scala   | 53 +++++------
 .../main/scala/kafka/tools/StorageTool.scala  | 36 +++----
 .../scala/kafka/tools/TestRaftServer.scala    | 26 ++---
 .../main/scala/kafka/utils/CoreUtils.scala    |  2 +-
 core/src/main/scala/kafka/utils/Logging.scala |  2 +-
 .../main/scala/kafka/zk/AdminZkClient.scala   | 11 +--
 .../main/scala/kafka/zk/KafkaZkClient.scala   | 46 ++++-----
 core/src/main/scala/kafka/zk/ZkData.scala     | 26 ++---
 .../scala/kafka/zk/ZkMigrationClient.scala    |  7 +-
 .../migration/ZkConfigMigrationClient.scala   |  6 +-
 .../ZkDelegationTokenMigrationClient.scala    |  2 +-
 .../kafka/zookeeper/ZooKeeperClient.scala     |  8 +-
 24 files changed, 259 insertions(+), 287 deletions(-)

diff --git a/core/src/main/scala/kafka/Kafka.scala b/core/src/main/scala/kafka/Kafka.scala
index efa6286051494..6b5389fbe3202 100755
--- a/core/src/main/scala/kafka/Kafka.scala
+++ b/core/src/main/scala/kafka/Kafka.scala
@@ -68,7 +68,7 @@ object Kafka extends Logging {
     config.migrationEnabled && config.interBrokerProtocolVersion.isApiForwardingEnabled
 
   private def buildServer(props: Properties): Server = {
-    val config = KafkaConfig.fromProps(props, false)
+    val config = KafkaConfig.fromProps(props, doLog = false)
     if (config.requiresZookeeper) {
       new KafkaServer(
         config,
diff --git a/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala b/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
index 6a50dde69d54f..136bb88b289ce 100755
--- a/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaMetricsReporter.scala
@@ -52,7 +52,7 @@ trait KafkaMetricsReporter {
 }
 
 object KafkaMetricsReporter {
-  val ReporterStarted: AtomicBoolean = new AtomicBoolean(false)
+  private val ReporterStarted: AtomicBoolean = new AtomicBoolean(false)
   private var reporters: ArrayBuffer[KafkaMetricsReporter] = _
 
   def startReporters(verifiableProps: VerifiableProperties): Seq[KafkaMetricsReporter] = {
diff --git a/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala b/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala
index 11602affb1f42..7c56cfca55bb6 100644
--- a/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala
+++ b/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala
@@ -17,7 +17,7 @@
 
 package kafka.metrics
 
-import java.nio.file.{Files, Paths}
+import java.nio.file.{Files, Path, Paths}
 import org.apache.kafka.common.utils.Time
 import org.slf4j.Logger
 
@@ -29,10 +29,10 @@ import scala.jdk.CollectionConverters._
  */
 class LinuxIoMetricsCollector(procRoot: String, val time: Time, val logger: Logger) {
   import LinuxIoMetricsCollector._
-  var lastUpdateMs = -1L
-  var cachedReadBytes = 0L
-  var cachedWriteBytes = 0L
-  val path = Paths.get(procRoot, "self", "io")
+  var lastUpdateMs: Long = -1L
+  var cachedReadBytes:Long = 0L
+  var cachedWriteBytes:Long = 0L
+  val path: Path = Paths.get(procRoot, "self", "io")
 
   def readBytes(): Long = this.synchronized {
     val curMs = time.milliseconds()
@@ -64,7 +64,7 @@ class LinuxIoMetricsCollector(procRoot: String, val time: Time, val logger: Logg
    * write_bytes: 0
    * cancelled_write_bytes: 0
    */
-  def updateValues(now: Long): Boolean = this.synchronized {
+  private def updateValues(now: Long): Boolean = this.synchronized {
     try {
       cachedReadBytes = -1
       cachedWriteBytes = -1
@@ -79,10 +79,9 @@ class LinuxIoMetricsCollector(procRoot: String, val time: Time, val logger: Logg
       lastUpdateMs = now
       true
     } catch {
-      case t: Throwable => {
+      case t: Throwable =>
         logger.warn("Unable to update IO metrics", t)
         false
-      }
     }
   }
 
@@ -97,6 +96,6 @@ class LinuxIoMetricsCollector(procRoot: String, val time: Time, val logger: Logg
 }
 
 object LinuxIoMetricsCollector {
-  val READ_BYTES_PREFIX = "read_bytes: "
-  val WRITE_BYTES_PREFIX = "write_bytes: "
+  private val READ_BYTES_PREFIX = "read_bytes: "
+  private val WRITE_BYTES_PREFIX = "write_bytes: "
 }
diff --git a/core/src/main/scala/kafka/migration/MigrationPropagator.scala b/core/src/main/scala/kafka/migration/MigrationPropagator.scala
index 7bf0fc3ff562e..678bc7a71922e 100644
--- a/core/src/main/scala/kafka/migration/MigrationPropagator.scala
+++ b/core/src/main/scala/kafka/migration/MigrationPropagator.scala
@@ -66,7 +66,7 @@ class MigrationPropagator(
     stateChangeLogger
   )
 
-  val requestBatch = new MigrationPropagatorBatch(
+  private val requestBatch = new MigrationPropagatorBatch(
     config,
     metadataProvider,
     () => _image.features().metadataVersion(),
@@ -103,11 +103,11 @@ class MigrationPropagator(
    * A very expensive function that creates a map with an entry for every partition that exists, from
    * (topic name, partition index) to partition registration.
    */
-  def materializePartitions(topicsImage: TopicsImage): util.Map[TopicPartition, PartitionRegistration] = {
+  private def materializePartitions(topicsImage: TopicsImage): util.Map[TopicPartition, PartitionRegistration] = {
     val result = new util.HashMap[TopicPartition, PartitionRegistration]()
-    topicsImage.topicsById().values().forEach(topic => {
-      topic.partitions().forEach((key, value) => result.put(new TopicPartition(topic.name(), key), value));
-    })
+    topicsImage.topicsById().values().forEach(topic =>
+      topic.partitions().forEach((key, value) => result.put(new TopicPartition(topic.name(), key), value))
+    )
     result
   }
 
diff --git a/core/src/main/scala/kafka/network/RequestChannel.scala b/core/src/main/scala/kafka/network/RequestChannel.scala
index 007049814cb06..679101e7d60da 100644
--- a/core/src/main/scala/kafka/network/RequestChannel.scala
+++ b/core/src/main/scala/kafka/network/RequestChannel.scala
@@ -21,7 +21,7 @@ import java.nio.ByteBuffer
 import java.util.concurrent._
 import com.fasterxml.jackson.databind.JsonNode
 import com.typesafe.scalalogging.Logger
-import com.yammer.metrics.core.Meter
+import com.yammer.metrics.core.{Histogram, Meter}
 import kafka.network
 import kafka.server.{KafkaConfig, RequestLocal}
 import kafka.utils.{Logging, NotNothing, Pool}
@@ -46,11 +46,11 @@ import scala.reflect.ClassTag
 object RequestChannel extends Logging {
   private val requestLogger = Logger("kafka.request.logger")
 
-  val RequestQueueSizeMetric = "RequestQueueSize"
-  val ResponseQueueSizeMetric = "ResponseQueueSize"
+  private val RequestQueueSizeMetric = "RequestQueueSize"
+  private val ResponseQueueSizeMetric = "ResponseQueueSize"
   val ProcessorMetricTag = "processor"
 
-  def isRequestLoggingEnabled: Boolean = requestLogger.underlying.isDebugEnabled
+  private def isRequestLoggingEnabled: Boolean = requestLogger.underlying.isDebugEnabled
 
   sealed trait BaseRequest
   case object ShutdownRequest extends BaseRequest
@@ -87,18 +87,18 @@ object RequestChannel extends Logging {
                 val envelope: Option[RequestChannel.Request] = None) extends BaseRequest {
     // These need to be volatile because the readers are in the network thread and the writers are in the request
     // handler threads or the purgatory threads
-    @volatile var requestDequeueTimeNanos = -1L
-    @volatile var apiLocalCompleteTimeNanos = -1L
-    @volatile var responseCompleteTimeNanos = -1L
-    @volatile var responseDequeueTimeNanos = -1L
-    @volatile var messageConversionsTimeNanos = 0L
-    @volatile var apiThrottleTimeMs = 0L
-    @volatile var temporaryMemoryBytes = 0L
+    @volatile var requestDequeueTimeNanos: Long = -1L
+    @volatile var apiLocalCompleteTimeNanos: Long = -1L
+    @volatile var responseCompleteTimeNanos: Long = -1L
+    @volatile var responseDequeueTimeNanos: Long = -1L
+    @volatile var messageConversionsTimeNanos: Long = 0L
+    @volatile var apiThrottleTimeMs: Long = 0L
+    @volatile var temporaryMemoryBytes: Long = 0L
     @volatile var recordNetworkThreadTimeCallback: Option[Long => Unit] = None
     @volatile var callbackRequestDequeueTimeNanos: Option[Long] = None
     @volatile var callbackRequestCompleteTimeNanos: Option[Long] = None
 
-    val session = new Session(context.principal, context.clientAddress)
+    val session: Session = new Session(context.principal, context.clientAddress)
 
     private val bodyAndSize: RequestAndSize = context.parseRequest(buffer)
 
@@ -110,7 +110,7 @@ object RequestChannel extends Logging {
 
     def header: RequestHeader = context.header
 
-    def sizeOfBodyInBytes: Int = bodyAndSize.size
+    private def sizeOfBodyInBytes: Int = bodyAndSize.size
 
     def sizeInBytes: Int = header.size + sizeOfBodyInBytes
 
@@ -293,7 +293,7 @@ object RequestChannel extends Logging {
       }
     }
 
-    override def toString = s"Request(processor=$processor, " +
+    override def toString: String = s"Request(processor=$processor, " +
       s"connectionId=${context.connectionId}, " +
       s"session=$session, " +
       s"listenerName=${context.listenerName}, " +
@@ -356,8 +356,8 @@ class RequestChannel(val queueSize: Int,
 
   private val requestQueue = new ArrayBlockingQueue[BaseRequest](queueSize)
   private val processors = new ConcurrentHashMap[Int, Processor]()
-  val requestQueueSizeMetricName = metricNamePrefix.concat(RequestQueueSizeMetric)
-  val responseQueueSizeMetricName = metricNamePrefix.concat(ResponseQueueSizeMetric)
+  private val requestQueueSizeMetricName = metricNamePrefix.concat(RequestQueueSizeMetric)
+  private val responseQueueSizeMetricName = metricNamePrefix.concat(ResponseQueueSizeMetric)
   private val callbackQueue = new ArrayBlockingQueue[BaseRequest](queueSize)
 
   metricsGroup.newGauge(requestQueueSizeMetricName, () => requestQueue.size)
@@ -511,24 +511,24 @@ class RequestChannel(val queueSize: Int,
 }
 
 object RequestMetrics {
-  val consumerFetchMetricName = ApiKeys.FETCH.name + "Consumer"
-  val followFetchMetricName = ApiKeys.FETCH.name + "Follower"
-
-  val verifyPartitionsInTxnMetricName = ApiKeys.ADD_PARTITIONS_TO_TXN.name + "Verification"
-
-  val RequestsPerSec = "RequestsPerSec"
-  val DeprecatedRequestsPerSec = "DeprecatedRequestsPerSec"
-  val RequestQueueTimeMs = "RequestQueueTimeMs"
-  val LocalTimeMs = "LocalTimeMs"
-  val RemoteTimeMs = "RemoteTimeMs"
-  val ThrottleTimeMs = "ThrottleTimeMs"
-  val ResponseQueueTimeMs = "ResponseQueueTimeMs"
-  val ResponseSendTimeMs = "ResponseSendTimeMs"
-  val TotalTimeMs = "TotalTimeMs"
-  val RequestBytes = "RequestBytes"
-  val MessageConversionsTimeMs = "MessageConversionsTimeMs"
-  val TemporaryMemoryBytes = "TemporaryMemoryBytes"
-  val ErrorsPerSec = "ErrorsPerSec"
+  val consumerFetchMetricName: String = ApiKeys.FETCH.name + "Consumer"
+  val followFetchMetricName: String = ApiKeys.FETCH.name + "Follower"
+
+  val verifyPartitionsInTxnMetricName: String = ApiKeys.ADD_PARTITIONS_TO_TXN.name + "Verification"
+
+  val RequestsPerSec: String = "RequestsPerSec"
+  val DeprecatedRequestsPerSec: String = "DeprecatedRequestsPerSec"
+  private val RequestQueueTimeMs = "RequestQueueTimeMs"
+  private val LocalTimeMs = "LocalTimeMs"
+  private val RemoteTimeMs = "RemoteTimeMs"
+  private val ThrottleTimeMs = "ThrottleTimeMs"
+  private val ResponseQueueTimeMs = "ResponseQueueTimeMs"
+  private val ResponseSendTimeMs = "ResponseSendTimeMs"
+  private val TotalTimeMs = "TotalTimeMs"
+  private val RequestBytes = "RequestBytes"
+  val MessageConversionsTimeMs: String = "MessageConversionsTimeMs"
+  val TemporaryMemoryBytes: String = "TemporaryMemoryBytes"
+  val ErrorsPerSec: String = "ErrorsPerSec"
 }
 
 private case class DeprecatedRequestRateKey(version: Short, clientInformation: ClientInformation)
@@ -539,34 +539,34 @@ class RequestMetrics(name: String) {
 
   private val metricsGroup = new KafkaMetricsGroup(this.getClass)
 
-  val tags = Map("request" -> name).asJava
-  val requestRateInternal = new Pool[Short, Meter]()
+  val tags: util.Map[String, String] = Map("request" -> name).asJava
+  private val requestRateInternal = new Pool[Short, Meter]()
   private val deprecatedRequestRateInternal = new Pool[DeprecatedRequestRateKey, Meter]()
   // time a request spent in a request queue
-  val requestQueueTimeHist = metricsGroup.newHistogram(RequestQueueTimeMs, true, tags)
+  val requestQueueTimeHist: Histogram = metricsGroup.newHistogram(RequestQueueTimeMs, true, tags)
   // time a request takes to be processed at the local broker
-  val localTimeHist = metricsGroup.newHistogram(LocalTimeMs, true, tags)
+  val localTimeHist: Histogram = metricsGroup.newHistogram(LocalTimeMs, true, tags)
   // time a request takes to wait on remote brokers (currently only relevant to fetch and produce requests)
-  val remoteTimeHist = metricsGroup.newHistogram(RemoteTimeMs, true, tags)
+  val remoteTimeHist: Histogram = metricsGroup.newHistogram(RemoteTimeMs, true, tags)
   // time a request is throttled, not part of the request processing time (throttling is done at the client level
   // for clients that support KIP-219 and by muting the channel for the rest)
-  val throttleTimeHist = metricsGroup.newHistogram(ThrottleTimeMs, true, tags)
+  val throttleTimeHist: Histogram = metricsGroup.newHistogram(ThrottleTimeMs, true, tags)
   // time a response spent in a response queue
-  val responseQueueTimeHist = metricsGroup.newHistogram(ResponseQueueTimeMs, true, tags)
+  val responseQueueTimeHist: Histogram = metricsGroup.newHistogram(ResponseQueueTimeMs, true, tags)
   // time to send the response to the requester
-  val responseSendTimeHist = metricsGroup.newHistogram(ResponseSendTimeMs, true, tags)
-  val totalTimeHist = metricsGroup.newHistogram(TotalTimeMs, true, tags)
+  val responseSendTimeHist: Histogram = metricsGroup.newHistogram(ResponseSendTimeMs, true, tags)
+  val totalTimeHist: Histogram = metricsGroup.newHistogram(TotalTimeMs, true, tags)
   // request size in bytes
-  val requestBytesHist = metricsGroup.newHistogram(RequestBytes, true, tags)
+  val requestBytesHist: Histogram = metricsGroup.newHistogram(RequestBytes, true, tags)
   // time for message conversions (only relevant to fetch and produce requests)
-  val messageConversionsTimeHist =
+  val messageConversionsTimeHist: Option[Histogram] =
     if (name == ApiKeys.FETCH.name || name == ApiKeys.PRODUCE.name)
       Some(metricsGroup.newHistogram(MessageConversionsTimeMs, true, tags))
     else
       None
   // Temporary memory allocated for processing request (only populated for fetch and produce requests)
   // This shows the memory allocated for compression/conversions excluding the actual request size
-  val tempMemoryBytesHist =
+  val tempMemoryBytesHist: Option[Histogram] =
     if (name == ApiKeys.FETCH.name || name == ApiKeys.PRODUCE.name)
       Some(metricsGroup.newHistogram(TemporaryMemoryBytes, true, tags))
     else
@@ -600,7 +600,7 @@ class RequestMetrics(name: String) {
     extendedTags
   }
 
-  class ErrorMeter(name: String, error: Errors) {
+  private class ErrorMeter(name: String, error: Errors) {
     private val tags = Map("request" -> name, "error" -> error.name).asJava
 
     @volatile private var meter: Meter = _
diff --git a/core/src/main/scala/kafka/network/SocketServer.scala b/core/src/main/scala/kafka/network/SocketServer.scala
index 50da12f5da6f1..e90af2e18ca02 100644
--- a/core/src/main/scala/kafka/network/SocketServer.scala
+++ b/core/src/main/scala/kafka/network/SocketServer.scala
@@ -85,7 +85,7 @@ class SocketServer(val config: KafkaConfig,
 
   private val maxQueuedRequests = config.queuedMaxRequests
 
-  protected val nodeId = config.brokerId
+  protected val nodeId: Int = config.brokerId
 
   private val logContext = new LogContext(s"[SocketServer listenerType=${apiVersionManager.listenerType}, nodeId=$nodeId] ")
 
@@ -239,7 +239,7 @@ class SocketServer(val config: KafkaConfig,
     enableFuture
   }
 
-  def createDataPlaneAcceptorAndProcessors(endpoint: EndPoint): Unit = synchronized {
+  private def createDataPlaneAcceptorAndProcessors(endpoint: EndPoint): Unit = synchronized {
     if (stopped) {
       throw new RuntimeException("Can't create new data plane acceptor and processors: SocketServer is stopped.")
     }
@@ -404,13 +404,13 @@ class SocketServer(val config: KafkaConfig,
 object SocketServer {
   val MetricsGroup = "socket-server-metrics"
 
-  val ReconfigurableConfigs = Set(
+  val ReconfigurableConfigs: Set[String] = Set(
     KafkaConfig.MaxConnectionsPerIpProp,
     KafkaConfig.MaxConnectionsPerIpOverridesProp,
     KafkaConfig.MaxConnectionsProp,
     KafkaConfig.MaxConnectionCreationRateProp)
 
-  val ListenerReconfigurableConfigs = Set(KafkaConfig.MaxConnectionsProp, KafkaConfig.MaxConnectionCreationRateProp)
+  val ListenerReconfigurableConfigs: Set[String] = Set(KafkaConfig.MaxConnectionsProp, KafkaConfig.MaxConnectionCreationRateProp)
 
   def closeSocket(
     channel: SocketChannel,
@@ -422,9 +422,9 @@ object SocketServer {
 }
 
 object DataPlaneAcceptor {
-  val ThreadPrefix = "data-plane"
-  val MetricPrefix = ""
-  val ListenerReconfigurableConfigs = Set(KafkaConfig.NumNetworkThreadsProp)
+  val ThreadPrefix: String = "data-plane"
+  val MetricPrefix: String = ""
+  val ListenerReconfigurableConfigs: Set[String] = Set(KafkaConfig.NumNetworkThreadsProp)
 }
 
 class DataPlaneAcceptor(socketServer: SocketServer,
@@ -481,7 +481,7 @@ class DataPlaneAcceptor(socketServer: SocketServer,
    */
   override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
     configs.forEach { (k, v) =>
-      if (reconfigurableConfigs.contains(k)) {
+      if (reconfigurableConfigs().contains(k)) {
         val newValue = v.asInstanceOf[Int]
         val oldValue = processors.length
         if (newValue != oldValue) {
@@ -501,8 +501,8 @@ class DataPlaneAcceptor(socketServer: SocketServer,
    * Reconfigures this instance with the given key-value pairs. The provided
    * map contains all configs including any reconfigurable configs that
    * may have changed since the object was initially configured using
-   * {@link Configurable# configure ( Map )}. This method will only be invoked if
-   * the configs have passed validation using {@link #validateReconfiguration ( Map )}.
+   * [[org.apache.kafka.common.Configurable#configure( Map )]]. This method will only be invoked if
+   * the configs have passed validation using [[validateReconfiguration( Map )]].
    */
   override def reconfigure(configs: util.Map[String, _]): Unit = {
     val newNumNetworkThreads = configs.get(KafkaConfig.NumNetworkThreadsProp).asInstanceOf[Int]
@@ -559,12 +559,6 @@ class ControlPlaneAcceptor(socketServer: SocketServer,
   override def metricPrefix(): String = ControlPlaneAcceptor.MetricPrefix
   override def threadPrefix(): String = ControlPlaneAcceptor.ThreadPrefix
 
-  def processorOpt(): Option[Processor] = {
-    if (processors.isEmpty)
-      None
-    else
-      Some(processors.apply(0))
-  }
 }
 
 /**
@@ -607,8 +601,8 @@ private[kafka] abstract class Acceptor(val socketServer: SocketServer,
     endPoint.port
   } else {
     serverChannel = openServerSocket(endPoint.host, endPoint.port, listenBacklogSize)
-    val newPort = serverChannel.socket().getLocalPort()
-    info(s"Opened wildcard endpoint ${endPoint.host}:${newPort}")
+    val newPort = serverChannel.socket().getLocalPort
+    info(s"Opened wildcard endpoint ${endPoint.host}:$newPort")
     newPort
   }
 
@@ -625,7 +619,7 @@ private[kafka] abstract class Acceptor(val socketServer: SocketServer,
   private var started = false
   private[network] val startedFuture = new CompletableFuture[Void]()
 
-  val thread = KafkaThread.nonDaemon(
+  val thread: KafkaThread = KafkaThread.nonDaemon(
     s"${threadPrefix()}-kafka-socket-acceptor-${endPoint.listenerName}-${endPoint.securityProtocol}-${endPoint.port}",
     this)
 
@@ -859,7 +853,7 @@ private[kafka] abstract class Acceptor(val socketServer: SocketServer,
   }
 
   def newProcessor(id: Int, listenerName: ListenerName, securityProtocol: SecurityProtocol): Processor = {
-    val name = s"${threadPrefix()}-kafka-network-thread-$nodeId-${endPoint.listenerName}-${endPoint.securityProtocol}-${id}"
+    val name = s"${threadPrefix()}-kafka-network-thread-$nodeId-${endPoint.listenerName}-${endPoint.securityProtocol}-$id"
     new Processor(id,
                   time,
                   config.socketRequestMaxBytes,
@@ -882,7 +876,7 @@ private[kafka] abstract class Acceptor(val socketServer: SocketServer,
 }
 
 private[kafka] object Processor {
-  val IdlePercentMetricName = "IdlePercent"
+  private val IdlePercentMetricName = "IdlePercent"
   val NetworkProcessorMetricTag = "networkProcessor"
   val ListenerMetricTag = "listener"
   val ConnectionQueueSize = 20
@@ -920,9 +914,9 @@ private[kafka] class Processor(
 ) extends Runnable with Logging {
   private val metricsGroup = new KafkaMetricsGroup(this.getClass)
 
-  val shouldRun = new AtomicBoolean(true)
+  val shouldRun: AtomicBoolean = new AtomicBoolean(true)
 
-  val thread = KafkaThread.nonDaemon(threadName, this)
+  val thread: KafkaThread = KafkaThread.nonDaemon(threadName, this)
 
   private object ConnectionId {
     def fromString(s: String): Option[ConnectionId] = s.split("-") match {
@@ -957,7 +951,7 @@ private[kafka] class Processor(
     Map(NetworkProcessorMetricTag -> id.toString).asJava
   )
 
-  val expiredConnectionsKilledCount = new CumulativeSum()
+  private val expiredConnectionsKilledCount = new CumulativeSum()
   private val expiredConnectionsKilledCountMetricName = metrics.metricName("expired-connections-killed-count", MetricsGroup, metricTags)
   metrics.addMetric(expiredConnectionsKilledCountMetricName, expiredConnectionsKilledCount)
 
@@ -1114,7 +1108,7 @@ private[kafka] class Processor(
     }
   }
 
-  protected def parseRequestHeader(buffer: ByteBuffer): RequestHeader = {
+  private def parseRequestHeader(buffer: ByteBuffer): RequestHeader = {
     val header = RequestHeader.parse(buffer)
     if (apiVersionManager.isApiEnabled(header.apiKey, header.apiVersion)) {
       header
diff --git a/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala b/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala
index 449574fad4a0c..3c330fb6f2ec2 100644
--- a/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala
+++ b/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala
@@ -25,7 +25,7 @@ import org.apache.kafka.server.util.timer.{Timer, TimerTask}
 object TimingWheelExpirationService {
   private val WorkTimeoutMs: Long = 200L
 
-  class TimerTaskCompletableFuture[T](delayMs: Long) extends TimerTask(delayMs) {
+  private class TimerTaskCompletableFuture[T](delayMs: Long) extends TimerTask(delayMs) {
     val future = new CompletableFuture[T]
     override def run(): Unit = {
       future.completeExceptionally(new TimeoutException(
diff --git a/core/src/main/scala/kafka/security/CredentialProvider.scala b/core/src/main/scala/kafka/security/CredentialProvider.scala
index 85951a526a658..6c58ff0a5dd3f 100644
--- a/core/src/main/scala/kafka/security/CredentialProvider.scala
+++ b/core/src/main/scala/kafka/security/CredentialProvider.scala
@@ -21,8 +21,6 @@ import java.util.{Collection, Properties}
 import org.apache.kafka.clients.admin.{ScramMechanism => AdminScramMechanism}
 import org.apache.kafka.common.security.authenticator.CredentialCache
 import org.apache.kafka.common.security.scram.ScramCredential
-import org.apache.kafka.common.config.ConfigDef
-import org.apache.kafka.common.config.ConfigDef._
 import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramMechanism}
 import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
 
@@ -62,12 +60,3 @@ class CredentialProvider(scramMechanisms: Collection[String], val tokenCache: De
     }
   }
 }
-
-object CredentialProvider {
-  def userCredentialConfigs: ConfigDef = {
-    ScramMechanism.values.foldLeft(new ConfigDef) {
-      (c, m) => c.define(m.mechanismName, Type.STRING, null, Importance.MEDIUM, s"User credentials for SCRAM mechanism ${m.mechanismName}")
-    }
-  }
-}
-
diff --git a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
index 34070d163783e..99642b33c2b90 100644
--- a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -48,23 +48,23 @@ import scala.util.{Failure, Random, Success, Try}
 object AclAuthorizer {
   // Optional override zookeeper cluster configuration where acls will be stored. If not specified,
   // acls will be stored in the same zookeeper where all other kafka broker metadata is stored.
-  val configPrefix = "authorizer."
-  val ZkUrlProp = s"${configPrefix}zookeeper.url"
-  val ZkConnectionTimeOutProp = s"${configPrefix}zookeeper.connection.timeout.ms"
-  val ZkSessionTimeOutProp = s"${configPrefix}zookeeper.session.timeout.ms"
-  val ZkMaxInFlightRequests = s"${configPrefix}zookeeper.max.in.flight.requests"
+  val configPrefix: String = "authorizer."
+  private val ZkUrlProp = s"${configPrefix}zookeeper.url"
+  private val ZkConnectionTimeOutProp = s"${configPrefix}zookeeper.connection.timeout.ms"
+  private val ZkSessionTimeOutProp = s"${configPrefix}zookeeper.session.timeout.ms"
+  private val ZkMaxInFlightRequests = s"${configPrefix}zookeeper.max.in.flight.requests"
 
   // Semi-colon separated list of users that will be treated as super users and will have access to all the resources
   // for all actions from all hosts, defaults to no super users.
-  val SuperUsersProp = "super.users"
+  val SuperUsersProp: String = "super.users"
   // If set to true when no acls are found for a resource, authorizer allows access to everyone. Defaults to false.
-  val AllowEveryoneIfNoAclIsFoundProp = "allow.everyone.if.no.acl.found"
+  val AllowEveryoneIfNoAclIsFoundProp: String = "allow.everyone.if.no.acl.found"
 
   case class VersionedAcls(acls: Set[AclEntry], zkVersion: Int) {
     def exists: Boolean = zkVersion != ZkVersion.UnknownVersion
   }
 
-  class AclSeqs(seqs: Seq[AclEntry]*) {
+  private class AclSeqs(seqs: Seq[AclEntry]*) {
     def find(p: AclEntry => Boolean): Option[AclEntry] = {
       // Lazily iterate through the inner `Seq` elements and stop as soon as we find a match
       val it = seqs.iterator.flatMap(_.find(p))
@@ -75,8 +75,8 @@ object AclAuthorizer {
     def isEmpty: Boolean = !seqs.exists(_.nonEmpty)
   }
 
-  val NoAcls = VersionedAcls(Set.empty, ZkVersion.UnknownVersion)
-  val WildcardHost = "*"
+  val NoAcls: VersionedAcls = VersionedAcls(Set.empty, ZkVersion.UnknownVersion)
+  val WildcardHost: String = "*"
 
   // Orders by resource type, then resource pattern type and finally reverse ordering by name.
   class ResourceOrdering extends Ordering[ResourcePattern] {
@@ -103,7 +103,7 @@ object AclAuthorizer {
     else {
       // start with the base config from the Kafka configuration
       // be sure to force creation since the zkSslClientEnable property in the kafkaConfig could be false
-      val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(kafkaConfig, true)
+      val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(kafkaConfig, forceZkSslClientEnable = true)
       // add in any prefixed overlays
       KafkaConfig.ZkSslConfigToSystemPropertyMap.forKeyValue { (kafkaProp, sysProp) =>
         configMap.get(AclAuthorizer.configPrefix + kafkaProp).foreach { prefixedValue =>
@@ -148,7 +148,7 @@ object AclAuthorizer {
     }
   }
 
-  def getAclsFromZk(zkClient: KafkaZkClient, resource: ResourcePattern): VersionedAcls = {
+  private def getAclsFromZk(zkClient: KafkaZkClient, resource: ResourcePattern): VersionedAcls = {
     zkClient.getVersionedAclsForResource(resource)
   }
 }
@@ -525,7 +525,7 @@ class AclAuthorizer extends Authorizer with Logging {
     if (authorized) AuthorizationResult.ALLOWED else AuthorizationResult.DENIED
   }
 
-  def isSuperUser(principal: KafkaPrincipal): Boolean = {
+  private def isSuperUser(principal: KafkaPrincipal): Boolean = {
     if (superUsers.contains(principal)) {
       authorizerLogger.debug(s"principal = $principal is a super user, allowing operation without checking acls.")
       true
@@ -596,7 +596,7 @@ class AclAuthorizer extends Authorizer with Logging {
     }
   }
 
-  def logAuditMessage(requestContext: AuthorizableRequestContext, action: Action, authorized: Boolean): Unit = {
+  private def logAuditMessage(requestContext: AuthorizableRequestContext, action: Action, authorized: Boolean): Unit = {
     def logMessage: String = {
       val principal = requestContext.principal
       val operation = SecurityUtils.operationName(action.operation)
@@ -753,7 +753,7 @@ class AclAuthorizer extends Authorizer with Logging {
     }
   }
 
-  object AclChangedNotificationHandler extends AclChangeNotificationHandler {
+  private object AclChangedNotificationHandler extends AclChangeNotificationHandler {
     override def processNotification(resource: ResourcePattern): Unit = {
       processAclChangeNotification(resource)
     }
diff --git a/core/src/main/scala/kafka/security/authorizer/AclEntry.scala b/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
index 9e2d49fc883c0..be54746c99117 100644
--- a/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
@@ -33,19 +33,19 @@ object AclEntry {
   val WildcardHost: String = "*"
   val WildcardResource: String = ResourcePattern.WILDCARD_RESOURCE
 
-  val ResourceSeparator = ":"
+  val ResourceSeparator: String = ":"
   val ResourceTypes: Set[ResourceType] = ResourceType.values.toSet
     .filterNot(t => t == ResourceType.UNKNOWN || t == ResourceType.ANY)
   val AclOperations: Set[AclOperation] = AclOperation.values.toSet
     .filterNot(t => t == AclOperation.UNKNOWN || t == AclOperation.ANY)
 
-  val PrincipalKey = "principal"
-  val PermissionTypeKey = "permissionType"
-  val OperationKey = "operation"
-  val HostsKey = "host"
-  val VersionKey = "version"
-  val CurrentVersion = 1
-  val AclsKey = "acls"
+  private val PrincipalKey = "principal"
+  private val PermissionTypeKey = "permissionType"
+  private val OperationKey = "operation"
+  private val HostsKey = "host"
+  val VersionKey: String = "version"
+  val CurrentVersion: Int = 1
+  private val AclsKey = "acls"
 
   def apply(principal: KafkaPrincipal,
             permissionType: AclPermissionType,
diff --git a/core/src/main/scala/kafka/serializer/Decoder.scala b/core/src/main/scala/kafka/serializer/Decoder.scala
index 1ad554100812a..ce166689cfa6e 100644
--- a/core/src/main/scala/kafka/serializer/Decoder.scala
+++ b/core/src/main/scala/kafka/serializer/Decoder.scala
@@ -43,7 +43,7 @@ class DefaultDecoder(props: VerifiableProperties = null) extends Decoder[Array[B
  * an optional property serializer.encoding to control this.
  */
 class StringDecoder(props: VerifiableProperties = null) extends Decoder[String] {
-  val encoding =
+  val encoding: String =
     if (props == null)
       StandardCharsets.UTF_8.name()
     else
diff --git a/core/src/main/scala/kafka/tools/ConsoleProducer.scala b/core/src/main/scala/kafka/tools/ConsoleProducer.scala
index 37a8c807c8052..0f12a49b00afa 100644
--- a/core/src/main/scala/kafka/tools/ConsoleProducer.scala
+++ b/core/src/main/scala/kafka/tools/ConsoleProducer.scala
@@ -21,7 +21,7 @@ import java.io._
 import java.nio.charset.StandardCharsets
 import java.util.Properties
 import java.util.regex.Pattern
-import joptsimple.{OptionException, OptionParser, OptionSet}
+import joptsimple.{OptionException, OptionParser, OptionSet, OptionSpec}
 import kafka.common.MessageReader
 import kafka.utils.Implicits._
 import kafka.utils.{Exit, Logging, ToolsUtils}
@@ -33,6 +33,7 @@ import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 import org.apache.kafka.tools.api.RecordReader
 
+import java.lang
 import scala.annotation.nowarn
 
 @nowarn("cat=deprecation")
@@ -104,7 +105,7 @@ object ConsoleProducer extends Logging {
         System.err.println(e.getMessage)
         Exit.exit(1)
       case e: Exception =>
-        e.printStackTrace
+        e.printStackTrace()
         Exit.exit(1)
     }
   }
@@ -174,81 +175,81 @@ object ConsoleProducer extends Logging {
   }
 
   class ProducerConfig(args: Array[String]) extends CommandDefaultOptions(args) {
-    val topicOpt = parser.accepts("topic", "REQUIRED: The topic id to produce messages to.")
+    val topicOpt: OptionSpec[String] = parser.accepts("topic", "REQUIRED: The topic id to produce messages to.")
       .withRequiredArg
       .describedAs("topic")
       .ofType(classOf[String])
-    val brokerListOpt = parser.accepts("broker-list", "DEPRECATED, use --bootstrap-server instead; ignored if --bootstrap-server is specified.  The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
+    private val brokerListOpt = parser.accepts("broker-list", "DEPRECATED, use --bootstrap-server instead; ignored if --bootstrap-server is specified.  The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
       .withRequiredArg
       .describedAs("broker-list")
       .ofType(classOf[String])
-    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED unless --broker-list(deprecated) is specified. The server(s) to connect to. The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
+    val bootstrapServerOpt: OptionSpec[String]  = parser.accepts("bootstrap-server", "REQUIRED unless --broker-list(deprecated) is specified. The server(s) to connect to. The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
       .requiredUnless("broker-list")
       .withRequiredArg
       .describedAs("server to connect to")
       .ofType(classOf[String])
-    val syncOpt = parser.accepts("sync", "If set message send requests to the brokers are synchronously, one at a time as they arrive.")
-    val compressionCodecOpt = parser.accepts("compression-codec", "The compression codec: either 'none', 'gzip', 'snappy', 'lz4', or 'zstd'." +
+    private val syncOpt = parser.accepts("sync", "If set message send requests to the brokers are synchronously, one at a time as they arrive.")
+    val compressionCodecOpt: OptionSpec[String]  = parser.accepts("compression-codec", "The compression codec: either 'none', 'gzip', 'snappy', 'lz4', or 'zstd'." +
                                                                   "If specified without value, then it defaults to 'gzip'")
                                     .withOptionalArg()
                                     .describedAs("compression-codec")
                                     .ofType(classOf[String])
-    val batchSizeOpt = parser.accepts("batch-size", "Number of messages to send in a single batch if they are not being sent synchronously. "+
+    val batchSizeOpt: OptionSpec[Integer] = parser.accepts("batch-size", "Number of messages to send in a single batch if they are not being sent synchronously. "+
        "please note that this option will be replaced if max-partition-memory-bytes is also set")
       .withRequiredArg
       .describedAs("size")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(16 * 1024)
-    val messageSendMaxRetriesOpt = parser.accepts("message-send-max-retries", "Brokers can fail receiving the message for multiple reasons, " +
+    val messageSendMaxRetriesOpt: OptionSpec[Integer] = parser.accepts("message-send-max-retries", "Brokers can fail receiving the message for multiple reasons, " +
       "and being unavailable transiently is just one of them. This property specifies the number of retries before the producer give up and drop this message. " +
       "This is the option to control `retries` in producer configs.")
       .withRequiredArg
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(3)
-    val retryBackoffMsOpt = parser.accepts("retry-backoff-ms", "Before each retry, the producer refreshes the metadata of relevant topics. " +
+    val retryBackoffMsOpt: OptionSpec[lang.Long] = parser.accepts("retry-backoff-ms", "Before each retry, the producer refreshes the metadata of relevant topics. " +
       "Since leader election takes a bit of time, this property specifies the amount of time that the producer waits before refreshing the metadata. " +
       "This is the option to control `retry.backoff.ms` in producer configs.")
       .withRequiredArg
       .ofType(classOf[java.lang.Long])
       .defaultsTo(100)
-    val sendTimeoutOpt = parser.accepts("timeout", "If set and the producer is running in asynchronous mode, this gives the maximum amount of time" +
+    val sendTimeoutOpt: OptionSpec[lang.Long] = parser.accepts("timeout", "If set and the producer is running in asynchronous mode, this gives the maximum amount of time" +
       " a message will queue awaiting sufficient batch size. The value is given in ms. " +
       "This is the option to control `linger.ms` in producer configs.")
       .withRequiredArg
       .describedAs("timeout_ms")
       .ofType(classOf[java.lang.Long])
       .defaultsTo(1000)
-    val requestRequiredAcksOpt = parser.accepts("request-required-acks", "The required `acks` of the producer requests")
+    val requestRequiredAcksOpt: OptionSpec[String] = parser.accepts("request-required-acks", "The required `acks` of the producer requests")
       .withRequiredArg
       .describedAs("request required acks")
       .ofType(classOf[java.lang.String])
       .defaultsTo("-1")
-    val requestTimeoutMsOpt = parser.accepts("request-timeout-ms", "The ack timeout of the producer requests. Value must be non-negative and non-zero.")
+    val requestTimeoutMsOpt: OptionSpec[Integer] = parser.accepts("request-timeout-ms", "The ack timeout of the producer requests. Value must be non-negative and non-zero.")
       .withRequiredArg
       .describedAs("request timeout ms")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(1500)
-    val metadataExpiryMsOpt = parser.accepts("metadata-expiry-ms",
+    val metadataExpiryMsOpt: OptionSpec[lang.Long] = parser.accepts("metadata-expiry-ms",
       "The period of time in milliseconds after which we force a refresh of metadata even if we haven't seen any leadership changes. " +
         "This is the option to control `metadata.max.age.ms` in producer configs.")
       .withRequiredArg
       .describedAs("metadata expiration interval")
       .ofType(classOf[java.lang.Long])
       .defaultsTo(5*60*1000L)
-    val maxBlockMsOpt = parser.accepts("max-block-ms",
+    val maxBlockMsOpt: OptionSpec[lang.Long] = parser.accepts("max-block-ms",
       "The max time that the producer will block for during a send request.")
       .withRequiredArg
       .describedAs("max block on send")
       .ofType(classOf[java.lang.Long])
       .defaultsTo(60*1000L)
-    val maxMemoryBytesOpt = parser.accepts("max-memory-bytes",
+    val maxMemoryBytesOpt: OptionSpec[lang.Long] = parser.accepts("max-memory-bytes",
       "The total memory used by the producer to buffer records waiting to be sent to the server. " +
         "This is the option to control `buffer.memory` in producer configs.")
       .withRequiredArg
       .describedAs("total memory in bytes")
       .ofType(classOf[java.lang.Long])
       .defaultsTo(32 * 1024 * 1024L)
-    val maxPartitionMemoryBytesOpt = parser.accepts("max-partition-memory-bytes",
+    val maxPartitionMemoryBytesOpt: OptionSpec[Integer] = parser.accepts("max-partition-memory-bytes",
       "The buffer size allocated for a partition. When records are received which are smaller than this size the producer " +
         "will attempt to optimistically group them together until this size is reached. " +
         "This is the option to control `batch.size` in producer configs.")
@@ -256,19 +257,19 @@ object ConsoleProducer extends Logging {
       .describedAs("memory in bytes per partition")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(16 * 1024)
-    val messageReaderOpt = parser.accepts("line-reader", "The class name of the class to use for reading lines from standard in. " +
+    private val messageReaderOpt = parser.accepts("line-reader", "The class name of the class to use for reading lines from standard in. " +
       "By default each line is read as a separate message.")
       .withRequiredArg
       .describedAs("reader_class")
       .ofType(classOf[java.lang.String])
       .defaultsTo(classOf[LineMessageReader].getName)
-    val socketBufferSizeOpt = parser.accepts("socket-buffer-size", "The size of the tcp RECV size. " +
+    val socketBufferSizeOpt: OptionSpec[Integer] = parser.accepts("socket-buffer-size", "The size of the tcp RECV size. " +
       "This is the option to control `send.buffer.bytes` in producer configs.")
       .withRequiredArg
       .describedAs("size")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(1024*100)
-    val propertyOpt = parser.accepts("property",
+    private val propertyOpt = parser.accepts("property",
       """A mechanism to pass user-defined properties in the form key=value to the message reader. This allows custom configuration for a user-defined message reader.
         |Default properties include:
         | parse.key=false
@@ -291,15 +292,15 @@ object ConsoleProducer extends Logging {
       .withRequiredArg
       .describedAs("prop")
       .ofType(classOf[String])
-    val readerConfigOpt = parser.accepts("reader-config", s"Config properties file for the message reader. Note that $propertyOpt takes precedence over this config.")
+    val readerConfigOpt: OptionSpec[String] = parser.accepts("reader-config", s"Config properties file for the message reader. Note that $propertyOpt takes precedence over this config.")
       .withRequiredArg
       .describedAs("config file")
       .ofType(classOf[String])
-    val producerPropertyOpt = parser.accepts("producer-property", "A mechanism to pass user-defined properties in the form key=value to the producer. ")
+    private val producerPropertyOpt = parser.accepts("producer-property", "A mechanism to pass user-defined properties in the form key=value to the producer. ")
             .withRequiredArg
             .describedAs("producer_prop")
             .ofType(classOf[String])
-    val producerConfigOpt = parser.accepts("producer.config", s"Producer config properties file. Note that $producerPropertyOpt takes precedence over this config.")
+    val producerConfigOpt: OptionSpec[String]  = parser.accepts("producer.config", s"Producer config properties file. Note that $producerPropertyOpt takes precedence over this config.")
       .withRequiredArg
       .describedAs("config file")
       .ofType(classOf[String])
@@ -310,24 +311,24 @@ object ConsoleProducer extends Logging {
 
     CommandLineUtils.checkRequiredArgs(parser, options, topicOpt)
 
-    val topic = options.valueOf(topicOpt)
+    val topic: String = options.valueOf(topicOpt)
 
-    val bootstrapServer = options.valueOf(bootstrapServerOpt)
-    val brokerList = options.valueOf(brokerListOpt)
+    val bootstrapServer: String = options.valueOf(bootstrapServerOpt)
+    val brokerList: String = options.valueOf(brokerListOpt)
 
-    val brokerHostsAndPorts = options.valueOf(if (options.has(bootstrapServerOpt)) bootstrapServerOpt else brokerListOpt)
+    val brokerHostsAndPorts: String = options.valueOf(if (options.has(bootstrapServerOpt)) bootstrapServerOpt else brokerListOpt)
     ToolsUtils.validatePortOrDie(parser, brokerHostsAndPorts)
 
-    val sync = options.has(syncOpt)
-    val compressionCodecOptionValue = options.valueOf(compressionCodecOpt)
-    val compressionCodec = if (options.has(compressionCodecOpt))
+    val sync: Boolean = options.has(syncOpt)
+    private val compressionCodecOptionValue = options.valueOf(compressionCodecOpt)
+    val compressionCodec: String = if (options.has(compressionCodecOpt))
                              if (compressionCodecOptionValue == null || compressionCodecOptionValue.isEmpty)
                                CompressionType.GZIP.name
                              else compressionCodecOptionValue
                            else CompressionType.NONE.name
-    val readerClass = options.valueOf(messageReaderOpt)
-    val cmdLineProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(propertyOpt))
-    val extraProducerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(producerPropertyOpt))
+    val readerClass: String = options.valueOf(messageReaderOpt)
+    val cmdLineProps: Properties = CommandLineUtils.parseKeyValueArgs(options.valuesOf(propertyOpt))
+    val extraProducerProps: Properties = CommandLineUtils.parseKeyValueArgs(options.valuesOf(producerPropertyOpt))
 
     def tryParse(parser: OptionParser, args: Array[String]): OptionSet = {
       try
@@ -341,17 +342,17 @@ object ConsoleProducer extends Logging {
 
   class LineMessageReader extends RecordReader {
     var topic: String = _
-    var parseKey = false
-    var keySeparator = "\t"
-    var parseHeaders = false
-    var headersDelimiter = "\t"
-    var headersSeparator = ","
-    var headersKeySeparator = ":"
-    var ignoreError = false
-    var lineNumber = 0
-    var printPrompt = System.console != null
-    var headersSeparatorPattern: Pattern = _
-    var nullMarker: String = _
+    var parseKey: Boolean = false
+    var keySeparator: String = "\t"
+    var parseHeaders: Boolean = false
+    private var headersDelimiter = "\t"
+    var headersSeparator: String = ","
+    private var headersKeySeparator = ":"
+    private var ignoreError = false
+    private var lineNumber = 0
+    private val printPrompt = System.console != null
+    private var headersSeparatorPattern: Pattern = _
+    private var nullMarker: String = _
 
     override def configure(props: java.util.Map[String, _]): Unit = {
       topic = props.get("topic").toString
diff --git a/core/src/main/scala/kafka/tools/DumpLogSegments.scala b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
index 8a37e65922078..8aa5995612cc0 100755
--- a/core/src/main/scala/kafka/tools/DumpLogSegments.scala
+++ b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
@@ -347,9 +347,9 @@ object DumpLogSegments {
   }
 
   class TimeIndexDumpErrors {
-    val misMatchesForTimeIndexFilesMap = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
-    val outOfOrderTimestamp = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
-    val shallowOffsetNotFound = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
+    val misMatchesForTimeIndexFilesMap: mutable.Map[String, ArrayBuffer[(Long, Long)]] = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
+    val outOfOrderTimestamp: mutable.Map[String, ArrayBuffer[(Long, Long)]] = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
+    val shallowOffsetNotFound: mutable.Map[String, ArrayBuffer[(Long, Long)]] = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
 
     def recordMismatchTimeIndex(file: File, indexTimestamp: Long, logTimestamp: Long): Unit = {
       val misMatchesSeq = misMatchesForTimeIndexFilesMap.getOrElse(file.getAbsolutePath, new ArrayBuffer[(Long, Long)]())
@@ -374,21 +374,19 @@ object DumpLogSegments {
 
     def printErrors(): Unit = {
       misMatchesForTimeIndexFilesMap.foreach {
-        case (fileName, listOfMismatches) => {
+        case (fileName, listOfMismatches) =>
           System.err.println("Found timestamp mismatch in :" + fileName)
           listOfMismatches.foreach(m => {
             System.err.println("  Index timestamp: %d, log timestamp: %d".format(m._1, m._2))
           })
-        }
       }
 
       outOfOrderTimestamp.foreach {
-        case (fileName, outOfOrderTimestamps) => {
+        case (fileName, outOfOrderTimestamps) =>
           System.err.println("Found out of order timestamp in :" + fileName)
           outOfOrderTimestamps.foreach(m => {
             System.err.println("  Index timestamp: %d, Previously indexed timestamp: %d".format(m._1, m._2))
           })
-        }
       }
 
       shallowOffsetNotFound.values.foreach { listOfShallowOffsetNotFound =>
@@ -413,7 +411,7 @@ object DumpLogSegments {
   }
 
   private class ClusterMetadataLogMessageParser extends MessageParser[String, String] {
-    val metadataRecordSerde = new MetadataRecordSerde()
+    private val metadataRecordSerde = new MetadataRecordSerde()
 
     override def parse(record: Record): (Option[String], Option[String]) = {
       val output = try {
@@ -425,11 +423,10 @@ object DumpLogSegments {
         json.set("version", new IntNode(messageAndVersion.version()))
         json.set("data", MetadataJsonConverters.writeJson(
           messageAndVersion.message(), messageAndVersion.version()))
-        json.toString()
+        json.toString
       } catch {
-        case e: Throwable => {
+        case e: Throwable =>
           s"Error at ${record.offset}, skipping. ${e.getMessage}"
-        }
       }
       // No keys for metadata records
       (None, Some(output))
@@ -437,39 +434,39 @@ object DumpLogSegments {
   }
 
   private class DumpLogSegmentsOptions(args: Array[String]) extends CommandDefaultOptions(args) {
-    val printOpt = parser.accepts("print-data-log", "if set, printing the messages content when dumping data logs. Automatically set if any decoder option is specified.")
-    val verifyOpt = parser.accepts("verify-index-only", "if set, just verify the index log without printing its content.")
-    val indexSanityOpt = parser.accepts("index-sanity-check", "if set, just checks the index sanity without printing its content. " +
+    private val printOpt = parser.accepts("print-data-log", "if set, printing the messages content when dumping data logs. Automatically set if any decoder option is specified.")
+    private val verifyOpt = parser.accepts("verify-index-only", "if set, just verify the index log without printing its content.")
+    private val indexSanityOpt = parser.accepts("index-sanity-check", "if set, just checks the index sanity without printing its content. " +
       "This is the same check that is executed on broker startup to determine if an index needs rebuilding or not.")
-    val filesOpt = parser.accepts("files", "REQUIRED: The comma separated list of data and index log files to be dumped.")
+    private val filesOpt = parser.accepts("files", "REQUIRED: The comma separated list of data and index log files to be dumped.")
       .withRequiredArg
       .describedAs("file1, file2, ...")
       .ofType(classOf[String])
-    val maxMessageSizeOpt = parser.accepts("max-message-size", "Size of largest message.")
+    private val maxMessageSizeOpt = parser.accepts("max-message-size", "Size of largest message.")
       .withRequiredArg
       .describedAs("size")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(5 * 1024 * 1024)
-    val maxBytesOpt = parser.accepts("max-bytes", "Limit the amount of total batches read in bytes avoiding reading the whole .log file(s).")
+    private val maxBytesOpt = parser.accepts("max-bytes", "Limit the amount of total batches read in bytes avoiding reading the whole .log file(s).")
        .withRequiredArg
        .describedAs("size")
        .ofType(classOf[java.lang.Integer])
        .defaultsTo(Integer.MAX_VALUE)
-    val deepIterationOpt = parser.accepts("deep-iteration", "if set, uses deep instead of shallow iteration. Automatically set if print-data-log is enabled.")
-    val valueDecoderOpt = parser.accepts("value-decoder-class", "if set, used to deserialize the messages. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
+    private val deepIterationOpt = parser.accepts("deep-iteration", "if set, uses deep instead of shallow iteration. Automatically set if print-data-log is enabled.")
+    private val valueDecoderOpt = parser.accepts("value-decoder-class", "if set, used to deserialize the messages. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
       .withOptionalArg()
       .ofType(classOf[java.lang.String])
       .defaultsTo("kafka.serializer.StringDecoder")
-    val keyDecoderOpt = parser.accepts("key-decoder-class", "if set, used to deserialize the keys. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
+    private val keyDecoderOpt = parser.accepts("key-decoder-class", "if set, used to deserialize the keys. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
       .withOptionalArg()
       .ofType(classOf[java.lang.String])
       .defaultsTo("kafka.serializer.StringDecoder")
-    val offsetsOpt = parser.accepts("offsets-decoder", "if set, log data will be parsed as offset data from the " +
+    private val offsetsOpt = parser.accepts("offsets-decoder", "if set, log data will be parsed as offset data from the " +
       "__consumer_offsets topic.")
-    val transactionLogOpt = parser.accepts("transaction-log-decoder", "if set, log data will be parsed as " +
+    private val transactionLogOpt = parser.accepts("transaction-log-decoder", "if set, log data will be parsed as " +
       "transaction metadata from the __transaction_state topic.")
-    val clusterMetadataOpt = parser.accepts("cluster-metadata-decoder", "if set, log data will be parsed as cluster metadata records.")
-    val skipRecordMetadataOpt = parser.accepts("skip-record-metadata", "whether to skip printing metadata for each record.")
+    private val clusterMetadataOpt = parser.accepts("cluster-metadata-decoder", "if set, log data will be parsed as cluster metadata records.")
+    private val skipRecordMetadataOpt = parser.accepts("skip-record-metadata", "whether to skip printing metadata for each record.")
     options = parser.parse(args : _*)
 
     def messageParser: MessageParser[_, _] =
@@ -492,13 +489,13 @@ object DumpLogSegments {
       options.has(valueDecoderOpt) ||
       options.has(keyDecoderOpt)
 
-    lazy val skipRecordMetadata = options.has(skipRecordMetadataOpt)
+    lazy val skipRecordMetadata: Boolean = options.has(skipRecordMetadataOpt)
     lazy val isDeepIteration: Boolean = options.has(deepIterationOpt) || shouldPrintDataLog
     lazy val verifyOnly: Boolean = options.has(verifyOpt)
     lazy val indexSanityOnly: Boolean = options.has(indexSanityOpt)
-    lazy val files = options.valueOf(filesOpt).split(",")
-    lazy val maxMessageSize = options.valueOf(maxMessageSizeOpt).intValue()
-    lazy val maxBytes = options.valueOf(maxBytesOpt).intValue()
+    lazy val files: Array[String] = options.valueOf(filesOpt).split(",")
+    lazy val maxMessageSize: Int = options.valueOf(maxMessageSizeOpt).intValue()
+    lazy val maxBytes: Int = options.valueOf(maxBytesOpt).intValue()
 
     def checkArgs(): Unit = CommandLineUtils.checkRequiredArgs(parser, options, filesOpt)
 
diff --git a/core/src/main/scala/kafka/tools/StorageTool.scala b/core/src/main/scala/kafka/tools/StorageTool.scala
index 30d836370a9ac..d328d1df8718f 100644
--- a/core/src/main/scala/kafka/tools/StorageTool.scala
+++ b/core/src/main/scala/kafka/tools/StorageTool.scala
@@ -64,11 +64,11 @@ object StorageTool extends Logging {
           if (!metadataVersion.isKRaftSupported) {
             throw new TerseFailure(s"Must specify a valid KRaft metadata version of at least 3.0.")
           }
-          if (!metadataVersion.isProduction()) {
+          if (!metadataVersion.isProduction) {
             if (config.get.unstableMetadataVersionsEnabled) {
-              System.out.println(s"WARNING: using pre-production metadata version ${metadataVersion}.")
+              System.out.println(s"WARNING: using pre-production metadata version $metadataVersion.")
             } else {
-              throw new TerseFailure(s"Metadata version ${metadataVersion} is not ready for production use yet.")
+              throw new TerseFailure(s"Metadata version $metadataVersion is not ready for production use yet.")
             }
           }
           val metaProperties = new MetaProperties.Builder().
@@ -78,8 +78,8 @@ object StorageTool extends Logging {
             build()
           val metadataRecords : ArrayBuffer[ApiMessageAndVersion] = ArrayBuffer()
           getUserScramCredentialRecords(namespace).foreach(userScramCredentialRecords => {
-            if (!metadataVersion.isScramSupported()) {
-              throw new TerseFailure(s"SCRAM is only supported in metadataVersion IBP_3_5_IV2 or later.");
+            if (!metadataVersion.isScramSupported) {
+              throw new TerseFailure(s"SCRAM is only supported in metadataVersion IBP_3_5_IV2 or later.")
             }
             for (record <- userScramCredentialRecords) {
               metadataRecords.append(new ApiMessageAndVersion(record, 0.toShort))
@@ -151,7 +151,7 @@ object StorageTool extends Logging {
     directories.toSeq
   }
 
-  def configToSelfManagedMode(config: KafkaConfig): Boolean = config.processRoles.nonEmpty
+  private def configToSelfManagedMode(config: KafkaConfig): Boolean = config.processRoles.nonEmpty
 
   def getMetadataVersion(
     namespace: Namespace,
@@ -167,7 +167,7 @@ object StorageTool extends Logging {
       .getOrElse(defaultValue)
   }
 
-  def getUserScramCredentialRecord(
+  private def getUserScramCredentialRecord(
     mechanism: String,
     config: String
   ) : UserScramCredentialRecord = {
@@ -242,7 +242,7 @@ object StorageTool extends Logging {
     val saltedPassword = getSaltedPassword(argMap, scramMechanism, salt, iterations)
 
     val myrecord = try {
-      val formatter = new ScramFormatter(scramMechanism);
+      val formatter = new ScramFormatter(scramMechanism)
 
       new UserScramCredentialRecord()
            .setName(name)
@@ -348,8 +348,8 @@ object StorageTool extends Logging {
 
       prevMetadata.foreach { prev =>
         val sortedOutput = new util.TreeMap[String, String]()
-        prev.toProperties().entrySet().forEach(e => sortedOutput.put(e.getKey.toString, e.getValue.toString))
-        stream.println(s"Found metadata: ${sortedOutput}")
+        prev.toProperties.entrySet.forEach(e => sortedOutput.put(e.getKey.toString, e.getValue.toString))
+        stream.println(s"Found metadata: $sortedOutput")
         stream.println("")
       }
 
@@ -375,7 +375,7 @@ object StorageTool extends Logging {
     val metadataRecords = new util.ArrayList[ApiMessageAndVersion]
     metadataRecords.add(new ApiMessageAndVersion(new FeatureLevelRecord().
                         setName(MetadataVersion.FEATURE_NAME).
-                        setFeatureLevel(metadataVersion.featureLevel()), 0.toShort));
+                        setFeatureLevel(metadataVersion.featureLevel()), 0.toShort))
 
     metadataOptionalArguments.foreach { metadataArguments =>
       for (record <- metadataArguments) metadataRecords.add(record)
@@ -427,21 +427,21 @@ object StorageTool extends Logging {
       throw new TerseFailure("No log directories found in the configuration.")
     }
     val loader = new MetaPropertiesEnsemble.Loader()
-    directories.foreach(loader.addLogDir(_))
+    directories.foreach(loader.addLogDir)
     val metaPropertiesEnsemble = loader.load()
     metaPropertiesEnsemble.verify(metaProperties.clusterId(), metaProperties.nodeId(),
       util.EnumSet.noneOf(classOf[VerificationFlag]))
 
-    System.out.println(s"metaPropertiesEnsemble=${metaPropertiesEnsemble}")
+    System.out.println(s"metaPropertiesEnsemble=$metaPropertiesEnsemble")
     val copier = new MetaPropertiesEnsemble.Copier(metaPropertiesEnsemble)
     if (!(ignoreFormatted || copier.logDirProps().isEmpty)) {
       val firstLogDir = copier.logDirProps().keySet().iterator().next()
-      throw new TerseFailure(s"Log directory ${firstLogDir} is already formatted. " +
+      throw new TerseFailure(s"Log directory $firstLogDir is already formatted. " +
         "Use --ignore-formatted to ignore this directory and format the others.")
     }
     if (!copier.errorLogDirs().isEmpty) {
       val firstLogDir = copier.errorLogDirs().iterator().next()
-      throw new TerseFailure(s"I/O error trying to read log directory ${firstLogDir}.")
+      throw new TerseFailure(s"I/O error trying to read log directory $firstLogDir.")
     }
     if (metaPropertiesEnsemble.emptyLogDirs().isEmpty) {
       stream.println("All of the log directories are already formatted.")
@@ -450,14 +450,14 @@ object StorageTool extends Logging {
         copier.setLogDirProps(logDir, new MetaProperties.Builder(metaProperties).
           setDirectoryId(copier.generateValidDirectoryId()).
           build())
-        copier.setPreWriteHandler((logDir, isNew, metaProperties) => {
-          stream.println(s"Formatting ${logDir} with metadata.version ${metadataVersion}.")
+        copier.setPreWriteHandler((logDir, _, _) => {
+          stream.println(s"Formatting $logDir with metadata.version $metadataVersion.")
           Files.createDirectories(Paths.get(logDir))
           val bootstrapDirectory = new BootstrapDirectory(logDir, Optional.empty())
           bootstrapDirectory.writeBinaryFile(bootstrapMetadata)
         })
         copier.setWriteErrorHandler((logDir, e) => {
-          throw new TerseFailure(s"Error while writing meta.properties file ${logDir}: ${e.getMessage}")
+          throw new TerseFailure(s"Error while writing meta.properties file $logDir: ${e.getMessage}")
         })
         copier.writeLogDirChanges()
       })
diff --git a/core/src/main/scala/kafka/tools/TestRaftServer.scala b/core/src/main/scala/kafka/tools/TestRaftServer.scala
index f2cdc26723a74..420b935dc299e 100644
--- a/core/src/main/scala/kafka/tools/TestRaftServer.scala
+++ b/core/src/main/scala/kafka/tools/TestRaftServer.scala
@@ -19,7 +19,7 @@ package kafka.tools
 
 import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
 import java.util.concurrent.{CompletableFuture, CountDownLatch, LinkedBlockingDeque, TimeUnit}
-import joptsimple.OptionException
+import joptsimple.{OptionException, OptionSpec}
 import kafka.network.{DataPlaneAcceptor, SocketServer}
 import kafka.raft.{KafkaRaftManager, RaftManager}
 import kafka.security.CredentialProvider
@@ -68,7 +68,7 @@ class TestRaftServer(
   var credentialProvider: CredentialProvider = _
   var tokenCache: DelegationTokenCache = _
   var dataPlaneRequestHandlerPool: KafkaRequestHandlerPool = _
-  var workloadGenerator: RaftWorkloadGenerator = _
+  private var workloadGenerator: RaftWorkloadGenerator = _
   var raftManager: KafkaRaftManager[Array[Byte]] = _
 
   def startup(): Unit = {
@@ -142,7 +142,7 @@ class TestRaftServer(
     shutdownLatch.await()
   }
 
-  class RaftWorkloadGenerator(
+  private class RaftWorkloadGenerator(
     raftManager: RaftManager[Array[Byte]],
     time: Time,
     recordsPerSec: Int,
@@ -150,12 +150,12 @@ class TestRaftServer(
   ) extends ShutdownableThread("raft-workload-generator")
     with RaftClient.Listener[Array[Byte]] {
 
-    sealed trait RaftEvent
-    case class HandleClaim(epoch: Int) extends RaftEvent
-    case object HandleResign extends RaftEvent
-    case class HandleCommit(reader: BatchReader[Array[Byte]]) extends RaftEvent
-    case class HandleSnapshot(reader: SnapshotReader[Array[Byte]]) extends RaftEvent
-    case object Shutdown extends RaftEvent
+    private sealed trait RaftEvent
+    private case class HandleClaim(epoch: Int) extends RaftEvent
+    private case object HandleResign extends RaftEvent
+    private case class HandleCommit(reader: BatchReader[Array[Byte]]) extends RaftEvent
+    private case class HandleSnapshot(reader: SnapshotReader[Array[Byte]]) extends RaftEvent
+    private case object Shutdown extends RaftEvent
 
     private val eventQueue = new LinkedBlockingDeque[RaftEvent]()
     private val stats = new WriteStats(metrics, time, printIntervalMs = 5000)
@@ -411,20 +411,20 @@ object TestRaftServer extends Logging {
     }
   }
 
-  class TestRaftServerOptions(args: Array[String]) extends CommandDefaultOptions(args) {
-    val configOpt = parser.accepts("config", "Required configured file")
+  private class TestRaftServerOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val configOpt: OptionSpec[String] = parser.accepts("config", "Required configured file")
       .withRequiredArg
       .describedAs("filename")
       .ofType(classOf[String])
 
-    val throughputOpt = parser.accepts("throughput",
+    val throughputOpt: OptionSpec[Int] = parser.accepts("throughput",
       "The number of records per second the leader will write to the metadata topic")
       .withRequiredArg
       .describedAs("records/sec")
       .ofType(classOf[Int])
       .defaultsTo(5000)
 
-    val recordSizeOpt = parser.accepts("record-size", "The size of each record")
+    val recordSizeOpt: OptionSpec[Int] = parser.accepts("record-size", "The size of each record")
       .withRequiredArg
       .describedAs("size in bytes")
       .ofType(classOf[Int])
diff --git a/core/src/main/scala/kafka/utils/CoreUtils.scala b/core/src/main/scala/kafka/utils/CoreUtils.scala
index 6af1f3594e2e7..0b445ed1a3cd5 100755
--- a/core/src/main/scala/kafka/utils/CoreUtils.scala
+++ b/core/src/main/scala/kafka/utils/CoreUtils.scala
@@ -159,7 +159,7 @@ object CoreUtils {
     listenerListToEndPoints(listeners, securityProtocolMap, requireDistinctPorts = true)
   }
 
-  def checkDuplicateListenerPorts(endpoints: Seq[EndPoint], listeners: String): Unit = {
+  private def checkDuplicateListenerPorts(endpoints: Seq[EndPoint], listeners: String): Unit = {
     val distinctPorts = endpoints.map(_.port).distinct
     require(distinctPorts.size == endpoints.map(_.port).size, s"Each listener must have a different port, listeners: $listeners")
   }
diff --git a/core/src/main/scala/kafka/utils/Logging.scala b/core/src/main/scala/kafka/utils/Logging.scala
index 0221821e7473a..dd83e90336099 100755
--- a/core/src/main/scala/kafka/utils/Logging.scala
+++ b/core/src/main/scala/kafka/utils/Logging.scala
@@ -40,7 +40,7 @@ private object Logging {
 
 trait Logging {
 
-  protected lazy val logger = Logger(LoggerFactory.getLogger(loggerName))
+  protected lazy val logger: Logger = Logger(LoggerFactory.getLogger(loggerName))
 
   protected var logIdent: String = _
 
diff --git a/core/src/main/scala/kafka/zk/AdminZkClient.scala b/core/src/main/scala/kafka/zk/AdminZkClient.scala
index e023503275b03..72925a036b10b 100644
--- a/core/src/main/scala/kafka/zk/AdminZkClient.scala
+++ b/core/src/main/scala/kafka/zk/AdminZkClient.scala
@@ -452,7 +452,7 @@ class AdminZkClient(zkClient: KafkaZkClient,
    * @param ip ip for which configs are being validated
    * @param configs properties to validate for the IP
    */
-  def validateIpConfig(ip: String, configs: Properties): Unit = {
+  private def validateIpConfig(ip: String, configs: Properties): Unit = {
     if (!DynamicConfig.Ip.isValidIpEntity(ip))
       throw new AdminOperationException(s"$ip is not a valid IP or resolvable host.")
     DynamicConfig.Ip.validate(configs)
@@ -528,7 +528,7 @@ class AdminZkClient(zkClient: KafkaZkClient,
     * only verifies that the provided config does not contain any static configs.
     * @param configs configs to validate
     */
-  def validateBrokerConfig(configs: Properties): Unit = {
+  private def validateBrokerConfig(configs: Properties): Unit = {
     DynamicConfig.Broker.validate(configs)
   }
 
@@ -561,13 +561,6 @@ class AdminZkClient(zkClient: KafkaZkClient,
     zkClient.getEntityConfigs(rootEntityType, sanitizedEntityName)
   }
 
-  /**
-   * Gets all topic configs
-   * @return The successfully gathered configs of all topics
-   */
-  def getAllTopicConfigs(): Map[String, Properties] =
-    zkClient.getAllTopicsInCluster().map(topic => (topic, fetchEntityConfig(ConfigType.TOPIC, topic))).toMap
-
   /**
    * Gets all the entity configs for a given entityType
    * @param entityType entityType for which configs are being fetched
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index c38f866fb0e20..26c9453bef211 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -235,7 +235,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
           case Code.OK =>
             info(s"Successfully registered KRaft controller $kraftControllerId with ZK epoch $newControllerEpoch")
             // First op is always SetData on /controller_epoch
-            val setDataResult = response.zkOpResults(0).rawOpResult.asInstanceOf[SetDataResult]
+            val setDataResult = response.zkOpResults.head.rawOpResult.asInstanceOf[SetDataResult]
             SuccessfulRegistrationResult(newControllerEpoch, setDataResult.getStat.getVersion)
           case Code.BADVERSION =>
             info(s"The ZK controller epoch changed $failureSuffix")
@@ -590,7 +590,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 
   /**
    * Checks the topic existence
-   * @param topicName
+   * @param topicName the name of the topic to check
    * @return true if topic exists else false
    */
   def topicExists(topicName: String): Boolean = {
@@ -646,7 +646,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   def setTopicAssignment(topic: String,
                          topicId: Option[Uuid],
                          assignment: Map[TopicPartition, ReplicaAssignment],
-                         expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion) = {
+                         expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion): Unit = {
     val setDataResponse = setTopicAssignmentRaw(topic, topicId, assignment, expectedControllerEpochZkVersion)
     setDataResponse.maybeThrow()
   }
@@ -893,7 +893,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 
   /**
    * Gets all the child nodes at a given zk node path
-   * @param path
+   * @param path the path to check
    * @return list of child node names
    */
   def getChildren(path : String): Seq[String] = {
@@ -958,7 +958,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 
   /**
    * Checks if topic is marked for deletion
-   * @param topic
+   * @param topic the name of the topic to check
    * @return true if topic is marked for deletion, else false
    */
   def isTopicMarkedForDeletion(topic: String): Boolean = {
@@ -1046,7 +1046,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
    * @param reassignment the reassignment to set on the reassignment znode.
    * @throws KeeperException if there is an error while creating the znode.
    */
-  def createPartitionReassignment(reassignment: Map[TopicPartition, Seq[Int]])  = {
+  def createPartitionReassignment(reassignment: Map[TopicPartition, Seq[Int]]): Unit = {
     createRecursive(ReassignPartitionsZNode.path, ReassignPartitionsZNode.encode(reassignment))
   }
 
@@ -1184,7 +1184,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 
   /**
    * Creates preferred replica election znode with partitions undergoing election
-   * @param partitions
+   * @param partitions the set of partitions
    * @throws KeeperException if there is an error while creating the znode
    */
   def createPreferredReplicaElection(partitions: Set[TopicPartition]): Unit = {
@@ -1229,7 +1229,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   }
 
 
-  def getControllerRegistration: Option[ZKControllerRegistration] = {
+  private def getControllerRegistration: Option[ZKControllerRegistration] = {
     val getDataRequest = GetDataRequest(ControllerZNode.path)
     val getDataResponse = retryRequestUntilConnected(getDataRequest)
     getDataResponse.resultCode match {
@@ -1692,7 +1692,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }
   }
 
-  def createInitialMigrationState(initialState: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
+  private def createInitialMigrationState(initialState: ZkMigrationLeadershipState): ZkMigrationLeadershipState = {
     val createRequest = CreateRequest(
       MigrationZNode.path,
       MigrationZNode.encode(initialState),
@@ -1771,7 +1771,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     * Pre-create top level paths in ZK if needed.
     */
   def createTopLevelPaths(): Unit = {
-    ZkData.PersistentZkPaths.foreach(makeSurePersistentPathExists(_))
+    ZkData.PersistentZkPaths.foreach(makeSurePersistentPathExists)
   }
 
   /**
@@ -1803,7 +1803,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   }
 
   def deleteFeatureZNode(): Unit = {
-    deletePath(FeatureZNode.path, ZkVersion.MatchAnyVersion, false)
+    deletePath(FeatureZNode.path, ZkVersion.MatchAnyVersion, recursiveDelete = false)
   }
 
   private def setConsumerOffset(group: String, topicPartition: TopicPartition, offset: Long): SetDataResponse = {
@@ -1812,7 +1812,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     retryRequestUntilConnected(setDataRequest)
   }
 
-  private def createConsumerOffset(group: String, topicPartition: TopicPartition, offset: Long) = {
+  private def createConsumerOffset(group: String, topicPartition: TopicPartition, offset: Long): Unit = {
     val path = ConsumerOffset.path(group, topicPartition.topic, topicPartition.partition)
     createRecursive(path, ConsumerOffset.encode(offset))
   }
@@ -1848,11 +1848,11 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }
   }
 
-  private[kafka] def createRecursive(path: String, data: Array[Byte] = null, throwIfPathExists: Boolean = true) = {
+  private[kafka] def createRecursive(path: String, data: Array[Byte] = null, throwIfPathExists: Boolean = true): Unit = {
 
     def parentPath(path: String): String = {
       val indexOfLastSlash = path.lastIndexOf("/")
-      if (indexOfLastSlash == -1) throw new IllegalArgumentException(s"Invalid path ${path}")
+      if (indexOfLastSlash == -1) throw new IllegalArgumentException(s"Invalid path $path")
       path.substring(0, indexOfLastSlash)
     }
 
@@ -1994,10 +1994,10 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
                 case Some(value) =>
                   val failedPayload = MigrationZNode.decode(value, version, -1)
                   throw new RuntimeException(
-                    s"Conditional update on KRaft Migration ZNode failed. Expected zkVersion = ${version}. The failed " +
-                    s"write was: ${failedPayload}. This indicates that another KRaft controller is making writes to ZooKeeper.")
+                    s"Conditional update on KRaft Migration ZNode failed. Expected zkVersion = $version. The failed " +
+                    s"write was: $failedPayload. This indicates that another KRaft controller is making writes to ZooKeeper.")
                 case None =>
-                  throw new RuntimeException(s"Check op on KRaft Migration ZNode failed. Expected zkVersion = ${version}. " +
+                  throw new RuntimeException(s"Check op on KRaft Migration ZNode failed. Expected zkVersion = $version. " +
                     s"This indicates that another KRaft controller is making writes to ZooKeeper.")
               }
             } else if (errorCode == Code.OK) {
@@ -2010,7 +2010,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
             throw new RuntimeException(s"Got migration result for incorrect path $path")
           }
         case _ => throw new RuntimeException(
-          s"Expected either CheckResult, SetDataResult, or ErrorResult for migration op, but saw ${migrationResult}")
+          s"Expected either CheckResult, SetDataResult, or ErrorResult for migration op, but saw $migrationResult")
       }
     }
 
@@ -2097,7 +2097,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     stat
   }
 
-  private def isZKSessionIdDiffFromCurrentZKSessionId(): Boolean = {
+  private def isZKSessionIdDiffFromCurrentZKSessionId: Boolean = {
     zooKeeperClient.sessionId != currentZooKeeperSessionId
   }
 
@@ -2106,7 +2106,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   }
 
   private[zk] def shouldReCreateEphemeralZNode(ephemeralOwnerId: Long): Boolean = {
-    isZKSessionTheEphemeralOwner(ephemeralOwnerId) && isZKSessionIdDiffFromCurrentZKSessionId()
+    isZKSessionTheEphemeralOwner(ephemeralOwnerId) && isZKSessionIdDiffFromCurrentZKSessionId
   }
 
   private def updateCurrentZKSessionId(newSessionId: Long): Unit = {
@@ -2125,7 +2125,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
           val setDataResult = response.zkOpResults(1).rawOpResult.asInstanceOf[SetDataResult]
           setDataResult.getStat
         case Code.NODEEXISTS =>
-          getAfterNodeExists()
+          getAfterNodeExists
         case code =>
           error(s"Error while creating ephemeral at $path with return code: $code")
           throw KeeperException.create(code)
@@ -2166,7 +2166,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
       }
     }
 
-    private def getAfterNodeExists(): Stat = {
+    private def getAfterNodeExists: Stat = {
       val getDataRequest = GetDataRequest(path)
       val getDataResponse = retryRequestUntilConnected(getDataRequest)
       val ephemeralOwnerId = getDataResponse.stat.getEphemeralOwner
@@ -2245,7 +2245,7 @@ object KafkaZkClient {
      * changed in 3.6.0.
      */
     if (zkClientConfig.getProperty(ZKConfig.JUTE_MAXBUFFER) == null)
-      zkClientConfig.setProperty(ZKConfig.JUTE_MAXBUFFER, ((4096 * 1024).toString))
+      zkClientConfig.setProperty(ZKConfig.JUTE_MAXBUFFER, (4096 * 1024).toString)
 
     if (createChrootIfNecessary) {
       val chrootIndex = connectString.indexOf("/")
diff --git a/core/src/main/scala/kafka/zk/ZkData.scala b/core/src/main/scala/kafka/zk/ZkData.scala
index 40a74114ff3da..3e830f97f39bd 100644
--- a/core/src/main/scala/kafka/zk/ZkData.scala
+++ b/core/src/main/scala/kafka/zk/ZkData.scala
@@ -185,7 +185,7 @@ object BrokerIdZNode {
       broker.rack, broker.features)
   }
 
-  def featuresAsJavaMap(brokerInfo: JsonObject): util.Map[String, util.Map[String, java.lang.Short]] = {
+  private def featuresAsJavaMap(brokerInfo: JsonObject): util.Map[String, util.Map[String, java.lang.Short]] = {
     FeatureZNode.asJavaMap(brokerInfo
       .get(FeaturesKey)
       .flatMap(_.to[Option[Map[String, Map[String, Int]]]])
@@ -461,7 +461,7 @@ object IsrChangeNotificationSequenceZNode {
       }
     }
   }.map(_.toSet).getOrElse(Set.empty)
-  def sequenceNumber(path: String) = path.substring(path.lastIndexOf(SequenceNumberPrefix) + SequenceNumberPrefix.length)
+  def sequenceNumber(path: String): String = path.substring(path.lastIndexOf(SequenceNumberPrefix) + SequenceNumberPrefix.length)
 }
 
 object LogDirEventNotificationZNode {
@@ -470,15 +470,15 @@ object LogDirEventNotificationZNode {
 
 object LogDirEventNotificationSequenceZNode {
   val SequenceNumberPrefix = "log_dir_event_"
-  val LogDirFailureEvent = 1
+  private val LogDirFailureEvent = 1
   def path(sequenceNumber: String) = s"${LogDirEventNotificationZNode.path}/$SequenceNumberPrefix$sequenceNumber"
-  def encode(brokerId: Int) = {
+  def encode(brokerId: Int): Array[Byte] = {
     Json.encodeAsBytes(Map("version" -> 1, "broker" -> brokerId, "event" -> LogDirFailureEvent).asJava)
   }
   def decode(bytes: Array[Byte]): Option[Int] = Json.parseBytes(bytes).map { js =>
     js.asJsonObject("broker").to[Int]
   }
-  def sequenceNumber(path: String) = path.substring(path.lastIndexOf(SequenceNumberPrefix) + SequenceNumberPrefix.length)
+  def sequenceNumber(path: String): String = path.substring(path.lastIndexOf(SequenceNumberPrefix) + SequenceNumberPrefix.length)
 }
 
 object AdminZNode {
@@ -561,14 +561,14 @@ object ConsumerPathZNode {
 }
 
 object ConsumerOffset {
-  def path(group: String, topic: String, partition: Integer) = s"${ConsumerPathZNode.path}/${group}/offsets/${topic}/${partition}"
+  def path(group: String, topic: String, partition: Integer) = s"${ConsumerPathZNode.path}/$group/offsets/$topic/$partition"
   def encode(offset: Long): Array[Byte] = offset.toString.getBytes(UTF_8)
   def decode(bytes: Array[Byte]): Option[Long] = Option(bytes).map(new String(_, UTF_8).toLong)
 }
 
 object ZkVersion {
-  val MatchAnyVersion = -1 // if used in a conditional set, matches any version (the value should match ZooKeeper codebase)
-  val UnknownVersion = -2  // Version returned from get if node does not exist (internal constant for Kafka codebase, unused value in ZK)
+  val MatchAnyVersion: Int = -1 // if used in a conditional set, matches any version (the value should match ZooKeeper codebase)
+  val UnknownVersion: Int = -2  // Version returned from get if node does not exist (internal constant for Kafka codebase, unused value in ZK)
 }
 
 object ZkStat {
@@ -794,7 +794,7 @@ object ClusterIdZNode {
 
   def fromJson(clusterIdJson:  Array[Byte]): String = {
     Json.parseBytes(clusterIdJson).map(_.asJsonObject("id").to[String]).getOrElse {
-      throw new KafkaException(s"Failed to parse the cluster id json $clusterIdJson")
+      throw new KafkaException(s"Failed to parse the cluster id json ${clusterIdJson.mkString("Array(", ", ", ")")}")
     }
   }
 }
@@ -844,7 +844,7 @@ object DelegationTokenChangeNotificationZNode {
 object DelegationTokenChangeNotificationSequenceZNode {
   val SequenceNumberPrefix = "token_change_"
   def createPath = s"${DelegationTokenChangeNotificationZNode.path}/$SequenceNumberPrefix"
-  def deletePath(sequenceNode: String) = s"${DelegationTokenChangeNotificationZNode.path}/${sequenceNode}"
+  def deletePath(sequenceNode: String) = s"${DelegationTokenChangeNotificationZNode.path}/$sequenceNode"
   def encode(tokenId : String): Array[Byte] = tokenId.getBytes(UTF_8)
   def decode(bytes: Array[Byte]): String = new String(bytes, UTF_8)
 }
@@ -1078,7 +1078,7 @@ object MigrationZNode {
 object ZkData {
 
   // Important: it is necessary to add any new top level Zookeeper path to the Seq
-  val SecureRootPaths = Seq(AdminZNode.path,
+  val SecureRootPaths: Seq[String] = Seq(AdminZNode.path,
     BrokersZNode.path,
     ClusterZNode.path,
     ConfigZNode.path,
@@ -1093,7 +1093,7 @@ object ZkData {
     FeatureZNode.path) ++ ZkAclStore.securePaths
 
   // These are persistent ZK paths that should exist on kafka broker startup.
-  val PersistentZkPaths = Seq(
+  val PersistentZkPaths: Seq[String] = Seq(
     ConsumerPathZNode.path, // old consumer path
     BrokerIdsZNode.path,
     TopicsZNode.path,
@@ -1105,7 +1105,7 @@ object ZkData {
     LogDirEventNotificationZNode.path
   ) ++ ConfigType.ALL.asScala.map(ConfigEntityTypeZNode.path)
 
-  val SensitiveRootPaths = Seq(
+  val SensitiveRootPaths: Seq[String] = Seq(
     ConfigEntityTypeZNode.path(ConfigType.USER),
     ConfigEntityTypeZNode.path(ConfigType.BROKER),
     DelegationTokensZNode.path
diff --git a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
index 5e06cd5e918af..ee960bd35f8bf 100644
--- a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
@@ -46,7 +46,7 @@ import scala.jdk.CollectionConverters._
 
 object ZkMigrationClient {
 
-  val MaxBatchSize = 100
+  private val MaxBatchSize = 100
 
   def apply(
     zkClient: KafkaZkClient,
@@ -295,17 +295,16 @@ class ZkMigrationClient(
     })
   }
 
-  def migrateDelegationTokens(
+  private def migrateDelegationTokens(
     recordConsumer: Consumer[util.List[ApiMessageAndVersion]]
   ): Unit = wrapZkException {
     val batch = new util.ArrayList[ApiMessageAndVersion]()
     val tokens = zkClient.getChildren(DelegationTokensZNode.path)
     for (tokenId <- tokens) {
       zkClient.getDelegationTokenInfo(tokenId) match {
-        case Some(tokenInformation) => {
+        case Some(tokenInformation) =>
           val newDelegationTokenData = new DelegationTokenData(tokenInformation)
           batch.add(new ApiMessageAndVersion(newDelegationTokenData.toRecord(), 0.toShort))
-        }
         case None =>
       }
     }
diff --git a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
index 0f37d142bcbe9..846a599875bb7 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
@@ -112,7 +112,7 @@ class ZkConfigMigrationClient(
       // Taken from ZkAdminManager
       val components = name.split("/")
       if (components.size != 3 || components(1) != "clients")
-        throw new IllegalArgumentException(s"Unexpected config path: ${name}")
+        throw new IllegalArgumentException(s"Unexpected config path: $name")
       val entity = List(
         buildEntityData(ClientQuotaEntity.USER, components(0)),
         buildEntityData(ClientQuotaEntity.CLIENT_ID, components(2))
@@ -278,7 +278,7 @@ class ZkConfigMigrationClient(
     quotas.forEach { case (key, value) =>
       val configKey = configKeys.get(key)
       if (configKey == null) {
-        throw new MigrationClientException(s"Invalid configuration key ${key}")
+        throw new MigrationClientException(s"Invalid configuration key $key")
       } else {
         configKey.`type` match {
           case ConfigDef.Type.DOUBLE =>
@@ -290,7 +290,7 @@ class ZkConfigMigrationClient(
             else
               (value + epsilon).toInt
             if ((intValue.toDouble - value).abs > epsilon)
-              throw new InvalidRequestException(s"Configuration ${key} must be a ${configKey.`type`} value")
+              throw new InvalidRequestException(s"Configuration $key must be a ${configKey.`type`} value")
             props.setProperty(key, intValue.toString)
           case _ =>
             throw new MigrationClientException(s"Unexpected config type ${configKey.`type`}")
diff --git a/core/src/main/scala/kafka/zk/migration/ZkDelegationTokenMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkDelegationTokenMigrationClient.scala
index 5b68f221acb19..77b301a751a1b 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkDelegationTokenMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkDelegationTokenMigrationClient.scala
@@ -36,7 +36,7 @@ class ZkDelegationTokenMigrationClient(
 
   val adminZkClient = new AdminZkClient(zkClient)
 
-  override def getDelegationTokens(): java.util.List[String] = {
+  override def getDelegationTokens: java.util.List[String] = {
       zkClient.getChildren(DelegationTokensZNode.path).asJava
   }
 
diff --git a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
index 04c9184791cf8..40648c57cb95a 100755
--- a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
+++ b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
@@ -179,7 +179,7 @@ class ZooKeeperClient(connectString: String,
     // Safe to cast as we always create a response of the right type
     def callback(response: AsyncResponse): Unit = processResponse(response.asInstanceOf[Req#Response])
 
-    def responseMetadata(sendTimeMs: Long) = new ResponseMetadata(sendTimeMs, receivedTimeMs = time.hiResClockMs())
+    def responseMetadata(sendTimeMs: Long) = ResponseMetadata(sendTimeMs, receivedTimeMs = time.hiResClockMs())
 
     val sendTimeMs = time.hiResClockMs()
 
@@ -348,7 +348,7 @@ class ZooKeeperClient(connectString: String,
       zNodeChildChangeHandlers.clear()
       stateChangeHandlers.clear()
       zooKeeper.close()
-      metricNames.foreach(metricsGroup.removeMetric(_))
+      metricNames.foreach(metricsGroup.removeMetric)
     }
     info("Closed.")
   }
@@ -365,7 +365,7 @@ class ZooKeeperClient(connectString: String,
   private def reinitialize(): Unit = {
     // Initialization callbacks are invoked outside of the lock to avoid deadlock potential since their completion
     // may require additional Zookeeper requests, which will block to acquire the initialization lock
-    stateChangeHandlers.values.foreach(callBeforeInitializingSession _)
+    stateChangeHandlers.values.foreach(callBeforeInitializingSession)
 
     inWriteLock(initializationLock) {
       if (!connectionState.isAlive) {
@@ -386,7 +386,7 @@ class ZooKeeperClient(connectString: String,
       }
     }
 
-    stateChangeHandlers.values.foreach(callAfterInitializingSession _)
+    stateChangeHandlers.values.foreach(callAfterInitializingSession)
   }
 
   /**

From 6eea40882d4767a2a45d81e7a009a9eafaa15d00 Mon Sep 17 00:00:00 2001
From: Ahmed Sobeh <ahmed.sobeh@aiven.io>
Date: Tue, 20 Feb 2024 04:19:53 +0100
Subject: [PATCH 054/258] KAFKA-15349: ducker-ak should fail fast when gradlew
 systemTestLibs fails (#15391)

In this modification, if ./gradlew systemTestLibs fails, the script will output an error message and terminate execution using the die function. This ensures that the script fails fast and prompts the user to address the error before continuing.

Reviewers: Luke Chen <showuon@gmail.com>
---
 tests/docker/ducker-ak | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/docker/ducker-ak b/tests/docker/ducker-ak
index 8047b8fa63e86..4add3172022fd 100755
--- a/tests/docker/ducker-ak
+++ b/tests/docker/ducker-ak
@@ -477,7 +477,7 @@ ducker_test() {
     done
     
     must_pushd "${kafka_dir}"
-    (test -f ./gradlew || gradle) && ./gradlew systemTestLibs
+    ( (test -f ./gradlew || gradle) && ./gradlew systemTestLibs ) || die "ducker_test: Failed to build system test libraries, please check the error log."
     must_popd
     if [[ "${debug}" -eq 1 ]]; then
         local ducktape_cmd="python3 -m debugpy --listen 0.0.0.0:${debugpy_port} --wait-for-client /usr/local/bin/ducktape"

From a26a1d847f1884a519561e7a4fb4cd13e051c824 Mon Sep 17 00:00:00 2001
From: runom <jeatommy@gmail.com>
Date: Tue, 20 Feb 2024 12:23:06 +0900
Subject: [PATCH 055/258] MINOR: fix
 MetricsTest.testBrokerTopicMetricsBytesInOut (#14744)

The assertion to check BytesOut doesn't include replication was performed before replication occurred.
This PR fixed the position of the assertion.

Reviewers: Luke Chen <showuon@gmail.com>
---
 core/src/test/scala/unit/kafka/metrics/MetricsTest.scala | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala b/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
index ecee4b03c2500..ed307eef8fc0b 100644
--- a/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
+++ b/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
@@ -197,15 +197,14 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
     val initialBytesIn = TestUtils.meterCount(bytesIn)
     val initialBytesOut = TestUtils.meterCount(bytesOut)
 
-    // BytesOut doesn't include replication, so it shouldn't have changed
-    assertEquals(initialBytesOut, TestUtils.meterCount(bytesOut))
-
     // Produce a few messages to make the metrics tick
     TestUtils.generateAndProduceMessages(brokers, topic, nMessages)
 
     assertTrue(TestUtils.meterCount(replicationBytesIn) > initialReplicationBytesIn)
     assertTrue(TestUtils.meterCount(replicationBytesOut) > initialReplicationBytesOut)
     assertTrue(TestUtils.meterCount(bytesIn) > initialBytesIn)
+    // BytesOut doesn't include replication, so it shouldn't have changed
+    assertEquals(initialBytesOut, TestUtils.meterCount(bytesOut))
 
     // Consume messages to make bytesOut tick
     TestUtils.consumeTopicRecords(brokers, topic, nMessages)

From 5854139cd8a75fec387183189c6cd1f32e891e57 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Tue, 20 Feb 2024 10:48:36 +0100
Subject: [PATCH 056/258] KAFKA-16243: Make sure that we do not exceed max poll
 interval inside poll (#15372)

The consumer keeps a poll timer, which is used to ensure liveness of the application thread. The poll timer automatically updates while the Consumer.poll(Duration) method is blocked, while the newer consumer only updates the poll timer when a new call to Consumer.poll(Duration) is issued. This means that the kafka-console-consumer.sh tools, which uses a very long timeout by default, works differently with the new consumer, with the consumer proactively rejoining the group during long poll timeouts.

This change solves the problem by (a) repeatedly sending PollApplicationEvents to the background thread, not just on the first call of poll and (b) making sure that the application thread doesn't block for so long that it runs out of max.poll.interval.

An integration test is added to make sure that we do not rejoin the group when a long poll timeout is used with a low max.poll.interval.

Reviewers: Lianet Magrans <lianetmr@gmail.com>, Andrew Schofield <aschofield@confluent.io>, Bruno Cadonna <cadonna@apache.org>
---
 .../internals/AsyncKafkaConsumer.java         |  6 +++--
 .../internals/HeartbeatRequestManager.java    | 13 ++++++++--
 .../kafka/api/PlaintextConsumerTest.scala     | 26 +++++++++++++++++++
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index 28d26a83be6b6..6781e0b73cb80 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -698,9 +698,11 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
                 throw new IllegalStateException("Consumer is not subscribed to any topics or assigned any partitions");
             }
 
-            applicationEventHandler.add(new PollApplicationEvent(timer.currentTimeMs()));
-
             do {
+
+                // Make sure to let the background thread know that we are still polling.
+                applicationEventHandler.add(new PollApplicationEvent(timer.currentTimeMs()));
+
                 // We must not allow wake-ups between polling for fetches and returning the records.
                 // If the polled fetches are not empty the consumed position has already been updated in the polling
                 // of the fetches. A wakeup between returned fetches and returning records would lead to never
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index a6b0b62c4ebe2..7408f584cdb4f 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -232,13 +232,22 @@ public MembershipManager membershipManager() {
      * are sent, so blocking for longer than the heartbeat interval might mean the application thread is not
      * responsive to changes.
      *
+     * Similarly, we may have to unblock the application thread to send a `PollApplicationEvent` to make sure
+     * our poll timer will not expire while we are polling.
+     *
      * <p>In the event that heartbeats are currently being skipped, this still returns the next heartbeat
      * delay rather than {@code Long.MAX_VALUE} so that the application thread remains responsive.
      */
     @Override
     public long maximumTimeToWait(long currentTimeMs) {
-        boolean heartbeatNow = membershipManager.shouldHeartbeatNow() && !heartbeatRequestState.requestInFlight();
-        return heartbeatNow ? 0L : heartbeatRequestState.nextHeartbeatMs(currentTimeMs);
+        pollTimer.update(currentTimeMs);
+        if (
+            pollTimer.isExpired() ||
+                (membershipManager.shouldHeartbeatNow() && !heartbeatRequestState.requestInFlight())
+        ) {
+            return 0L;
+        }
+        return Math.min(pollTimer.remainingMs() / 2, heartbeatRequestState.nextHeartbeatMs(currentTimeMs));
     }
 
     /**
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 553188900cdcf..c96288703fb99 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -263,6 +263,32 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     ensureNoRebalance(consumer, listener)
   }
 
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollIntervalMsShorterThanPollTimeout(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
+
+    val consumer = createConsumer()
+    val listener = new TestConsumerReassignmentListener
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
+
+    val callsToAssignedAfterFirstRebalance = listener.callsToAssigned
+
+    consumer.poll(Duration.ofMillis(2000))
+
+    // If the poll poll above times out, it would trigger a rebalance.
+    // Leave some time for the rebalance to happen and check for the rebalance event.
+    consumer.poll(Duration.ofMillis(500))
+    consumer.poll(Duration.ofMillis(500))
+
+    assertEquals(callsToAssignedAfterFirstRebalance, listener.callsToAssigned)
+  }
+
+
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testAutoCommitOnClose(quorum: String, groupProtocol: String): Unit = {

From ead2431c37ace9255df88ffe819bb905311af088 Mon Sep 17 00:00:00 2001
From: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
Date: Tue, 20 Feb 2024 10:25:17 +0000
Subject: [PATCH 057/258] MINOR: Remove unwanted debug line in
 LogDirFailureTest (#15371)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Justine Olshan <jolshan@confluent.io>, Igor Soarez <soarez@apple.com>
---
 core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
index 57be102bea042..54a10d78c76e8 100644
--- a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
@@ -169,10 +169,6 @@ class LogDirFailureTest extends IntegrationTestHarness {
   }
 
   def testProduceAfterLogDirFailureOnLeader(failureType: LogDirFailureType, quorum: String): Unit = {
-    if (isKRaftTest()) {
-      val value = configs.map(c => c.brokerId -> c.logDirs.contains(c.metadataLogDir))
-      logger.warn(s">>>>>> ${value.mkString(",")}")
-    }
     val consumer = createConsumer()
     subscribeAndWaitForAssignment(topic, consumer)
 

From 4c70581eb63fe74494fbabf5a90e87c38e17996d Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Tue, 20 Feb 2024 12:24:32 -0800
Subject: [PATCH 058/258] KAFKA-15770: IQv2 must return immutable position
 (#15219)

ConsistencyVectorIntegrationTest failed frequently because the return
Position from IQv2 is not immutable while the test assume immutability.
To return a Position with a QueryResult that does not change, we need to
deep copy the Position object.

Reviewers: John Roesler <john@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 ...tDualSchemaRocksDBSegmentedBytesStore.java |  47 +++----
 .../AbstractRocksDBSegmentedBytesStore.java   |  30 +++--
 .../state/internals/AbstractSegments.java     |   6 +
 .../state/internals/CachingKeyValueStore.java |  69 ++++++-----
 .../internals/InMemoryKeyValueStore.java      |  30 +++--
 .../state/internals/InMemorySessionStore.java |  50 ++++----
 .../state/internals/InMemoryWindowStore.java  |  66 +++++-----
 .../state/internals/KeyValueSegment.java      |   3 +
 .../state/internals/KeyValueSegments.java     |   2 +-
 .../internals/LogicalKeyValueSegments.java    |   6 +
 .../state/internals/MemoryLRUCache.java       |  34 ++---
 .../streams/state/internals/RocksDBStore.java |  47 +++----
 .../internals/RocksDBTimestampedStore.java    |  54 ++++----
 .../internals/RocksDBVersionedStore.java      | 116 ++++++++++--------
 .../state/internals/StoreQueryUtils.java      |  44 +++----
 .../state/internals/TimestampedSegment.java   |   3 +
 .../state/internals/TimestampedSegments.java  |   2 +-
 .../ConsistencyVectorIntegrationTest.java     |  40 +++---
 .../integration/IQv2IntegrationTest.java      |  26 ++--
 .../StoreUpgradeIntegrationTest.java          |   3 +-
 .../state/internals/KeyValueSegmentTest.java  |  21 ++--
 .../LogicalKeyValueSegmentsTest.java          |   2 +
 .../state/internals/SegmentIteratorTest.java  |   5 +-
 .../internals/TimestampedSegmentTest.java     |  21 ++--
 24 files changed, 403 insertions(+), 324 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStore.java
index 441c17201b4dd..284385a47de9f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStore.java
@@ -193,17 +193,19 @@ public void put(final Bytes rawBaseKey,
             expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
             LOG.warn("Skipping record for expired segment.");
         } else {
-            StoreQueryUtils.updatePosition(position, stateStoreContext);
-
-            // Put to index first so that if put to base failed, when we iterate index, we will
-            // find no base value. If put to base first but putting to index fails, when we iterate
-            // index, we can't find the key but if we iterate over base store, we can find the key
-            // which lead to inconsistency.
-            if (hasIndex()) {
-                final KeyValue<Bytes, byte[]> indexKeyValue = getIndexKeyValue(rawBaseKey, value);
-                segment.put(indexKeyValue.key, indexKeyValue.value);
+            synchronized (position) {
+                StoreQueryUtils.updatePosition(position, stateStoreContext);
+
+                // Put to index first so that if put to base failed, when we iterate index, we will
+                // find no base value. If put to base first but putting to index fails, when we iterate
+                // index, we can't find the key but if we iterate over base store, we can find the key
+                // which lead to inconsistency.
+                if (hasIndex()) {
+                    final KeyValue<Bytes, byte[]> indexKeyValue = getIndexKeyValue(rawBaseKey, value);
+                    segment.put(indexKeyValue.key, indexKeyValue.value);
+                }
+                segment.put(rawBaseKey, value);
             }
-            segment.put(rawBaseKey, value);
         }
     }
 
@@ -254,11 +256,12 @@ public void init(final ProcessorContext context,
             metrics
         );
 
-        segments.openExisting(context, observedStreamTime);
-
         final File positionCheckpointFile = new File(context.stateDir(), name() + ".position");
         this.positionCheckpoint = new OffsetCheckpoint(positionCheckpointFile);
         this.position = StoreQueryUtils.readPositionFromCheckpoint(positionCheckpoint);
+        segments.setPosition(this.position);
+
+        segments.openExisting(context, observedStreamTime);
 
         // register and possibly restore the state from the logs
         stateStoreContext.register(
@@ -310,16 +313,18 @@ List<S> getSegments() {
 
     // Visible for testing
     void restoreAllInternal(final Collection<ConsumerRecord<byte[], byte[]>> records) {
-        try {
-            final Map<S, WriteBatch> writeBatchMap = getWriteBatches(records);
-            for (final Map.Entry<S, WriteBatch> entry : writeBatchMap.entrySet()) {
-                final S segment = entry.getKey();
-                final WriteBatch batch = entry.getValue();
-                segment.write(batch);
-                batch.close();
+        synchronized (position) {
+            try {
+                final Map<S, WriteBatch> writeBatchMap = getWriteBatches(records);
+                for (final Map.Entry<S, WriteBatch> entry : writeBatchMap.entrySet()) {
+                    final S segment = entry.getKey();
+                    final WriteBatch batch = entry.getValue();
+                    segment.write(batch);
+                    batch.close();
+                }
+            } catch (final RocksDBException e) {
+                throw new ProcessorStateException("Error restoring batch to store " + this.name, e);
             }
-        } catch (final RocksDBException e) {
-            throw new ProcessorStateException("Error restoring batch to store " + this.name, e);
         }
     }
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
index 79f4f46098980..348d6bb18632a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
@@ -266,8 +266,10 @@ public void put(final Bytes key,
             expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
             LOG.warn("Skipping record for expired segment.");
         } else {
-            StoreQueryUtils.updatePosition(position, stateStoreContext);
-            segment.put(key, value);
+            synchronized (position) {
+                StoreQueryUtils.updatePosition(position, stateStoreContext);
+                segment.put(key, value);
+            }
         }
     }
 
@@ -308,11 +310,11 @@ public void init(final ProcessorContext context,
                 metrics
         );
 
-        segments.openExisting(this.context, observedStreamTime);
-
         final File positionCheckpointFile = new File(context.stateDir(), name() + ".position");
         this.positionCheckpoint = new OffsetCheckpoint(positionCheckpointFile);
         this.position = StoreQueryUtils.readPositionFromCheckpoint(positionCheckpoint);
+        segments.setPosition(position);
+        segments.openExisting(this.context, observedStreamTime);
 
         // register and possibly restore the state from the logs
         stateStoreContext.register(
@@ -363,16 +365,18 @@ List<S> getSegments() {
 
     // Visible for testing
     void restoreAllInternal(final Collection<ConsumerRecord<byte[], byte[]>> records) {
-        try {
-            final Map<S, WriteBatch> writeBatchMap = getWriteBatches(records);
-            for (final Map.Entry<S, WriteBatch> entry : writeBatchMap.entrySet()) {
-                final S segment = entry.getKey();
-                final WriteBatch batch = entry.getValue();
-                segment.write(batch);
-                batch.close();
+        synchronized (position) {
+            try {
+                final Map<S, WriteBatch> writeBatchMap = getWriteBatches(records);
+                for (final Map.Entry<S, WriteBatch> entry : writeBatchMap.entrySet()) {
+                    final S segment = entry.getKey();
+                    final WriteBatch batch = entry.getValue();
+                    segment.write(batch);
+                    batch.close();
+                }
+            } catch (final RocksDBException e) {
+                throw new ProcessorStateException("Error restoring batch to store " + this.name, e);
             }
-        } catch (final RocksDBException e) {
-            throw new ProcessorStateException("Error restoring batch to store " + this.name, e);
         }
     }
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractSegments.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractSegments.java
index 5c31894f6e316..495c3423b8f88 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractSegments.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractSegments.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.streams.errors.ProcessorStateException;
 import org.apache.kafka.streams.processor.ProcessorContext;
+import org.apache.kafka.streams.query.Position;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -43,6 +44,7 @@ abstract class AbstractSegments<S extends Segment> implements Segments<S> {
     private final long retentionPeriod;
     private final long segmentInterval;
     private final SimpleDateFormat formatter;
+    Position position;
 
     AbstractSegments(final String name, final long retentionPeriod, final long segmentInterval) {
         this.name = name;
@@ -53,6 +55,10 @@ abstract class AbstractSegments<S extends Segment> implements Segments<S> {
         this.formatter.setTimeZone(new SimpleTimeZone(0, "UTC"));
     }
 
+    public void setPosition(final Position position) {
+        this.position = position;
+    }
+
     @Override
     public long segmentId(final long timestamp) {
         return timestamp / segmentInterval;
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java
index 36f08412828d6..d89672daff243 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingKeyValueStore.java
@@ -120,8 +120,13 @@ public void init(final StateStoreContext context,
     public Position getPosition() {
         // We return the merged position since the query uses the merged position as well
         final Position mergedPosition = Position.emptyPosition();
-        mergedPosition.merge(position);
-        mergedPosition.merge(wrapped().getPosition());
+        final Position wrappedPosition = wrapped().getPosition();
+        synchronized (position) {
+            synchronized (wrappedPosition) {
+                mergedPosition.merge(position);
+                mergedPosition.merge(wrappedPosition);
+            }
+        }
         return mergedPosition;
     }
 
@@ -183,25 +188,27 @@ private <R> QueryResult<R> runKeyQuery(final Query<R> query,
 
         final Bytes key = keyQuery.getKey();
 
-        if (context.cache() != null) {
-            final LRUCacheEntry lruCacheEntry = context.cache().get(cacheName, key);
-            if (lruCacheEntry != null) {
-                final byte[] rawValue;
-                if (timestampedSchema && !WrappedStateStore.isTimestamped(wrapped()) && !StoreQueryUtils.isAdapter(wrapped())) {
-                    rawValue = ValueAndTimestampDeserializer.rawValue(lruCacheEntry.value());
-                } else {
-                    rawValue = lruCacheEntry.value();
+        synchronized (mergedPosition) {
+            if (context.cache() != null) {
+                final LRUCacheEntry lruCacheEntry = context.cache().get(cacheName, key);
+                if (lruCacheEntry != null) {
+                    final byte[] rawValue;
+                    if (timestampedSchema && !WrappedStateStore.isTimestamped(wrapped()) && !StoreQueryUtils.isAdapter(wrapped())) {
+                        rawValue = ValueAndTimestampDeserializer.rawValue(lruCacheEntry.value());
+                    } else {
+                        rawValue = lruCacheEntry.value();
+                    }
+                    result = (QueryResult<R>) QueryResult.forResult(rawValue);
                 }
-                result = (QueryResult<R>) QueryResult.forResult(rawValue);
             }
-        }
 
-        // We don't need to check the position at the state store since we already performed the check on
-        // the merged position above
-        if (result == null) {
-            result = wrapped().query(query, PositionBound.unbounded(), config);
+            // We don't need to check the position at the state store since we already performed the check on
+            // the merged position above
+            if (result == null) {
+                result = wrapped().query(query, PositionBound.unbounded(), config);
+            }
+            result.setPosition(mergedPosition.copy());
         }
-        result.setPosition(mergedPosition);
         return result;
     }
 
@@ -276,19 +283,21 @@ public void put(final Bytes key,
 
     private void putInternal(final Bytes key,
                              final byte[] value) {
-        context.cache().put(
-            cacheName,
-            key,
-            new LRUCacheEntry(
-                value,
-                context.headers(),
-                true,
-                context.offset(),
-                context.timestamp(),
-                context.partition(),
-                context.topic()));
-
-        StoreQueryUtils.updatePosition(position, context);
+        synchronized (position) {
+            context.cache().put(
+                cacheName,
+                key,
+                new LRUCacheEntry(
+                    value,
+                    context.headers(),
+                    true,
+                    context.offset(),
+                    context.timestamp(),
+                    context.partition(),
+                    context.topic()));
+
+            StoreQueryUtils.updatePosition(position, context);
+        }
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
index 7599bff82b366..652db52f0c73b 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
@@ -79,13 +79,15 @@ public void init(final ProcessorContext context,
             context.register(
                 root,
                 (RecordBatchingStateRestoreCallback) records -> {
-                    for (final ConsumerRecord<byte[], byte[]> record : records) {
-                        put(Bytes.wrap(record.key()), record.value());
-                        ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
-                            record,
-                            consistencyEnabled,
-                            position
-                        );
+                    synchronized (position) {
+                        for (final ConsumerRecord<byte[], byte[]> record : records) {
+                            put(Bytes.wrap(record.key()), record.value());
+                            ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                                record,
+                                consistencyEnabled,
+                                position
+                            );
+                        }
                     }
                 }
             );
@@ -152,13 +154,15 @@ public synchronized byte[] putIfAbsent(final Bytes key, final byte[] value) {
 
     // the unlocked implementation of put method, to avoid multiple lock/unlock cost in `putAll` method
     private void putInternal(final Bytes key, final byte[] value) {
-        if (value == null) {
-            map.remove(key);
-        } else {
-            map.put(key, value);
-        }
+        synchronized (position) {
+            if (value == null) {
+                map.remove(key);
+            } else {
+                map.put(key, value);
+            }
 
-        StoreQueryUtils.updatePosition(position, context);
+            StoreQueryUtils.updatePosition(position, context);
+        }
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java
index 579abc3678275..67a6ea46b6c3a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java
@@ -124,13 +124,15 @@ public void init(final ProcessorContext context, final StateStore root) {
             context.register(
                 root,
                 (RecordBatchingStateRestoreCallback) records -> {
-                    for (final ConsumerRecord<byte[], byte[]> record : records) {
-                        put(SessionKeySchema.from(Bytes.wrap(record.key())), record.value());
-                        ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
-                            record,
-                            consistencyEnabled,
-                            position
-                        );
+                    synchronized (position) {
+                        for (final ConsumerRecord<byte[], byte[]> record : records) {
+                            put(SessionKeySchema.from(Bytes.wrap(record.key())), record.value());
+                            ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                                record,
+                                consistencyEnabled,
+                                position
+                            );
+                        }
                     }
                 }
             );
@@ -157,25 +159,27 @@ public void put(final Windowed<Bytes> sessionKey, final byte[] aggregate) {
         final long windowEndTimestamp = sessionKey.window().end();
         observedStreamTime = Math.max(observedStreamTime, windowEndTimestamp);
 
-        if (windowEndTimestamp <= observedStreamTime - retentionPeriod) {
-            // The provided context is not required to implement InternalProcessorContext,
-            // If it doesn't, we can't record this metric (in fact, we wouldn't have even initialized it).
-            if (expiredRecordSensor != null && context != null) {
-                expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
-            }
-            LOG.warn("Skipping record for expired segment.");
-        } else {
-            if (aggregate != null) {
-                endTimeMap.computeIfAbsent(windowEndTimestamp, t -> new ConcurrentSkipListMap<>());
-                final ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>> keyMap = endTimeMap.get(windowEndTimestamp);
-                keyMap.computeIfAbsent(sessionKey.key(), t -> new ConcurrentSkipListMap<>());
-                keyMap.get(sessionKey.key()).put(sessionKey.window().start(), aggregate);
+        synchronized (position) {
+            if (windowEndTimestamp <= observedStreamTime - retentionPeriod) {
+                // The provided context is not required to implement InternalProcessorContext,
+                // If it doesn't, we can't record this metric (in fact, we wouldn't have even initialized it).
+                if (expiredRecordSensor != null && context != null) {
+                    expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
+                }
+                LOG.warn("Skipping record for expired segment.");
             } else {
-                remove(sessionKey);
+                if (aggregate != null) {
+                    endTimeMap.computeIfAbsent(windowEndTimestamp, t -> new ConcurrentSkipListMap<>());
+                    final ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>> keyMap = endTimeMap.get(windowEndTimestamp);
+                    keyMap.computeIfAbsent(sessionKey.key(), t -> new ConcurrentSkipListMap<>());
+                    keyMap.get(sessionKey.key()).put(sessionKey.window().start(), aggregate);
+                } else {
+                    remove(sessionKey);
+                }
             }
-        }
 
-        StoreQueryUtils.updatePosition(position, stateStoreContext);
+            StoreQueryUtils.updatePosition(position, stateStoreContext);
+        }
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryWindowStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryWindowStore.java
index 2ddeadc3585c1..2f85a1b4cc6ae 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryWindowStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryWindowStore.java
@@ -123,17 +123,19 @@ public void init(final ProcessorContext context, final StateStore root) {
             context.register(
                 root,
                 (RecordBatchingStateRestoreCallback) records -> {
-                    for (final ConsumerRecord<byte[], byte[]> record : records) {
-                        put(
-                            Bytes.wrap(extractStoreKeyBytes(record.key())),
-                            record.value(),
-                            extractStoreTimestamp(record.key())
-                        );
-                        ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
-                            record,
-                            consistencyEnabled,
-                            position
-                        );
+                    synchronized (position) {
+                        for (final ConsumerRecord<byte[], byte[]> record : records) {
+                            put(
+                                Bytes.wrap(extractStoreKeyBytes(record.key())),
+                                record.value(),
+                                extractStoreTimestamp(record.key())
+                            );
+                            ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                                record,
+                                consistencyEnabled,
+                                position
+                            );
+                        }
                     }
                 }
             );
@@ -158,28 +160,30 @@ public void put(final Bytes key, final byte[] value, final long windowStartTimes
         removeExpiredSegments();
         observedStreamTime = Math.max(observedStreamTime, windowStartTimestamp);
 
-        if (windowStartTimestamp <= observedStreamTime - retentionPeriod) {
-            expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
-            LOG.warn("Skipping record for expired segment.");
-        } else {
-            if (value != null) {
-                maybeUpdateSeqnumForDups();
-                final Bytes keyBytes = retainDuplicates ? wrapForDups(key, seqnum) : key;
-                segmentMap.computeIfAbsent(windowStartTimestamp, t -> new ConcurrentSkipListMap<>());
-                segmentMap.get(windowStartTimestamp).put(keyBytes, value);
-            } else if (!retainDuplicates) {
-                // Skip if value is null and duplicates are allowed since this delete is a no-op
-                segmentMap.computeIfPresent(windowStartTimestamp, (t, kvMap) -> {
-                    kvMap.remove(key);
-                    if (kvMap.isEmpty()) {
-                        segmentMap.remove(windowStartTimestamp);
-                    }
-                    return kvMap;
-                });
+        synchronized (position) {
+            if (windowStartTimestamp <= observedStreamTime - retentionPeriod) {
+                expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
+                LOG.warn("Skipping record for expired segment.");
+            } else {
+                if (value != null) {
+                    maybeUpdateSeqnumForDups();
+                    final Bytes keyBytes = retainDuplicates ? wrapForDups(key, seqnum) : key;
+                    segmentMap.computeIfAbsent(windowStartTimestamp, t -> new ConcurrentSkipListMap<>());
+                    segmentMap.get(windowStartTimestamp).put(keyBytes, value);
+                } else if (!retainDuplicates) {
+                    // Skip if value is null and duplicates are allowed since this delete is a no-op
+                    segmentMap.computeIfPresent(windowStartTimestamp, (t, kvMap) -> {
+                        kvMap.remove(key);
+                        if (kvMap.isEmpty()) {
+                            segmentMap.remove(windowStartTimestamp);
+                        }
+                        return kvMap;
+                    });
+                }
             }
-        }
 
-        StoreQueryUtils.updatePosition(position, stateStoreContext);
+            StoreQueryUtils.updatePosition(position, stateStoreContext);
+        }
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegment.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegment.java
index 66c55fc9d92cb..1cf631c37243f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegment.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegment.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 
 import java.io.File;
@@ -31,9 +32,11 @@ class KeyValueSegment extends RocksDBStore implements Comparable<KeyValueSegment
     KeyValueSegment(final String segmentName,
                     final String windowName,
                     final long id,
+                    final Position position,
                     final RocksDBMetricsRecorder metricsRecorder) {
         super(segmentName, windowName, metricsRecorder);
         this.id = id;
+        this.position = position;
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegments.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegments.java
index 2dccaea4df235..c08c106f155d3 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegments.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/KeyValueSegments.java
@@ -42,7 +42,7 @@ public KeyValueSegment getOrCreateSegment(final long segmentId,
             return segments.get(segmentId);
         } else {
             final KeyValueSegment newSegment =
-                new KeyValueSegment(segmentName(segmentId), name, segmentId, metricsRecorder);
+                new KeyValueSegment(segmentName(segmentId), name, segmentId, position, metricsRecorder);
 
             if (segments.put(segmentId, newSegment) != null) {
                 throw new IllegalStateException("KeyValueSegment already exists. Possible concurrent access.");
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java
index e099b0b81a6e6..2b445ed41cdd0 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java
@@ -20,6 +20,7 @@
 import java.util.Map;
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.internals.ProcessorContextUtils;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 
 /**
@@ -53,6 +54,11 @@ public class LogicalKeyValueSegments extends AbstractSegments<LogicalKeyValueSeg
         this.physicalStore = new RocksDBStore(name, parentDir, metricsRecorder, false);
     }
 
+    @Override
+    public void setPosition(final Position position) {
+        this.physicalStore.position = position;
+    }
+
     @Override
     public LogicalKeyValueSegment getOrCreateSegment(final long segmentId,
                                                      final ProcessorContext context) {
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryLRUCache.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryLRUCache.java
index 4744bb7ae63cc..1d43576228edf 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryLRUCache.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MemoryLRUCache.java
@@ -108,13 +108,15 @@ public void init(final StateStoreContext context, final StateStore root) {
             root,
             (RecordBatchingStateRestoreCallback) records -> {
                 restoring = true;
-                for (final ConsumerRecord<byte[], byte[]> record : records) {
-                    put(Bytes.wrap(record.key()), record.value());
-                    ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
-                        record,
-                        consistencyEnabled,
-                        position
-                    );
+                synchronized (position) {
+                    for (final ConsumerRecord<byte[], byte[]> record : records) {
+                        put(Bytes.wrap(record.key()), record.value());
+                        ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                            record,
+                            consistencyEnabled,
+                            position
+                        );
+                    }
                 }
                 restoring = false;
             }
@@ -147,12 +149,14 @@ public synchronized byte[] get(final Bytes key) {
     @Override
     public synchronized void put(final Bytes key, final byte[] value) {
         Objects.requireNonNull(key);
-        if (value == null) {
-            delete(key);
-        } else {
-            this.map.put(key, value);
+        synchronized (position) {
+            if (value == null) {
+                delete(key);
+            } else {
+                this.map.put(key, value);
+            }
+            StoreQueryUtils.updatePosition(position, context);
         }
-        StoreQueryUtils.updatePosition(position, context);
     }
 
     @Override
@@ -175,8 +179,10 @@ public void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
     @Override
     public synchronized byte[] delete(final Bytes key) {
         Objects.requireNonNull(key);
-        StoreQueryUtils.updatePosition(position, context);
-        return this.map.remove(key);
+        synchronized (position) {
+            StoreQueryUtils.updatePosition(position, context);
+            return this.map.remove(key);
+        }
     }
 
     /**
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
index b09f8e1ddad12..0dfae3499c10f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
@@ -391,9 +391,11 @@ public synchronized void put(final Bytes key,
                                  final byte[] value) {
         Objects.requireNonNull(key, "key cannot be null");
         validateStoreOpen();
-        cfAccessor.put(dbAccessor, key.get(), value);
 
-        StoreQueryUtils.updatePosition(position, context);
+        synchronized (position) {
+            cfAccessor.put(dbAccessor, key.get(), value);
+            StoreQueryUtils.updatePosition(position, context);
+        }
     }
 
     @Override
@@ -409,12 +411,14 @@ public synchronized byte[] putIfAbsent(final Bytes key,
 
     @Override
     public void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
-        try (final WriteBatch batch = new WriteBatch()) {
-            cfAccessor.prepareBatch(entries, batch);
-            write(batch);
-            StoreQueryUtils.updatePosition(position, context);
-        } catch (final RocksDBException e) {
-            throw new ProcessorStateException("Error while batch writing to store " + name, e);
+        synchronized (position) {
+            try (final WriteBatch batch = new WriteBatch()) {
+                cfAccessor.prepareBatch(entries, batch);
+                write(batch);
+                StoreQueryUtils.updatePosition(position, context);
+            } catch (final RocksDBException e) {
+                throw new ProcessorStateException("Error while batch writing to store " + name, e);
+            }
         }
     }
 
@@ -980,21 +984,22 @@ public void close() {
     }
 
     void restoreBatch(final Collection<ConsumerRecord<byte[], byte[]>> records) {
-        try (final WriteBatch batch = new WriteBatch()) {
-            for (final ConsumerRecord<byte[], byte[]> record : records) {
-                ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
-                    record,
-                    consistencyEnabled,
-                    position
-                );
-                // If version headers are not present or version is V0
-                cfAccessor.addToBatch(record.key(), record.value(), batch);
+        synchronized (position) {
+            try (final WriteBatch batch = new WriteBatch()) {
+                for (final ConsumerRecord<byte[], byte[]> record : records) {
+                    ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                        record,
+                        consistencyEnabled,
+                        position
+                    );
+                    // If version headers are not present or version is V0
+                    cfAccessor.addToBatch(record.key(), record.value(), batch);
+                }
+                write(batch);
+            } catch (final RocksDBException e) {
+                throw new ProcessorStateException("Error restoring batch to store " + name, e);
             }
-            write(batch);
-        } catch (final RocksDBException e) {
-            throw new ProcessorStateException("Error restoring batch to store " + name, e);
         }
-
     }
 
     // for testing
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStore.java
index d4ba8ccb072fd..7b75649ba71ef 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStore.java
@@ -101,32 +101,34 @@ private DualColumnFamilyAccessor(final ColumnFamilyHandle oldColumnFamily,
         public void put(final DBAccessor accessor,
                         final byte[] key,
                         final byte[] valueWithTimestamp) {
-            if (valueWithTimestamp == null) {
-                try {
-                    accessor.delete(oldColumnFamily, key);
-                } catch (final RocksDBException e) {
-                    // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
-                    throw new ProcessorStateException("Error while removing key from store " + name, e);
-                }
-                try {
-                    accessor.delete(newColumnFamily, key);
-                } catch (final RocksDBException e) {
-                    // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
-                    throw new ProcessorStateException("Error while removing key from store " + name, e);
-                }
-            } else {
-                try {
-                    accessor.delete(oldColumnFamily, key);
-                } catch (final RocksDBException e) {
-                    // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
-                    throw new ProcessorStateException("Error while removing key from store " + name, e);
-                }
-                try {
-                    accessor.put(newColumnFamily, key, valueWithTimestamp);
-                    StoreQueryUtils.updatePosition(position, context);
-                } catch (final RocksDBException e) {
-                    // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
-                    throw new ProcessorStateException("Error while putting key/value into store " + name, e);
+            synchronized (position) {
+                if (valueWithTimestamp == null) {
+                    try {
+                        accessor.delete(oldColumnFamily, key);
+                    } catch (final RocksDBException e) {
+                        // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
+                        throw new ProcessorStateException("Error while removing key from store " + name, e);
+                    }
+                    try {
+                        accessor.delete(newColumnFamily, key);
+                    } catch (final RocksDBException e) {
+                        // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
+                        throw new ProcessorStateException("Error while removing key from store " + name, e);
+                    }
+                } else {
+                    try {
+                        accessor.delete(oldColumnFamily, key);
+                    } catch (final RocksDBException e) {
+                        // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
+                        throw new ProcessorStateException("Error while removing key from store " + name, e);
+                    }
+                    try {
+                        accessor.put(newColumnFamily, key, valueWithTimestamp);
+                        StoreQueryUtils.updatePosition(position, context);
+                    } catch (final RocksDBException e) {
+                        // String format is happening in wrapping stores. So formatted message is thrown from wrapping stores.
+                        throw new ProcessorStateException("Error while putting key/value into store " + name, e);
+                    }
                 }
             }
         }
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStore.java
index 3379d3cebc974..52ee40b65af3b 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStore.java
@@ -130,24 +130,27 @@ public long put(final Bytes key, final byte[] value, final long timestamp) {
         Objects.requireNonNull(key, "key cannot be null");
         validateStoreOpen();
 
-        if (timestamp < observedStreamTime - gracePeriod) {
-            expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
-            LOG.warn("Skipping record for expired put.");
-            return PUT_RETURN_CODE_NOT_PUT;
-        }
-        observedStreamTime = Math.max(observedStreamTime, timestamp);
-
-        final long foundTs = doPut(
-            versionedStoreClient,
-            observedStreamTime,
-            key,
-            value,
-            timestamp
-        );
+        synchronized (position) {
+            if (timestamp < observedStreamTime - gracePeriod) {
+                expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
+                LOG.warn("Skipping record for expired put.");
+                StoreQueryUtils.updatePosition(position, stateStoreContext);
+                return PUT_RETURN_CODE_NOT_PUT;
+            }
+            observedStreamTime = Math.max(observedStreamTime, timestamp);
+
+            final long foundTs = doPut(
+                versionedStoreClient,
+                observedStreamTime,
+                key,
+                value,
+                timestamp
+            );
 
-        StoreQueryUtils.updatePosition(position, stateStoreContext);
+            StoreQueryUtils.updatePosition(position, stateStoreContext);
 
-        return foundTs;
+            return foundTs;
+        }
     }
 
     @Override
@@ -155,26 +158,28 @@ public VersionedRecord<byte[]> delete(final Bytes key, final long timestamp) {
         Objects.requireNonNull(key, "key cannot be null");
         validateStoreOpen();
 
-        if (timestamp < observedStreamTime - gracePeriod) {
-            expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
-            LOG.warn("Skipping record for expired delete.");
-            return null;
-        }
+        synchronized (position) {
+            if (timestamp < observedStreamTime - gracePeriod) {
+                expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
+                LOG.warn("Skipping record for expired delete.");
+                return null;
+            }
 
-        final VersionedRecord<byte[]> existingRecord = get(key, timestamp);
+            final VersionedRecord<byte[]> existingRecord = get(key, timestamp);
 
-        observedStreamTime = Math.max(observedStreamTime, timestamp);
-        doPut(
-            versionedStoreClient,
-            observedStreamTime,
-            key,
-            null,
-            timestamp
-        );
+            observedStreamTime = Math.max(observedStreamTime, timestamp);
+            doPut(
+                versionedStoreClient,
+                observedStreamTime,
+                key,
+                null,
+                timestamp
+            );
 
-        StoreQueryUtils.updatePosition(position, stateStoreContext);
+            StoreQueryUtils.updatePosition(position, stateStoreContext);
 
-        return existingRecord;
+            return existingRecord;
+        }
     }
 
     @Override
@@ -361,11 +366,11 @@ public void init(final ProcessorContext context, final StateStore root) {
 
         metricsRecorder.init(ProcessorContextUtils.getMetricsImpl(context), context.taskId());
 
-        segmentStores.openExisting(context, observedStreamTime);
-
         final File positionCheckpointFile = new File(context.stateDir(), name() + ".position");
-        this.positionCheckpoint = new OffsetCheckpoint(positionCheckpointFile);
-        this.position = StoreQueryUtils.readPositionFromCheckpoint(positionCheckpoint);
+        positionCheckpoint = new OffsetCheckpoint(positionCheckpointFile);
+        position = StoreQueryUtils.readPositionFromCheckpoint(positionCheckpoint);
+        segmentStores.setPosition(position);
+        segmentStores.openExisting(context, observedStreamTime);
 
         // register and possibly restore the state from the logs
         stateStoreContext.register(
@@ -408,36 +413,39 @@ void restoreBatch(final Collection<ConsumerRecord<byte[], byte[]>> records) {
         // "segment" entry -- restoring a single changelog entry could require loading multiple
         // records into memory. how high this memory amplification will be is very much dependent
         // on the specific workload and the value of the "segment interval" parameter.
-        for (final ConsumerRecord<byte[], byte[]> record : records) {
-            if (record.timestamp() < observedStreamTime - gracePeriod) {
-                // record is older than grace period and was therefore never written to the store
-                continue;
-            }
-            // advance observed stream time as usual, for use in deciding whether records have
-            // exceeded the store's grace period and should be dropped.
-            observedStreamTime = Math.max(observedStreamTime, record.timestamp());
+        synchronized (position) {
+            for (final ConsumerRecord<byte[], byte[]> record : records) {
+                if (record.timestamp() < observedStreamTime - gracePeriod) {
+                    // record is older than grace period and was therefore never written to the store
+                    continue;
+                }
+                // advance observed stream time as usual, for use in deciding whether records have
+                // exceeded the store's grace period and should be dropped.
+                observedStreamTime = Math.max(observedStreamTime, record.timestamp());
 
-            ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
+                ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
                     record,
                     consistencyEnabled,
                     position
-            );
+                );
 
-            // put records to write buffer
-            doPut(
+                // put records to write buffer
+                doPut(
                     restoreClient,
                     endOfBatchStreamTime,
                     new Bytes(record.key()),
                     record.value(),
                     record.timestamp()
-            );
-        }
+                );
+            }
 
-        try {
-            restoreWriteBuffer.flush();
-        } catch (final RocksDBException e) {
-            throw new ProcessorStateException("Error restoring batch to store " + name, e);
+            try {
+                restoreWriteBuffer.flush();
+            } catch (final RocksDBException e) {
+                throw new ProcessorStateException("Error restoring batch to store " + name, e);
+            }
         }
+
     }
 
     private void validateStoreOpen() {
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/StoreQueryUtils.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/StoreQueryUtils.java
index 1609e8b2c5d5c..fa2081ad25bcd 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/StoreQueryUtils.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/StoreQueryUtils.java
@@ -125,28 +125,30 @@ public static <R> QueryResult<R> handleBasicQueries(
         final QueryResult<R> result;
 
         final QueryHandler handler = QUERY_HANDLER_MAP.get(query.getClass());
-        if (handler == null) {
-            result = QueryResult.forUnknownQueryType(query, store);
-        } else if (context == null || !isPermitted(position, positionBound, context.taskId().partition())) {
-            result = QueryResult.notUpToBound(
-                position,
-                positionBound,
-                context == null ? null : context.taskId().partition()
-            );
-        } else {
-            result = (QueryResult<R>) handler.apply(
-                query,
-                positionBound,
-                config,
-                store
-            );
-        }
-        if (config.isCollectExecutionInfo()) {
-            result.addExecutionInfo(
-                "Handled in " + store.getClass() + " in " + (System.nanoTime() - start) + "ns"
-            );
+        synchronized (position) {
+            if (handler == null) {
+                result = QueryResult.forUnknownQueryType(query, store);
+            } else if (context == null || !isPermitted(position, positionBound, context.taskId().partition())) {
+                result = QueryResult.notUpToBound(
+                    position,
+                    positionBound,
+                    context == null ? null : context.taskId().partition()
+                );
+            } else {
+                result = (QueryResult<R>) handler.apply(
+                    query,
+                    positionBound,
+                    config,
+                    store
+                );
+            }
+            if (config.isCollectExecutionInfo()) {
+                result.addExecutionInfo(
+                    "Handled in " + store.getClass() + " in " + (System.nanoTime() - start) + "ns"
+                );
+            }
+            result.setPosition(position.copy());
         }
-        result.setPosition(position);
         return result;
     }
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegment.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegment.java
index f0e4cf6132e09..1bf07afab779e 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegment.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegment.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 
 import java.io.File;
@@ -31,9 +32,11 @@ class TimestampedSegment extends RocksDBTimestampedStore implements Comparable<T
     TimestampedSegment(final String segmentName,
                        final String windowName,
                        final long id,
+                       final Position position,
                        final RocksDBMetricsRecorder metricsRecorder) {
         super(segmentName, windowName, metricsRecorder);
         this.id = id;
+        this.position = position;
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegments.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegments.java
index 2ea8da53ca72e..eabaca6096e5f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegments.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/TimestampedSegments.java
@@ -42,7 +42,7 @@ public TimestampedSegment getOrCreateSegment(final long segmentId,
             return segments.get(segmentId);
         } else {
             final TimestampedSegment newSegment =
-                new TimestampedSegment(segmentName(segmentId), name, segmentId, metricsRecorder);
+                new TimestampedSegment(segmentName(segmentId), name, segmentId, position, metricsRecorder);
 
             if (segments.put(segmentId, newSegment) != null) {
                 throw new IllegalStateException("TimestampedSegment already exists. Possible concurrent access.");
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/ConsistencyVectorIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/ConsistencyVectorIntegrationTest.java
index 3e809467b0879..e5d8e5ae1330a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/ConsistencyVectorIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/ConsistencyVectorIntegrationTest.java
@@ -38,8 +38,6 @@
 import org.apache.kafka.streams.query.StateQueryRequest;
 import org.apache.kafka.streams.query.StateQueryResult;
 import org.apache.kafka.streams.state.KeyValueStore;
-import org.apache.kafka.streams.state.QueryableStoreType;
-import org.apache.kafka.streams.state.ReadOnlyKeyValueStore;
 import org.apache.kafka.test.IntegrationTest;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
@@ -54,7 +52,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Objects;
 import java.util.Properties;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
@@ -63,7 +60,6 @@
 
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.safeUniqueTestName;
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
-import static org.apache.kafka.streams.state.QueryableStoreTypes.keyValueStore;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -76,6 +72,8 @@ public class ConsistencyVectorIntegrationTest {
     private static int port = 0;
     private static final String INPUT_TOPIC_NAME = "input-topic";
     private static final String TABLE_NAME = "source-table";
+    private static final int KEY = 1;
+    private static final int NUMBER_OF_MESSAGES = 100;
 
     public final EmbeddedKafkaCluster cluster = new EmbeddedKafkaCluster(NUM_BROKERS);
 
@@ -101,12 +99,9 @@ public void after() {
 
     @Test
     public void shouldHaveSamePositionBoundActiveAndStandBy() throws Exception {
-        final int batch1NumMessages = 100;
-        final int key = 1;
         final Semaphore semaphore = new Semaphore(0);
 
         final StreamsBuilder builder = new StreamsBuilder();
-        Objects.requireNonNull(TABLE_NAME, "name cannot be null");
 
         builder.table(INPUT_TOPIC_NAME, Consumed.with(Serdes.Integer(), Serdes.Integer()),
                       Materialized.<Integer, Integer, KeyValueStore<Bytes, byte[]>>as(TABLE_NAME)
@@ -123,30 +118,30 @@ public void shouldHaveSamePositionBoundActiveAndStandBy() throws Exception {
         try {
             startApplicationAndWaitUntilRunning(kafkaStreamsList);
 
-            produceValueRange(key, 0, batch1NumMessages);
+            produceValueRange();
 
             // Assert that all messages in the first batch were processed in a timely manner
-            assertThat(semaphore.tryAcquire(batch1NumMessages, 60, TimeUnit.SECONDS), is(equalTo(true)));
-
-            final QueryableStoreType<ReadOnlyKeyValueStore<Integer, Integer>> queryableStoreType = keyValueStore();
+            assertThat(
+                "Did not process all message in time.",
+                semaphore.tryAcquire(NUMBER_OF_MESSAGES, 120, TimeUnit.SECONDS), is(equalTo(true))
+            );
 
             // Assert that both active and standby have the same position bound
             final StateQueryRequest<Integer> request =
                 StateQueryRequest
                     .inStore(TABLE_NAME)
-                    .withQuery(KeyQuery.<Integer, Integer>withKey(key))
+                    .withQuery(KeyQuery.<Integer, Integer>withKey(KEY))
                     .withPositionBound(PositionBound.unbounded());
 
-            checkPosition(batch1NumMessages, request, kafkaStreams1);
-            checkPosition(batch1NumMessages, request, kafkaStreams2);
+            checkPosition(request, kafkaStreams1);
+            checkPosition(request, kafkaStreams2);
         } finally {
             kafkaStreams1.close();
             kafkaStreams2.close();
         }
     }
 
-    private void checkPosition(final int batch1NumMessages,
-                               final StateQueryRequest<Integer> request,
+    private void checkPosition(final StateQueryRequest<Integer> request,
                                final KafkaStreams kafkaStreams1) throws InterruptedException {
         final long maxWaitMs = TestUtils.DEFAULT_MAX_WAIT_MS;
         final long expectedEnd = System.currentTimeMillis() + maxWaitMs;
@@ -170,7 +165,7 @@ private void checkPosition(final int batch1NumMessages,
                     )
                 );
 
-                if (queryResult.getResult() == batch1NumMessages - 1) {
+                if (queryResult.getResult() == NUMBER_OF_MESSAGES - 1) {
                     // we're at the end of the input.
                     return;
                 }
@@ -187,14 +182,13 @@ private void checkPosition(final int batch1NumMessages,
 
     }
 
-
     private KafkaStreams createKafkaStreams(final StreamsBuilder builder, final Properties config) {
         final KafkaStreams streams = new KafkaStreams(builder.build(config), config);
         streamsToCleanup.add(streams);
         return streams;
     }
 
-    private void produceValueRange(final int key, final int start, final int endExclusive) {
+    private void produceValueRange() {
         final Properties producerProps = new Properties();
         producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
         producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, IntegerSerializer.class);
@@ -202,8 +196,8 @@ private void produceValueRange(final int key, final int start, final int endExcl
 
         IntegrationTestUtils.produceKeyValuesSynchronously(
             INPUT_TOPIC_NAME,
-            IntStream.range(start, endExclusive)
-                     .mapToObj(i -> KeyValue.pair(key, i))
+            IntStream.range(0, NUMBER_OF_MESSAGES)
+                     .mapToObj(i -> KeyValue.pair(KEY, i))
                      .collect(Collectors.toList()),
             producerProps,
             mockTime
@@ -216,8 +210,8 @@ private Properties streamsConfiguration(final String safeTestName) {
         config.put(StreamsConfig.APPLICATION_SERVER_CONFIG, "localhost:" + (++port));
         config.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
         config.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
-        config.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
-        config.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
+        config.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
+        config.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
         config.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
         config.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 100);
         config.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 200);
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/IQv2IntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/IQv2IntegrationTest.java
index becbe301ecb0d..edfc27e458f08 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/IQv2IntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/IQv2IntegrationTest.java
@@ -319,28 +319,36 @@ public KeyValueStore<Bytes, byte[]> get() {
 
                         @Override
                         public void put(final Bytes key, final byte[] value) {
-                            map.put(key, value);
-                            StoreQueryUtils.updatePosition(position,  context);
+                            synchronized (position) {
+                                map.put(key, value);
+                                StoreQueryUtils.updatePosition(position, context);
+                            }
                         }
 
                         @Override
                         public byte[] putIfAbsent(final Bytes key, final byte[] value) {
-                            StoreQueryUtils.updatePosition(position,  context);
-                            return map.putIfAbsent(key, value);
+                            synchronized (position) {
+                                StoreQueryUtils.updatePosition(position, context);
+                                return map.putIfAbsent(key, value);
+                            }
                         }
 
                         @Override
                         public void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
-                            StoreQueryUtils.updatePosition(position,  context);
-                            for (final KeyValue<Bytes, byte[]> entry : entries) {
-                                map.put(entry.key, entry.value);
+                            synchronized (position) {
+                                StoreQueryUtils.updatePosition(position, context);
+                                for (final KeyValue<Bytes, byte[]> entry : entries) {
+                                    map.put(entry.key, entry.value);
+                                }
                             }
                         }
 
                         @Override
                         public byte[] delete(final Bytes key) {
-                            StoreQueryUtils.updatePosition(position,  context);
-                            return map.remove(key);
+                            synchronized (position) {
+                                StoreQueryUtils.updatePosition(position, context);
+                                return map.remove(key);
+                            }
                         }
 
                         @Override
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/StoreUpgradeIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/StoreUpgradeIntegrationTest.java
index c6ed805ac4c0b..5db747b1b46c9 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/StoreUpgradeIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/StoreUpgradeIntegrationTest.java
@@ -394,7 +394,8 @@ private <K> void verifyCountWithTimestamp(final K key,
                 }
             },
             60_000L,
-            "Could not get expected result in time.");
+            5_000L,
+            () -> "Could not get expected result in time.");
     }
 
     private <K> void verifyCountWithSurrogateTimestamp(final K key,
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/KeyValueSegmentTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/KeyValueSegmentTest.java
index e4543d2b56b8a..a8d2eb2bbce17 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/KeyValueSegmentTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/KeyValueSegmentTest.java
@@ -22,6 +22,7 @@
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 import org.apache.kafka.test.TestUtils;
 import org.junit.Before;
@@ -60,7 +61,7 @@ public void setUp() {
 
     @Test
     public void shouldDeleteStateDirectoryOnDestroy() throws Exception {
-        final KeyValueSegment segment = new KeyValueSegment("segment", "window", 0L, metricsRecorder);
+        final KeyValueSegment segment = new KeyValueSegment("segment", "window", 0L, Position.emptyPosition(),  metricsRecorder);
         final String directoryPath = TestUtils.tempDirectory().getAbsolutePath();
         final File directory = new File(directoryPath);
 
@@ -82,10 +83,10 @@ public void shouldDeleteStateDirectoryOnDestroy() throws Exception {
 
     @Test
     public void shouldBeEqualIfIdIsEqual() {
-        final KeyValueSegment segment = new KeyValueSegment("anyName", "anyName", 0L, metricsRecorder);
+        final KeyValueSegment segment = new KeyValueSegment("anyName", "anyName", 0L, Position.emptyPosition(), metricsRecorder);
         final KeyValueSegment segmentSameId =
-            new KeyValueSegment("someOtherName", "someOtherName", 0L, metricsRecorder);
-        final KeyValueSegment segmentDifferentId = new KeyValueSegment("anyName", "anyName", 1L, metricsRecorder);
+            new KeyValueSegment("someOtherName", "someOtherName", 0L, Position.emptyPosition(), metricsRecorder);
+        final KeyValueSegment segmentDifferentId = new KeyValueSegment("anyName", "anyName", 1L, Position.emptyPosition(), metricsRecorder);
 
         assertThat(segment, equalTo(segment));
         assertThat(segment, equalTo(segmentSameId));
@@ -98,10 +99,10 @@ public void shouldBeEqualIfIdIsEqual() {
 
     @Test
     public void shouldHashOnSegmentIdOnly() {
-        final KeyValueSegment segment = new KeyValueSegment("anyName", "anyName", 0L, metricsRecorder);
+        final KeyValueSegment segment = new KeyValueSegment("anyName", "anyName", 0L, Position.emptyPosition(), metricsRecorder);
         final KeyValueSegment segmentSameId =
-            new KeyValueSegment("someOtherName", "someOtherName", 0L, metricsRecorder);
-        final KeyValueSegment segmentDifferentId = new KeyValueSegment("anyName", "anyName", 1L, metricsRecorder);
+            new KeyValueSegment("someOtherName", "someOtherName", 0L, Position.emptyPosition(), metricsRecorder);
+        final KeyValueSegment segmentDifferentId = new KeyValueSegment("anyName", "anyName", 1L, Position.emptyPosition(), metricsRecorder);
 
         final Set<KeyValueSegment> set = new HashSet<>();
         assertTrue(set.add(segment));
@@ -113,9 +114,9 @@ public void shouldHashOnSegmentIdOnly() {
 
     @Test
     public void shouldCompareSegmentIdOnly() {
-        final KeyValueSegment segment1 = new KeyValueSegment("a", "C", 50L, metricsRecorder);
-        final KeyValueSegment segment2 = new KeyValueSegment("b", "B", 100L, metricsRecorder);
-        final KeyValueSegment segment3 = new KeyValueSegment("c", "A", 0L, metricsRecorder);
+        final KeyValueSegment segment1 = new KeyValueSegment("a", "C", 50L, Position.emptyPosition(), metricsRecorder);
+        final KeyValueSegment segment2 = new KeyValueSegment("b", "B", 100L, Position.emptyPosition(), metricsRecorder);
+        final KeyValueSegment segment3 = new KeyValueSegment("c", "A", 0L, Position.emptyPosition(), metricsRecorder);
 
         assertThat(segment1.compareTo(segment1), equalTo(0));
         assertThat(segment1.compareTo(segment2), equalTo(-1));
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegmentsTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegmentsTest.java
index edc53e4c7001a..a28c55f70ded3 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegmentsTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegmentsTest.java
@@ -32,6 +32,7 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.processor.internals.MockStreamsMetrics;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.KeyValueIterator;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 import org.apache.kafka.test.InternalMockProcessorContext;
@@ -69,6 +70,7 @@ public void setUp() {
             SEGMENT_INTERVAL,
             new RocksDBMetricsRecorder(METRICS_SCOPE, STORE_NAME)
         );
+        segments.setPosition(Position.emptyPosition());
         segments.openExisting(context, 0L);
     }
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/SegmentIteratorTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/SegmentIteratorTest.java
index 31ce3c696368a..70ca80edb4d36 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/SegmentIteratorTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/SegmentIteratorTest.java
@@ -23,6 +23,7 @@
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.processor.StateStoreContext;
 import org.apache.kafka.streams.processor.internals.MockStreamsMetrics;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 import org.apache.kafka.test.InternalMockProcessorContext;
 import org.apache.kafka.test.MockRecordCollector;
@@ -46,9 +47,9 @@ public class SegmentIteratorTest {
     private final RocksDBMetricsRecorder rocksDBMetricsRecorder =
         new RocksDBMetricsRecorder("metrics-scope", "store-name");
     private final KeyValueSegment segmentOne =
-        new KeyValueSegment("one", "one", 0, rocksDBMetricsRecorder);
+        new KeyValueSegment("one", "one", 0, Position.emptyPosition(), rocksDBMetricsRecorder);
     private final KeyValueSegment segmentTwo =
-        new KeyValueSegment("two", "window", 1, rocksDBMetricsRecorder);
+        new KeyValueSegment("two", "window", 1, Position.emptyPosition(), rocksDBMetricsRecorder);
     private final HasNextCondition hasNextCondition = Iterator::hasNext;
 
     private SegmentIterator<KeyValueSegment> iterator = null;
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedSegmentTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedSegmentTest.java
index 062fb808ed826..dd31d294c6147 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedSegmentTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/TimestampedSegmentTest.java
@@ -22,6 +22,7 @@
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
+import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 import org.apache.kafka.test.TestUtils;
 import org.junit.Before;
@@ -60,7 +61,7 @@ public void setUp() {
 
     @Test
     public void shouldDeleteStateDirectoryOnDestroy() throws Exception {
-        final TimestampedSegment segment = new TimestampedSegment("segment", "window", 0L, metricsRecorder);
+        final TimestampedSegment segment = new TimestampedSegment("segment", "window", 0L, Position.emptyPosition(), metricsRecorder);
         final String directoryPath = TestUtils.tempDirectory().getAbsolutePath();
         final File directory = new File(directoryPath);
 
@@ -82,11 +83,11 @@ public void shouldDeleteStateDirectoryOnDestroy() throws Exception {
 
     @Test
     public void shouldBeEqualIfIdIsEqual() {
-        final TimestampedSegment segment = new TimestampedSegment("anyName", "anyName", 0L, metricsRecorder);
+        final TimestampedSegment segment = new TimestampedSegment("anyName", "anyName", 0L, Position.emptyPosition(), metricsRecorder);
         final TimestampedSegment segmentSameId =
-            new TimestampedSegment("someOtherName", "someOtherName", 0L, metricsRecorder);
+            new TimestampedSegment("someOtherName", "someOtherName", 0L, Position.emptyPosition(), metricsRecorder);
         final TimestampedSegment segmentDifferentId =
-            new TimestampedSegment("anyName", "anyName", 1L, metricsRecorder);
+            new TimestampedSegment("anyName", "anyName", 1L, Position.emptyPosition(), metricsRecorder);
 
         assertThat(segment, equalTo(segment));
         assertThat(segment, equalTo(segmentSameId));
@@ -101,11 +102,11 @@ public void shouldBeEqualIfIdIsEqual() {
 
     @Test
     public void shouldHashOnSegmentIdOnly() {
-        final TimestampedSegment segment = new TimestampedSegment("anyName", "anyName", 0L, metricsRecorder);
+        final TimestampedSegment segment = new TimestampedSegment("anyName", "anyName", 0L, Position.emptyPosition(), metricsRecorder);
         final TimestampedSegment segmentSameId =
-            new TimestampedSegment("someOtherName", "someOtherName", 0L, metricsRecorder);
+            new TimestampedSegment("someOtherName", "someOtherName", 0L, Position.emptyPosition(), metricsRecorder);
         final TimestampedSegment segmentDifferentId =
-            new TimestampedSegment("anyName", "anyName", 1L, metricsRecorder);
+            new TimestampedSegment("anyName", "anyName", 1L, Position.emptyPosition(), metricsRecorder);
 
         final Set<TimestampedSegment> set = new HashSet<>();
         assertTrue(set.add(segment));
@@ -119,9 +120,9 @@ public void shouldHashOnSegmentIdOnly() {
 
     @Test
     public void shouldCompareSegmentIdOnly() {
-        final TimestampedSegment segment1 = new TimestampedSegment("a", "C", 50L, metricsRecorder);
-        final TimestampedSegment segment2 = new TimestampedSegment("b", "B", 100L, metricsRecorder);
-        final TimestampedSegment segment3 = new TimestampedSegment("c", "A", 0L, metricsRecorder);
+        final TimestampedSegment segment1 = new TimestampedSegment("a", "C", 50L, Position.emptyPosition(), metricsRecorder);
+        final TimestampedSegment segment2 = new TimestampedSegment("b", "B", 100L, Position.emptyPosition(), metricsRecorder);
+        final TimestampedSegment segment3 = new TimestampedSegment("c", "A", 0L, Position.emptyPosition(), metricsRecorder);
 
         assertThat(segment1.compareTo(segment1), equalTo(0));
         assertThat(segment1.compareTo(segment2), equalTo(-1));

From c1fdcb2a274d4fb822510f555a01d24f7b8f978a Mon Sep 17 00:00:00 2001
From: Artem Livshits <84364232+artemlivshits@users.noreply.github.com>
Date: Tue, 20 Feb 2024 15:13:03 -0800
Subject: [PATCH 059/258] MINOR: extend transaction unit test to validate drain
 (#15320)

There is a test already that checks that transactional messages are not
drained when partition is not added, this change just logically
completes the test to also show that messages can be drained after
parition is added.

Reviewers: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Tzu-Li (Gordon) Tai <tzulitai@apache.org>, Justine Olshan <jolshan@confluent.io>
---
 .../internals/TransactionManagerTest.java       | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index 5a2527deee99e..de02738a34ad6 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -2545,6 +2545,23 @@ public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedExc
         // We shouldn't drain batches which haven't been added to the transaction yet.
         assertTrue(drainedBatches.containsKey(node1.id()));
         assertTrue(drainedBatches.get(node1.id()).isEmpty());
+
+        // Let's now add the partition, flush and try to drain again.
+        transactionManager.maybeAddPartition(tp0);
+        accumulator.beginFlush();
+
+        drainedBatches = accumulator.drain(metadataMock, nodes, Integer.MAX_VALUE,
+            time.milliseconds());
+
+        // We still shouldn't drain batches because the partition call didn't complete yet.
+        assertTrue(drainedBatches.containsKey(node1.id()));
+        assertTrue(drainedBatches.get(node1.id()).isEmpty());
+        assertTrue(accumulator.hasUndrained());
+
+        // Now prepare response to complete the partition addition.
+        // We should now be able to drain the request.
+        prepareAddPartitionsToTxn(tp0, Errors.NONE);
+        runUntil(() -> !accumulator.hasUndrained());
     }
 
     @Test

From fcbfd3412eb746a0c81374eb55ad0f73de6b1e71 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Wed, 21 Feb 2024 09:01:53 +0100
Subject: [PATCH 060/258] KAFKA-16284: Fix performance regression in RocksDB
 (#15393)

A performance regression introduced in commit 5bc3aa4 reduces the write performance in RocksDB by ~3x. The bug is that we fail to pass the WriteOptions that disable the write-ahead log into the DB accessor.

For testing, the time to write 10 times 1 Million records into one RocksDB each were measured:

Before 5bc3aa4: 7954ms, 12933ms
After 5bc3aa4: 30345ms, 31992ms
After 5bc3aa4 with this fix: 8040ms, 10563ms
On current trunk with this fix: 9508ms, 10441ms

Reviewers: Bruno Cadonna <bruno@confluent.io>, Nick Telford <nick.telford@gmail.com>
---
 .../kafka/streams/state/internals/RocksDBStore.java  | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
index 0dfae3499c10f..2d3e440b13c9c 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
@@ -252,7 +252,7 @@ void openDB(final Map<String, Object> configs, final File stateDir) {
         // with the measurements from Rocks DB
         setupStatistics(configs, dbOptions);
         openRocksDB(dbOptions, columnFamilyOptions);
-        dbAccessor = new DirectDBAccessor(db, fOptions);
+        dbAccessor = new DirectDBAccessor(db, fOptions, wOptions);
         open = true;
 
         addValueProvidersToMetricsRecorder();
@@ -748,10 +748,12 @@ static class DirectDBAccessor implements DBAccessor {
 
         private final RocksDB db;
         private final FlushOptions flushOptions;
+        private final WriteOptions wOptions;
 
-        DirectDBAccessor(final RocksDB db, final FlushOptions flushOptions) {
+        DirectDBAccessor(final RocksDB db, final FlushOptions flushOptions, final WriteOptions wOptions) {
             this.db = db;
             this.flushOptions = flushOptions;
+            this.wOptions = wOptions;
         }
 
         @Override
@@ -771,17 +773,17 @@ public RocksIterator newIterator(final ColumnFamilyHandle columnFamily) {
 
         @Override
         public void put(final ColumnFamilyHandle columnFamily, final byte[] key, final byte[] value) throws RocksDBException {
-            db.put(columnFamily, key, value);
+            db.put(columnFamily, wOptions, key, value);
         }
 
         @Override
         public void delete(final ColumnFamilyHandle columnFamily, final byte[] key) throws RocksDBException {
-            db.delete(columnFamily, key);
+            db.delete(columnFamily, wOptions, key);
         }
 
         @Override
         public void deleteRange(final ColumnFamilyHandle columnFamily, final byte[] from, final byte[] to) throws RocksDBException {
-            db.deleteRange(columnFamily, from, to);
+            db.deleteRange(columnFamily, wOptions, from, to);
         }
 
         @Override

From 9118ad653fa4a076bf1b774c22aa015dfc2c2aab Mon Sep 17 00:00:00 2001
From: Paolo Patierno <ppatierno@live.com>
Date: Wed, 21 Feb 2024 10:43:05 +0100
Subject: [PATCH 061/258] Additional fix on the rollback migration
 documentation (#15317)

This is an additional PR to fix rollback migration documentation related to #15287

Reviewers: Luke Chen <showuon@gmail.com>
---
 docs/ops.html | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/docs/ops.html b/docs/ops.html
index cf039c752207c..7e9d9630e2574 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -4063,8 +4063,10 @@ <h3>Reverting to ZooKeeper mode During the Migration</h3>
           </ul>
         </td>
         <td>
-          It is important to perform the <code>zookeeper-shell.sh</code> step quickly, to minimize the amount of
-          time that the cluster lacks a controller.
+          It is important to perform the <code>zookeeper-shell.sh</code> step <b>quickly</b>, to minimize the amount of
+          time that the cluster lacks a controller. Until the <code> /controller</code> znode is deleted,
+          you can also ignore any errors in the broker log about failing to connect to the Kraft controller.
+          Those error logs should disappear after second roll to pure zookeeper mode.
         </td>
       </tr>
       <tr>
@@ -4072,7 +4074,8 @@ <h3>Reverting to ZooKeeper mode During the Migration</h3>
         <td>
           <ul>
             <li>
-              On each broker, remove the <code>process.roles</code> configuration, and
+              On each broker, remove the <code>process.roles</code> configuration,
+              replace the <code>node.id</code> with <code>broker.id</code> and
               restore the <code>zookeeper.connect</code> configuration to its previous value.
               If your cluster requires other ZooKeeper configurations for brokers, such as
               <code>zookeeper.ssl.protocol</code>, re-add those configurations as well.
@@ -4088,7 +4091,7 @@ <h3>Reverting to ZooKeeper mode During the Migration</h3>
             <li>
               On each broker, remove the <code>zookeeper.metadata.migration.enable</code>,
               <code>controller.listener.names</code>, and <code>controller.quorum.voters</code>
-              configurations. Replace <code>node.id</code> with <code>broker.id</code>.
+              configurations.
               Then perform a second rolling restart of all brokers.
             </li>
             <li>
@@ -4100,7 +4103,9 @@ <h3>Reverting to ZooKeeper mode During the Migration</h3>
           <ul>
             <li>
               It is important to perform the <code>zookeeper-shell.sh</code> step <b>quickly</b>, to minimize the amount of
-              time that the cluster lacks a controller.
+              time that the cluster lacks a controller. Until the <code> /controller</code> znode is deleted,
+              you can also ignore any errors in the broker log about failing to connect to the Kraft controller.
+              Those error logs should disappear after second roll to pure zookeeper mode.
             </li>
             <li>
               Make sure that on the first cluster roll, <code>zookeeper.metadata.migration.enable</code> remains set to

From 77ba06fa620ad9cc42af92f42354661308676135 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Wed, 21 Feb 2024 05:08:37 -0500
Subject: [PATCH 062/258] KAFKA-16033: Commit retry logic fixes (#15357)

This change modifies the commit manager for improved retry logic & fixing bugs:

- defines high level functions for each of the different types of commit: commitSync, commitAsync, autoCommitSync (used from consumer close), autoCommitAsync (on interval), autoCommitNow (before revocation).
 - moves retry logic to these caller functions, keeping a common response error handling that propagates errors that each caller functions retry as it needs.

Fixes the following issues:

- auto-commit before revocation should retry with latest consumed offsets
- auto-commit before revocation should only reset the timer once, when the rebalance completes
- StaleMemberEpoch error (fatal) is considered retriable only when committing offsets before revocation, where it is retried with backoff if the member has a valid epoch. All other commits will fail fatally on stale epoch. Note that auto commit on the interval (autoCommitAsync) does not have any specific retry logic for the stale epoch, but will effectively retry on the next interval (as it does for any other fatal error)
- fix duplicated and noisy logs for auto-commit

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/AsyncKafkaConsumer.java         |  24 +-
 .../internals/CommitRequestManager.java       | 669 ++++++++++--------
 .../internals/MembershipManagerImpl.java      |  12 +-
 .../internals/events/ApplicationEvent.java    |   2 +-
 .../events/ApplicationEventProcessor.java     |  32 +-
 .../events/AsyncCommitApplicationEvent.java   |  39 +
 .../events/CommitApplicationEvent.java        |  34 +-
 .../events/SyncCommitApplicationEvent.java    |  52 ++
 .../internals/AsyncKafkaConsumerTest.java     |  71 +-
 .../internals/CommitRequestManagerTest.java   | 291 +++++---
 .../internals/ConsumerNetworkThreadTest.java  |  26 +-
 .../internals/MembershipManagerImplTest.java  |   4 +-
 12 files changed, 719 insertions(+), 537 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index 6781e0b73cb80..d2e4788a1bc18 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -40,6 +40,7 @@
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
@@ -57,6 +58,7 @@
 import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.TopicMetadataApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
@@ -755,9 +757,8 @@ public void commitAsync(OffsetCommitCallback callback) {
     public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCommitCallback callback) {
         acquireAndEnsureOpen();
         try {
-            // Commit without timer to indicate that the commit should be triggered without
-            // waiting for a response.
-            CompletableFuture<Void> future = commit(offsets, false, Optional.empty());
+            AsyncCommitApplicationEvent asyncCommitEvent = new AsyncCommitApplicationEvent(offsets);
+            CompletableFuture<Void> future = commit(asyncCommitEvent);
             future.whenComplete((r, t) -> {
 
                 if (t == null) {
@@ -778,14 +779,12 @@ public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCo
         }
     }
 
-    // Visible for testing
-    CompletableFuture<Void> commit(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                   final boolean isWakeupable,
-                                   final Optional<Long> retryTimeoutMs) {
+    private CompletableFuture<Void> commit(final CommitApplicationEvent commitEvent) {
         maybeInvokeCommitCallbacks();
         maybeThrowFencedInstanceException();
         maybeThrowInvalidGroupIdException();
 
+        Map<TopicPartition, OffsetAndMetadata> offsets = commitEvent.offsets();
         log.debug("Committing offsets: {}", offsets);
         offsets.forEach(this::updateLastSeenEpochIfNewer);
 
@@ -793,11 +792,6 @@ CompletableFuture<Void> commit(final Map<TopicPartition, OffsetAndMetadata> offs
             return CompletableFuture.completedFuture(null);
         }
 
-        final CommitApplicationEvent commitEvent = new CommitApplicationEvent(offsets, retryTimeoutMs);
-        if (isWakeupable) {
-            // the task can only be woken up if the top level API call is commitSync
-            wakeupTrigger.setActiveTask(commitEvent.future());
-        }
         applicationEventHandler.add(commitEvent);
         return commitEvent.future();
     }
@@ -1344,9 +1338,9 @@ public void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration
         long commitStart = time.nanoseconds();
         try {
             Timer requestTimer = time.timer(timeout.toMillis());
-            // Commit with a timer to control how long the request should be retried until it
-            // gets a successful response or non-retriable error.
-            CompletableFuture<Void> commitFuture = commit(offsets, true, Optional.of(timeout.toMillis()));
+            SyncCommitApplicationEvent syncCommitEvent = new SyncCommitApplicationEvent(offsets, timeout.toMillis());
+            CompletableFuture<Void> commitFuture = commit(syncCommitEvent);
+            wakeupTrigger.setActiveTask(commitFuture);
             ConsumerUtils.getResult(commitFuture, requestTimer);
             interceptors.onCommit(offsets);
         } finally {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
index 704da62178c46..9206783d561be 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
@@ -27,6 +27,7 @@
 import org.apache.kafka.common.errors.DisconnectException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.errors.StaleMemberEpochException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.UnstableOffsetCommitException;
@@ -168,7 +169,7 @@ public NetworkClientDelegate.PollResult poll(final long currentTimeMs) {
             return drainPendingOffsetCommitRequests();
         }
 
-        maybeAutoCommitAllConsumedAsync();
+        maybeAutoCommitAsync();
         if (!pendingRequests.hasUnsentRequests())
             return EMPTY;
 
@@ -204,126 +205,313 @@ private static long findMinTime(final Collection<? extends RequestState> request
     }
 
     /**
-     * Generate a request to commit offsets if auto-commit is enabled. The request will be
-     * returned to be sent out on the next call to {@link #poll(long)}. This will only generate a
-     * request if there is no other commit request already in-flight, and if the commit interval
-     * has elapsed.
+     * Generate a request to commit consumed offsets. Add the request to the queue of pending
+     * requests to be sent out on the next call to {@link #poll(long)}. If there are empty
+     * offsets to commit, no request will be generated and a completed future will be returned.
      *
-     * @param offsets           Offsets to commit
-     * @param expirationTimeMs  Time until which the request will continue to be retried if it
-     *                          fails with a retriable error. If not present, the request will be
-     *                          sent but not retried.
-     * @param checkInterval     True if the auto-commit interval expiration should be checked for
-     *                          sending a request. If true, the request will be sent only if the
-     *                          auto-commit interval has expired. Pass false to
-     *                          send the auto-commit request regardless of the interval (ex.
-     *                          auto-commit before rebalance).
-     * @param retryOnStaleEpoch True if the request should be retried in case it fails with
-     *                          {@link Errors#STALE_MEMBER_EPOCH}.
-     * @return Future that will complete when a response is received for the request, or a
-     * completed future if no request is generated.
+     * @param requestState Commit request
+     * @return Future containing the offsets that were committed, or an error if the request
+     * failed.
      */
-    private CompletableFuture<Void> maybeAutoCommit(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                                    final Optional<Long> expirationTimeMs,
-                                                    boolean checkInterval,
-                                                    boolean retryOnStaleEpoch) {
-        if (!autoCommitEnabled()) {
-            log.debug("Skipping auto-commit because auto-commit config is not enabled.");
-            return CompletableFuture.completedFuture(null);
-        }
-
+    private CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> requestAutoCommit(final OffsetCommitRequestState requestState) {
         AutoCommitState autocommit = autoCommitState.get();
-        if (checkInterval && !autocommit.shouldAutoCommit()) {
-            return CompletableFuture.completedFuture(null);
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result;
+        if (requestState.offsets.isEmpty()) {
+            result = CompletableFuture.completedFuture(Collections.emptyMap());
+        } else {
+            autocommit.setInflightCommitStatus(true);
+            OffsetCommitRequestState request = pendingRequests.addOffsetCommitRequest(requestState);
+            result = request.future;
+            result.whenComplete(autoCommitCallback(request.offsets));
         }
-
-        autocommit.resetTimer();
-        autocommit.setInflightCommitStatus(true);
-        CompletableFuture<Void> result = addOffsetCommitRequest(offsets, expirationTimeMs, retryOnStaleEpoch)
-            .whenComplete(autoCommitCallback(offsets));
         return result;
     }
 
     /**
-     * If auto-commit is enabled, this will generate a commit offsets request for all assigned
-     * partitions and their current positions. Note on auto-commit timers: this will reset the
-     * auto-commit timer to the interval before issuing the async commit, and when the async commit
-     * completes, it will reset the auto-commit timer with the exponential backoff if the request
-     * failed with a retriable error.
-     *
-     * @return Future that will complete when a response is received for the request, or a
-     * completed future if no request is generated.
+     * If auto-commit is enabled, and the auto-commit interval has expired, this will generate and
+     * enqueue a request to commit all consumed offsets, and will reset the auto-commit timer to the
+     * interval. The request will be sent on the next call to {@link #poll(long)}.
+     * <p/>
+     * If the request completes with a retriable error, this will reset the auto-commit timer with
+     * the exponential backoff. If it fails with a non-retriable error, no action is taken, so
+     * the next commit will be generated when the interval expires.
      */
-    public CompletableFuture<Void> maybeAutoCommitAllConsumedAsync() {
-        if (!autoCommitEnabled()) {
-            // Early return to ensure that no action/logging is performed.
-            return CompletableFuture.completedFuture(null);
+    public void maybeAutoCommitAsync() {
+        if (autoCommitEnabled() && autoCommitState.get().shouldAutoCommit()) {
+            OffsetCommitRequestState requestState = createOffsetCommitRequest(
+                subscriptions.allConsumed(),
+                Optional.empty());
+            CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = requestAutoCommit(requestState);
+            // Reset timer to the interval (even if no request was generated), but ensure that if
+            // the request completes with a retriable error, the timer is reset to send the next
+            // auto-commit after the backoff expires.
+            resetAutoCommitTimer();
+            maybeResetTimerWithBackoff(result);
         }
-        Map<TopicPartition, OffsetAndMetadata> offsets = subscriptions.allConsumed();
-        CompletableFuture<Void> result = maybeAutoCommit(offsets, Optional.empty(), true, true);
-        result.whenComplete((__, error) -> {
+    }
+
+    /**
+     * Reset auto-commit timer to retry with backoff if the future failed with a RetriableCommitFailedException.
+     */
+    private void maybeResetTimerWithBackoff(final CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result) {
+        result.whenComplete((offsets, error) -> {
             if (error != null) {
                 if (error instanceof RetriableCommitFailedException) {
                     log.debug("Asynchronous auto-commit of offsets {} failed due to retriable error.", offsets, error);
                     resetAutoCommitTimer(retryBackoffMs);
                 } else {
-                    log.warn("Asynchronous auto-commit of offsets {} failed: {}", offsets, error.getMessage());
+                    log.debug("Asynchronous auto-commit of offsets {} failed: {}", offsets, error.getMessage());
                 }
             } else {
                 log.debug("Completed asynchronous auto-commit of offsets {}", offsets);
             }
         });
-
-        return result;
     }
 
     /**
-     * Commit consumed offsets if auto-commit is enabled. Retry while the timer is not expired,
-     * until the request succeeds or fails with a fatal error.
+     * Commit consumed offsets if auto-commit is enabled, regardless of the auto-commit interval.
+     * This is used for committing offsets before revoking partitions. This will retry committing
+     * the latest offsets until the request succeeds, fails with a fatal error, or the timeout
+     * expires. Note that this considers {@link Errors#STALE_MEMBER_EPOCH} as a retriable error,
+     * and will retry it including the latest member ID and epoch received from the broker.
+     *
+     * @return Future that will complete when the offsets are successfully committed. It will
+     * complete exceptionally if the commit fails with a non-retriable error, or if the retry
+     * timeout expires.
      */
-    public CompletableFuture<Void> maybeAutoCommitAllConsumedNow(
-        final Optional<Long> expirationTimeMs,
-        final boolean retryOnStaleEpoch) {
-        return maybeAutoCommit(subscriptions.allConsumed(), expirationTimeMs, false, retryOnStaleEpoch);
+    public CompletableFuture<Void> maybeAutoCommitSyncNow(final long retryExpirationTimeMs) {
+        if (!autoCommitEnabled()) {
+            return CompletableFuture.completedFuture(null);
+        }
+
+        CompletableFuture<Void> result = new CompletableFuture<>();
+        OffsetCommitRequestState requestState =
+            createOffsetCommitRequest(subscriptions.allConsumed(), Optional.of(retryExpirationTimeMs));
+        autoCommitSyncNowWithRetries(requestState, result);
+        return result;
     }
 
-    private BiConsumer<? super Void, ? super Throwable> autoCommitCallback(final Map<TopicPartition, OffsetAndMetadata> allConsumedOffsets) {
+    private void autoCommitSyncNowWithRetries(OffsetCommitRequestState requestAttempt,
+                                              CompletableFuture<Void> result) {
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> commitAttempt = requestAutoCommit(requestAttempt);
+        commitAttempt.whenComplete((committedOffsets, error) -> {
+            if (error == null) {
+                result.complete(null);
+            } else {
+                if (error instanceof RetriableException || isStaleEpochErrorAndValidEpochAvailable(error)) {
+                    if (error instanceof TimeoutException && requestAttempt.isExpired) {
+                        log.debug("Auto-commit sync timed out and won't be retried anymore");
+                        result.completeExceptionally(error);
+                    } else {
+                        // Make sure the auto-commit is retries with the latest offsets
+                        requestAttempt.offsets = subscriptions.allConsumed();
+                        requestAttempt.resetFuture();
+                        autoCommitSyncNowWithRetries(requestAttempt, result);
+                    }
+                } else {
+                    log.debug("Auto-commit sync failed with non-retriable error", error);
+                    result.completeExceptionally(error);
+                }
+            }
+        });
+    }
+
+    /**
+     * Clear the inflight auto-commit flag and log auto-commit completion status.
+     */
+    private BiConsumer<? super Map<TopicPartition, OffsetAndMetadata>, ? super Throwable> autoCommitCallback(final Map<TopicPartition, OffsetAndMetadata> allConsumedOffsets) {
         return (response, throwable) -> {
             autoCommitState.ifPresent(autoCommitState -> autoCommitState.setInflightCommitStatus(false));
             if (throwable == null) {
                 offsetCommitCallbackInvoker.enqueueInterceptorInvocation(allConsumedOffsets);
-                log.debug("Completed asynchronous auto-commit of offsets {}", allConsumedOffsets);
+                log.debug("Completed auto-commit of offsets {}", allConsumedOffsets);
             } else if (throwable instanceof RetriableCommitFailedException) {
-                log.debug("Asynchronous auto-commit of offsets {} failed due to retriable error: {}",
+                log.debug("Auto-commit of offsets {} failed due to retriable error: {}",
                         allConsumedOffsets, throwable.getMessage());
             } else {
-                log.warn("Asynchronous auto-commit of offsets {} failed", allConsumedOffsets, throwable);
+                log.warn("Auto-commit of offsets {} failed", allConsumedOffsets, throwable);
             }
         };
     }
 
     /**
-     * Handles {@link org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent}. It creates an
-     * {@link OffsetCommitRequestState} and enqueue it to send later.
+     * Generate a request to commit offsets without retrying, even if it fails with a retriable
+     * error. The generated request will be added to the queue to be sent on the next call to
+     * {@link #poll(long)}.
+     *
+     * @param offsets Offsets to commit per partition.
+     * @return Future that will complete when a response is received, successfully or
+     * exceptionally depending on the response. If the request fails with a retriable error, the
+     * future will be completed with a {@link RetriableCommitFailedException}.
      */
-    public CompletableFuture<Void> addOffsetCommitRequest(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                                          final Optional<Long> expirationTimeMs,
-                                                          final boolean retryOnStaleEpoch) {
+    public CompletableFuture<Void> commitAsync(final Map<TopicPartition, OffsetAndMetadata> offsets) {
         if (offsets.isEmpty()) {
             log.debug("Skipping commit of empty offsets");
             return CompletableFuture.completedFuture(null);
         }
-        return pendingRequests.addOffsetCommitRequest(offsets, expirationTimeMs, retryOnStaleEpoch).future;
+        OffsetCommitRequestState commitRequest = createOffsetCommitRequest(offsets, Optional.empty());
+        pendingRequests.addOffsetCommitRequest(commitRequest);
+
+        CompletableFuture<Void> asyncCommitResult = new CompletableFuture<>();
+        commitRequest.future.whenComplete((committedOffsets, error) -> {
+            if (error != null) {
+                asyncCommitResult.completeExceptionally(commitAsyncExceptionForError(error));
+            } else {
+                asyncCommitResult.complete(null);
+            }
+        });
+        return asyncCommitResult;
     }
 
     /**
-     * Handles {@link org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsApplicationEvent}. It creates an
-     * {@link OffsetFetchRequestState} and enqueue it to send later.
+     * Commit offsets, retrying on expected retriable errors while the retry timeout hasn't expired.
+     *
+     * @param offsets               Offsets to commit
+     * @param retryExpirationTimeMs Time until which the request will be retried if it fails with
+     *                              an expected retriable error.
+     * @return Future that will complete when a successful response
      */
-    public CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> addOffsetFetchRequest(
+    public CompletableFuture<Void> commitSync(final Map<TopicPartition, OffsetAndMetadata> offsets,
+                                              final long retryExpirationTimeMs) {
+        CompletableFuture<Void> result = new CompletableFuture<>();
+        OffsetCommitRequestState requestState = createOffsetCommitRequest(
+            offsets,
+            Optional.of(retryExpirationTimeMs));
+        commitSyncWithRetries(requestState, result);
+        return result;
+    }
+
+    private OffsetCommitRequestState createOffsetCommitRequest(final Map<TopicPartition, OffsetAndMetadata> offsets,
+                                                               final Optional<Long> expirationTimeMs) {
+        return jitter.isPresent() ?
+            new OffsetCommitRequestState(
+                offsets,
+                groupId,
+                groupInstanceId,
+                expirationTimeMs,
+                retryBackoffMs,
+                retryBackoffMaxMs,
+                jitter.getAsDouble(),
+                memberInfo) :
+            new OffsetCommitRequestState(
+                offsets,
+                groupId,
+                groupInstanceId,
+                expirationTimeMs,
+                retryBackoffMs,
+                retryBackoffMaxMs,
+                memberInfo);
+    }
+
+    private void commitSyncWithRetries(OffsetCommitRequestState requestAttempt,
+                                       CompletableFuture<Void> result) {
+        pendingRequests.addOffsetCommitRequest(requestAttempt);
+
+        // Retry the same commit request while it fails with RetriableException and the retry
+        // timeout hasn't expired.
+        requestAttempt.future.whenComplete((res, error) -> {
+            if (error == null) {
+                result.complete(null);
+            } else {
+                if (error instanceof RetriableException) {
+                    if (error instanceof TimeoutException && requestAttempt.isExpired) {
+                        log.info("OffsetCommit timeout expired so it won't be retried anymore");
+                        result.completeExceptionally(error);
+                    } else {
+                        requestAttempt.resetFuture();
+                        commitSyncWithRetries(requestAttempt, result);
+                    }
+                } else {
+                    result.completeExceptionally(commitSyncExceptionForError(error));
+                }
+            }
+        });
+    }
+
+    private Throwable commitSyncExceptionForError(Throwable error) {
+        if (error instanceof StaleMemberEpochException) {
+            return new CommitFailedException("OffsetCommit failed with stale member epoch."
+                + Errors.STALE_MEMBER_EPOCH.message());
+        }
+        return error;
+    }
+
+    private Throwable commitAsyncExceptionForError(Throwable error) {
+        if (error instanceof RetriableException) {
+            return new RetriableCommitFailedException(error.getMessage());
+        }
+        return error;
+    }
+
+    /**
+     * Enqueue a request to fetch committed offsets, that will be sent on the next call to {@link #poll(long)}.
+     *
+     * @param partitions       Partitions to fetch offsets for.
+     * @param expirationTimeMs Time until which the request should be retried if it fails
+     *                         with expected retriable errors.
+     * @return Future that will complete when a successful response is received, or the request
+     * fails and cannot be retried. Note that the request is retried whenever it fails with
+     * retriable expected error and the retry time hasn't expired.
+     */
+    public CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> fetchOffsets(
         final Set<TopicPartition> partitions,
         final long expirationTimeMs) {
-        return pendingRequests.addOffsetFetchRequest(partitions, expirationTimeMs);
+        if (partitions.isEmpty()) {
+            return CompletableFuture.completedFuture(Collections.emptyMap());
+        }
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = new CompletableFuture<>();
+        OffsetFetchRequestState request = createOffsetFetchRequest(partitions, expirationTimeMs);
+        fetchOffsetsWithRetries(request, result);
+        return result;
+    }
+
+    private OffsetFetchRequestState createOffsetFetchRequest(final Set<TopicPartition> partitions,
+                                                             final long expirationTimeMs) {
+        return jitter.isPresent() ?
+            new OffsetFetchRequestState(
+                partitions,
+                retryBackoffMs,
+                retryBackoffMaxMs,
+                expirationTimeMs,
+                jitter.getAsDouble(),
+                memberInfo) :
+            new OffsetFetchRequestState(
+                partitions,
+                retryBackoffMs,
+                retryBackoffMaxMs,
+                expirationTimeMs,
+                memberInfo);
+    }
+
+    private void fetchOffsetsWithRetries(final OffsetFetchRequestState fetchRequest,
+                                         final CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result) {
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> currentResult = pendingRequests.addOffsetFetchRequest(fetchRequest);
+
+        // Retry the same fetch request while it fails with RetriableException and the retry timeout hasn't expired.
+        currentResult.whenComplete((res, error) -> {
+            boolean inflightRemoved = pendingRequests.inflightOffsetFetches.remove(fetchRequest);
+            if (!inflightRemoved) {
+                log.warn("A duplicated, inflight, request was identified, but unable to find it in the " +
+                    "outbound buffer:" + fetchRequest);
+            }
+            if (error == null) {
+                result.complete(res);
+            } else {
+                if (error instanceof RetriableException || isStaleEpochErrorAndValidEpochAvailable(error)) {
+                    if (error instanceof TimeoutException && fetchRequest.isExpired) {
+                        result.completeExceptionally(error);
+                    } else {
+                        fetchRequest.resetFuture();
+                        fetchOffsetsWithRetries(fetchRequest, result);
+                    }
+                } else
+                    result.completeExceptionally(error);
+            }
+        });
+    }
+
+    private boolean isStaleEpochErrorAndValidEpochAvailable(Throwable error) {
+        return error instanceof StaleMemberEpochException && memberInfo.memberEpoch.isPresent();
     }
 
     public void updateAutoCommitTimer(final long currentTimeMs) {
@@ -393,18 +581,15 @@ public NetworkClientDelegate.PollResult drainPendingOffsetCommitRequests() {
     }
 
     private class OffsetCommitRequestState extends RetriableRequestState {
-        private final Map<TopicPartition, OffsetAndMetadata> offsets;
+        private Map<TopicPartition, OffsetAndMetadata> offsets;
         private final String groupId;
         private final Optional<String> groupInstanceId;
 
-        private final CompletableFuture<Void> future;
-
         /**
-         * Time until which the request should be retried if it fails with retriable
-         * errors. If not present, the request is triggered without waiting for a response or
-         * retrying.
+         * Future containing the offsets that were committed. It completes when a response is
+         * received for the commit request.
          */
-        private final Optional<Long> expirationTimeMs;
+        private CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future;
 
         OffsetCommitRequestState(final Map<TopicPartition, OffsetAndMetadata> offsets,
                                  final String groupId,
@@ -412,15 +597,13 @@ private class OffsetCommitRequestState extends RetriableRequestState {
                                  final Optional<Long> expirationTimeMs,
                                  final long retryBackoffMs,
                                  final long retryBackoffMaxMs,
-                                 final MemberInfo memberInfo,
-                                 final boolean retryOnStaleEpoch) {
+                                 final MemberInfo memberInfo) {
             super(logContext, CommitRequestManager.class.getSimpleName(), retryBackoffMs,
-                retryBackoffMaxMs, memberInfo, retryOnStaleEpoch);
+                retryBackoffMaxMs, memberInfo, expirationTimeMs);
             this.offsets = offsets;
             this.groupId = groupId;
             this.groupInstanceId = groupInstanceId;
             this.future = new CompletableFuture<>();
-            this.expirationTimeMs = expirationTimeMs;
         }
 
         // Visible for testing
@@ -431,15 +614,13 @@ private class OffsetCommitRequestState extends RetriableRequestState {
                                  final long retryBackoffMs,
                                  final long retryBackoffMaxMs,
                                  final double jitter,
-                                 final MemberInfo memberInfo,
-                                 final boolean retryOnStaleEpoch) {
+                                 final MemberInfo memberInfo) {
             super(logContext, CommitRequestManager.class.getSimpleName(), retryBackoffMs, 2,
-                retryBackoffMaxMs, jitter, memberInfo, retryOnStaleEpoch);
+                retryBackoffMaxMs, jitter, memberInfo, expirationTimeMs);
             this.offsets = offsets;
             this.groupId = groupId;
             this.groupInstanceId = groupInstanceId;
             this.future = new CompletableFuture<>();
-            this.expirationTimeMs = expirationTimeMs;
         }
 
         public NetworkClientDelegate.UnsentRequest toUnsentRequest() {
@@ -494,14 +675,16 @@ public void onResponse(final ClientResponse response) {
             for (OffsetCommitResponseData.OffsetCommitResponseTopic topic : commitResponse.data().topics()) {
                 for (OffsetCommitResponseData.OffsetCommitResponsePartition partition : topic.partitions()) {
                     TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex());
-                    OffsetAndMetadata offsetAndMetadata = offsets.get(tp);
-                    long offset = offsetAndMetadata.offset();
+
                     Errors error = Errors.forCode(partition.errorCode());
                     if (error == Errors.NONE) {
+                        OffsetAndMetadata offsetAndMetadata = offsets.get(tp);
+                        long offset = offsetAndMetadata.offset();
                         log.debug("OffsetCommit completed successfully for offset {} partition {}", offset, tp);
                         continue;
                     }
 
+                    onFailedAttempt(currentTimeMs);
                     if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                         future.completeExceptionally(GroupAuthorizationException.forGroupId(groupId));
                         return;
@@ -509,11 +692,11 @@ public void onResponse(final ClientResponse response) {
                         error == Errors.NOT_COORDINATOR ||
                         error == Errors.REQUEST_TIMED_OUT) {
                         coordinatorRequestManager.markCoordinatorUnknown(error.message(), currentTimeMs);
-                        maybeRetry(currentTimeMs, error.exception());
+                        future.completeExceptionally(error.exception());
                         return;
                     } else if (error == Errors.FENCED_INSTANCE_ID) {
                         String fencedError = "OffsetCommit failed due to group instance id fenced: " + groupInstanceId;
-                        log.error(fencedError, error.message());
+                        log.error(fencedError);
                         future.completeExceptionally(new CommitFailedException(fencedError));
                         return;
                     } else if (error == Errors.OFFSET_METADATA_TOO_LARGE ||
@@ -523,21 +706,15 @@ public void onResponse(final ClientResponse response) {
                     } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS ||
                         error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
                         // just retry
-                        maybeRetry(currentTimeMs, error.exception());
+                        future.completeExceptionally(error.exception());
                         return;
                     } else if (error == Errors.UNKNOWN_MEMBER_ID) {
-                        log.error("OffsetCommit failed with {} on partition {} for offset {}",
-                            error, tp, offset);
+                        log.error("OffsetCommit failed with {}", error);
                         future.completeExceptionally(new CommitFailedException("OffsetCommit " +
                             "failed with unknown member ID. " + error.message()));
                         return;
                     } else if (error == Errors.STALE_MEMBER_EPOCH) {
-                        if (maybeRetryWithNewMemberEpoch(currentTimeMs, error)) {
-                            log.debug("OffsetCommit failed with {} and will be retried with the " +
-                                "latest member ID and epoch.", error);
-                            return;
-                        }
-                        future.completeExceptionally(commitExceptionForStaleMemberEpoch());
+                        future.completeExceptionally(error.exception());
                         return;
                     } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
                         // Collect all unauthorized topics before failing
@@ -559,31 +736,6 @@ public void onResponse(final ClientResponse response) {
             }
         }
 
-        /**
-         * Enqueue the request to be retried with exponential backoff, if the request allows
-         * retries and the timer has not expired. Complete the request future exceptionally if
-         * the request won't be retried.
-         */
-        @Override
-        void maybeRetry(long currentTimeMs, Throwable throwable) {
-            if (!allowsRetries()) {
-                // Fail requests that do not allow retries (async requests), making sure to
-                // propagate a RetriableCommitException if the failure is retriable.
-                future.completeExceptionally(commitExceptionForRetriableError(throwable));
-                return;
-            }
-            if (isExpired(currentTimeMs)) {
-                // Fail requests that allowed retries (sync requests), but expired.
-                future.completeExceptionally(throwable);
-                return;
-            }
-
-            // Enqueue request to be retried with backoff. Note that this maintains the same
-            // timer of the initial request, so all the retries are time-bounded.
-            onFailedAttempt(currentTimeMs);
-            pendingRequests.addOffsetCommitRequest(this);
-        }
-
         @Override
         String requestDescription() {
             return "OffsetCommit request for offsets " + offsets;
@@ -594,49 +746,15 @@ CompletableFuture<?> future() {
             return future;
         }
 
-        private boolean isExpired(final long currentTimeMs) {
-            return expirationTimeMs.isPresent() && expirationTimeMs.get() <= currentTimeMs;
+        void resetFuture() {
+            future = new CompletableFuture<>();
         }
 
-        /**
-         * @return True if the requests allows to be retried (sync requests that provide an
-         * expiration time to bound the retries). False if the request does not allow to be
-         * retried on RetriableErrors (async requests that does not provide an expiration time
-         * for retries)
-         */
-        private boolean allowsRetries() {
-            return expirationTimeMs.isPresent();
-        }
-
-        /**
-         * Complete the request future with a TimeoutException if the request expired. No action
-         * taken if the request is still active.
-         *
-         * @return True if the request expired.
-         */
-        private boolean maybeExpire(final long currentTimeMs) {
-            if (isExpired(currentTimeMs)) {
-                future.completeExceptionally(new TimeoutException("OffsetCommit could not complete " +
-                    "before timeout expired."));
-                return true;
+        @Override
+        void removeRequest() {
+            if (!unsentOffsetCommitRequests().remove(this)) {
+                log.warn("OffsetCommit request to remove not found in the outbound buffer: {}", this);
             }
-            return false;
-        }
-
-        /**
-         * @return A RetriableCommitFailedException for async commit requests if the original
-         * Exception was a RetriableException. Return the original one in any other case.
-         */
-        private Throwable commitExceptionForRetriableError(Throwable throwable) {
-            if (!allowsRetries() && throwable instanceof RetriableException)
-                return new RetriableCommitFailedException(throwable);
-            return throwable;
-        }
-
-        private Throwable commitExceptionForStaleMemberEpoch() {
-            if (retryOnStaleEpoch)
-                return new RetriableCommitFailedException(Errors.STALE_MEMBER_EPOCH.exception());
-            return new CommitFailedException("OffsetCommit failed with stale member epoch." + Errors.STALE_MEMBER_EPOCH.message());
         }
     }
 
@@ -652,45 +770,34 @@ abstract class RetriableRequestState extends RequestState {
         final MemberInfo memberInfo;
 
         /**
-         * True if the request should be retried if it fails with {@link Errors#STALE_MEMBER_EPOCH}.
+         * Time until which the request should be retried if it fails with retriable
+         * errors. If not present, the request is triggered without waiting for a response or
+         * retrying.
          */
-        boolean retryOnStaleEpoch;
+        private final Optional<Long> expirationTimeMs;
+
+        /**
+         * True if the request expiration time has been reached. This is set when validating the
+         * request expiration on {@link #poll(long)} before sending it. It is used to know if a
+         * request should be retried on TimeoutException.
+         */
+        boolean isExpired;
 
         RetriableRequestState(LogContext logContext, String owner, long retryBackoffMs,
-                              long retryBackoffMaxMs, MemberInfo memberInfo, boolean retryOnStaleEpoch) {
+                              long retryBackoffMaxMs, MemberInfo memberInfo, Optional<Long> expirationTimeMs) {
             super(logContext, owner, retryBackoffMs, retryBackoffMaxMs);
             this.memberInfo = memberInfo;
-            this.retryOnStaleEpoch = retryOnStaleEpoch;
+            this.expirationTimeMs = expirationTimeMs;
         }
 
         // Visible for testing
         RetriableRequestState(LogContext logContext, String owner, long retryBackoffMs, int retryBackoffExpBase,
-                              long retryBackoffMaxMs, double jitter, MemberInfo memberInfo,
-                              boolean retryOnStaleEpoch) {
+                              long retryBackoffMaxMs, double jitter, MemberInfo memberInfo, Optional<Long> expirationTimeMs) {
             super(logContext, owner, retryBackoffMs, retryBackoffExpBase, retryBackoffMaxMs, jitter);
             this.memberInfo = memberInfo;
-            this.retryOnStaleEpoch = retryOnStaleEpoch;
-        }
-
-        /**
-         * Retry with backoff if the request failed with {@link Errors#STALE_MEMBER_EPOCH} and
-         * the member has valid epoch.
-         *
-         * @return True if the request has been enqueued to be retried with the latest member ID
-         * and epoch.
-         */
-        boolean maybeRetryWithNewMemberEpoch(long currentTimeMs, Errors responseError) {
-            if (retryOnStaleEpoch && memberInfo.memberEpoch.isPresent()) {
-                // Request failed with invalid epoch, but the member has a valid one, so
-                // retry the request with the latest ID/epoch.
-                maybeRetry(currentTimeMs, responseError.exception());
-                return true;
-            }
-            return false;
+            this.expirationTimeMs = expirationTimeMs;
         }
 
-        abstract void maybeRetry(long currentTimeMs, Throwable throwable);
-
         /**
          * @return String containing the request name and arguments, to be used for logging
          * purposes.
@@ -702,6 +809,19 @@ boolean maybeRetryWithNewMemberEpoch(long currentTimeMs, Errors responseError) {
          */
         abstract CompletableFuture<?> future();
 
+        /**
+         * Complete the request future with a TimeoutException if the request timeout has been
+         * reached, based on the provided current time.
+         */
+        void maybeExpire(long currentTimeMs) {
+            if (retryTimeoutExpired(currentTimeMs)) {
+                removeRequest();
+                isExpired = true;
+                future().completeExceptionally(new TimeoutException(requestDescription() +
+                    " could not complete before timeout expired."));
+            }
+        }
+
         /**
          * Build request with the given builder, including response handling logic.
          */
@@ -719,26 +839,29 @@ NetworkClientDelegate.UnsentRequest buildRequestWithResponseHandling(final Abstr
 
         private void handleClientResponse(final ClientResponse response,
                                           final Throwable error,
-                                          final long currentTimeMs) {
+                                          final long requestCompletionTimeMs) {
             try {
                 if (error == null) {
                     onResponse(response);
                 } else {
                     log.debug("{} completed with error", requestDescription(), error);
-                    handleCoordinatorDisconnect(error, currentTimeMs);
-                    if (error instanceof RetriableException) {
-                        maybeRetry(currentTimeMs, error);
-                    } else {
-                        future().completeExceptionally(error);
-                    }
+                    onFailedAttempt(requestCompletionTimeMs);
+                    handleCoordinatorDisconnect(error, requestCompletionTimeMs);
+                    future().completeExceptionally(error);
                 }
             } catch (Throwable t) {
-                log.error("Unexpected error handling response for ", requestDescription(), t);
+                log.error("Unexpected error handling response for {}", requestDescription(), t);
                 future().completeExceptionally(t);
             }
         }
 
         abstract void onResponse(final ClientResponse response);
+
+        boolean retryTimeoutExpired(long currentTimeMs) {
+            return expirationTimeMs.isPresent() && expirationTimeMs.get() <= currentTimeMs;
+        }
+
+        abstract void removeRequest();
     }
 
     class OffsetFetchRequestState extends RetriableRequestState {
@@ -748,12 +871,11 @@ class OffsetFetchRequestState extends RetriableRequestState {
          */
         public final Set<TopicPartition> requestedPartitions;
 
-        private final CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future;
-
         /**
-         * Time until which the request should be retried if it fails with retriable errors.
+         * Future with the result of the request. This can be reset using {@link #resetFuture()}
+         * to get a new result when the request is retried.
          */
-        private final long expirationTimeMs;
+        private CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future;
 
         public OffsetFetchRequestState(final Set<TopicPartition> partitions,
                                        final long retryBackoffMs,
@@ -761,10 +883,9 @@ public OffsetFetchRequestState(final Set<TopicPartition> partitions,
                                        final long expirationTimeMs,
                                        final MemberInfo memberInfo) {
             super(logContext, CommitRequestManager.class.getSimpleName(), retryBackoffMs,
-                retryBackoffMaxMs, memberInfo, true);
+                retryBackoffMaxMs, memberInfo, Optional.of(expirationTimeMs));
             this.requestedPartitions = partitions;
             this.future = new CompletableFuture<>();
-            this.expirationTimeMs = expirationTimeMs;
         }
 
         public OffsetFetchRequestState(final Set<TopicPartition> partitions,
@@ -774,10 +895,9 @@ public OffsetFetchRequestState(final Set<TopicPartition> partitions,
                                        final double jitter,
                                        final MemberInfo memberInfo) {
             super(logContext, CommitRequestManager.class.getSimpleName(), retryBackoffMs, 2,
-                retryBackoffMaxMs, jitter, memberInfo, true);
+                retryBackoffMaxMs, jitter, memberInfo, Optional.of(expirationTimeMs));
             this.requestedPartitions = partitions;
             this.future = new CompletableFuture<>();
-            this.expirationTimeMs = expirationTimeMs;
         }
 
         public boolean sameRequest(final OffsetFetchRequestState request) {
@@ -829,28 +949,22 @@ void onResponse(final ClientResponse response) {
         private void onFailure(final long currentTimeMs,
                                final Errors responseError) {
             log.debug("Offset fetch failed: {}", responseError.message());
+            onFailedAttempt(currentTimeMs);
             if (responseError == COORDINATOR_LOAD_IN_PROGRESS) {
-                maybeRetry(currentTimeMs, responseError.exception());
+                future.completeExceptionally(responseError.exception());
             } else if (responseError == Errors.UNKNOWN_MEMBER_ID) {
                 log.error("OffsetFetch failed with {} because the member is not part of the group" +
                     " anymore.", responseError);
                 future.completeExceptionally(responseError.exception());
             } else if (responseError == Errors.STALE_MEMBER_EPOCH) {
-                if (maybeRetryWithNewMemberEpoch(currentTimeMs, responseError)) {
-                    log.debug("OffsetFetch failed with {} but the consumer is still part" +
-                        " of the group, so the request will be retried with the latest " +
-                        "member ID and epoch.", responseError);
-                    return;
-                }
                 log.error("OffsetFetch failed with {} and the consumer is not part " +
                     "of the group anymore (it probably left the group, got fenced" +
                     " or failed). The request cannot be retried and will fail.", responseError);
                 future.completeExceptionally(responseError.exception());
-            } else if (responseError == Errors.NOT_COORDINATOR
-                || responseError == Errors.COORDINATOR_NOT_AVAILABLE) {
+            } else if (responseError == Errors.NOT_COORDINATOR || responseError == Errors.COORDINATOR_NOT_AVAILABLE) {
                 // Re-discover the coordinator and retry
                 coordinatorRequestManager.markCoordinatorUnknown("error response " + responseError.name(), currentTimeMs);
-                maybeRetry(currentTimeMs, responseError.exception());
+                future.completeExceptionally(responseError.exception());
             } else if (responseError == Errors.GROUP_AUTHORIZATION_FAILED) {
                 future.completeExceptionally(GroupAuthorizationException.forGroupId(groupId));
             } else {
@@ -860,21 +974,6 @@ private void onFailure(final long currentTimeMs,
             }
         }
 
-        /**
-         * Enqueue the request to be retried with exponential backoff if the time has not expired.
-         * This will fail the request future if the request is not retried.
-         */
-        @Override
-        void maybeRetry(long currentTimeMs, Throwable throwable) {
-            if (isExpired(currentTimeMs)) {
-                future.completeExceptionally(throwable);
-                return;
-            }
-            onFailedAttempt(currentTimeMs);
-            pendingRequests.inflightOffsetFetches.remove(this);
-            pendingRequests.addOffsetFetchRequest(this);
-        }
-
         @Override
         String requestDescription() {
             return "OffsetFetch request for partitions " + requestedPartitions;
@@ -885,25 +984,23 @@ CompletableFuture<?> future() {
             return future;
         }
 
-        private boolean isExpired(final long currentTimeMs) {
-            return expirationTimeMs <= currentTimeMs;
+        void resetFuture() {
+            future = new CompletableFuture<>();
         }
 
-        /**
-         * Complete the request future with a TimeoutException if the request expired. No action
-         * taken if the request is still active.
-         *
-         * @return True if the request expired.
-         */
-        private boolean maybeExpire(final long currentTimeMs) {
-            if (isExpired(currentTimeMs)) {
-                future.completeExceptionally(new TimeoutException("OffsetFetch request could not " +
-                    "complete before timeout expired."));
-                return true;
+        @Override
+        void removeRequest() {
+            if (!unsentOffsetFetchRequests().remove(this)) {
+                log.warn("OffsetFetch request to remove not found in the outbound buffer: {}", this);
             }
-            return false;
         }
 
+        /**
+         * Handle OffsetFetch response that has no group level errors. This will look for
+         * partition level errors and fail the future accordingly, also recording a failed request
+         * attempt. If no partition level errors are found, this will complete the future with the
+         * offsets contained in the response, and record a successful request attempt.
+         */
         private void onSuccess(final long currentTimeMs,
                                final OffsetFetchResponse response) {
             Set<String> unauthorizedTopics = null;
@@ -915,6 +1012,7 @@ private void onSuccess(final long currentTimeMs,
                 TopicPartition tp = entry.getKey();
                 OffsetFetchResponse.PartitionData partitionData = entry.getValue();
                 if (partitionData.hasError()) {
+                    onFailedAttempt(currentTimeMs);
                     Errors error = partitionData.error;
                     log.debug("Failed to fetch offset for partition {}: {}", tp, error.message());
 
@@ -954,8 +1052,10 @@ private void onSuccess(final long currentTimeMs,
                         ", this could be either " +
                         "transactional offsets waiting for completion, or " +
                         "normal offsets waiting for replication after appending to local log", unstableTxnOffsetTopicPartitions);
-                maybeRetry(currentTimeMs, new UnstableOffsetCommitException("There are unstable offsets for the requested topic partitions"));
+                future.completeExceptionally(new UnstableOffsetCommitException("There are " +
+                    "unstable offsets for the requested topic partitions"));
             } else {
+                onSuccessfulAttempt(currentTimeMs);
                 future.complete(offsets);
             }
         }
@@ -1003,51 +1103,16 @@ boolean hasUnsentRequests() {
             return !unsentOffsetCommits.isEmpty() || !unsentOffsetFetches.isEmpty();
         }
 
-        OffsetCommitRequestState addOffsetCommitRequest(
-            final Map<TopicPartition, OffsetAndMetadata> offsets,
-            final Optional<Long> expirationTimeMs,
-            final boolean retryOnStaleEpoch) {
-            // TODO: Dedupe committing the same offsets to the same partitions
-            OffsetCommitRequestState requestState = createOffsetCommitRequest(
-                offsets,
-                jitter,
-                expirationTimeMs,
-                retryOnStaleEpoch);
-            return addOffsetCommitRequest(requestState);
-        }
-
+        /**
+         * Add a commit request to the queue, so that it's sent out on the next call to
+         * {@link #poll(long)}. This is used from all commits (sync, async, auto-commit).
+         */
         OffsetCommitRequestState addOffsetCommitRequest(OffsetCommitRequestState request) {
             log.debug("Enqueuing OffsetCommit request for offsets: {}", request.offsets);
             unsentOffsetCommits.add(request);
             return request;
         }
 
-        OffsetCommitRequestState createOffsetCommitRequest(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                                           final OptionalDouble jitter,
-                                                           final Optional<Long> expirationTimeMs,
-                                                           final boolean retryOnStaleEpoch) {
-            return jitter.isPresent() ?
-                new OffsetCommitRequestState(
-                    offsets,
-                    groupId,
-                    groupInstanceId,
-                    expirationTimeMs,
-                    retryBackoffMs,
-                    retryBackoffMaxMs,
-                    jitter.getAsDouble(),
-                    memberInfo,
-                    retryOnStaleEpoch) :
-                new OffsetCommitRequestState(
-                    offsets,
-                    groupId,
-                    groupInstanceId,
-                    expirationTimeMs,
-                    retryBackoffMs,
-                    retryBackoffMaxMs,
-                    memberInfo,
-                    retryOnStaleEpoch);
-        }
-
         /**
          * <p>Adding an offset fetch request to the outgoing buffer.  If the same request was made, we chain the future
          * to the existing one.
@@ -1065,37 +1130,11 @@ private CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> addOffsetFetch
                 log.info("Duplicated OffsetFetchRequest: " + request.requestedPartitions);
                 dupe.orElseGet(() -> inflight.get()).chainFuture(request.future);
             } else {
-                // remove the request from the outbound buffer: inflightOffsetFetches
-                request.future.whenComplete((r, t) -> {
-                    if (!inflightOffsetFetches.remove(request)) {
-                        log.warn("A duplicated, inflight, request was identified, but unable to find it in the " +
-                                "outbound buffer:" + request);
-                    }
-                });
                 this.unsentOffsetFetches.add(request);
             }
             return request.future;
         }
 
-        private CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> addOffsetFetchRequest(final Set<TopicPartition> partitions,
-                                                                                                final long expirationTimeMs) {
-            OffsetFetchRequestState request = jitter.isPresent() ?
-                    new OffsetFetchRequestState(
-                            partitions,
-                            retryBackoffMs,
-                            retryBackoffMaxMs,
-                            expirationTimeMs,
-                            jitter.getAsDouble(),
-                            memberInfo) :
-                    new OffsetFetchRequestState(
-                            partitions,
-                            retryBackoffMs,
-                            retryBackoffMaxMs,
-                            expirationTimeMs,
-                            memberInfo);
-            return addOffsetFetchRequest(request);
-        }
-
         /**
          * Clear {@code unsentOffsetCommits} and moves all the sendable request in {@code
          * unsentOffsetFetches} to the {@code inflightOffsetFetches} to bookkeep all the inflight
@@ -1147,7 +1186,8 @@ List<NetworkClientDelegate.UnsentRequest> drain(final long currentTimeMs) {
          * futures with a TimeoutException.
          */
         private void failAndRemoveExpiredCommitRequests(final long currentTimeMs) {
-            unsentOffsetCommits.removeIf(req -> req.maybeExpire(currentTimeMs));
+            Queue<OffsetCommitRequestState> requestsToPurge = new LinkedList<>(unsentOffsetCommits);
+            requestsToPurge.forEach(req -> req.maybeExpire(currentTimeMs));
         }
 
         /**
@@ -1155,7 +1195,8 @@ private void failAndRemoveExpiredCommitRequests(final long currentTimeMs) {
          * futures with a TimeoutException.
          */
         private void failAndRemoveExpiredFetchRequests(final long currentTimeMs) {
-            unsentOffsetFetches.removeIf(req -> req.maybeExpire(currentTimeMs));
+            Queue<OffsetFetchRequestState> requestsToPurge = new LinkedList<>(unsentOffsetFetches);
+            requestsToPurge.forEach(req -> req.maybeExpire(currentTimeMs));
         }
 
         private void clearAll() {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 6556495e55b34..dd035506d4bea 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -837,12 +837,10 @@ void maybeReconcile() {
         // best effort to commit the offsets in the case where the epoch might have changed while
         // the current reconciliation is in process. Note this is using the rebalance timeout as
         // it is the limit enforced by the broker to complete the reconciliation process.
-        commitResult = commitRequestManager.maybeAutoCommitAllConsumedNow(
-            Optional.of(getExpirationTimeForTimeout(rebalanceTimeoutMs)),
-            true);
+        commitResult = commitRequestManager.maybeAutoCommitSyncNow(getExpirationTimeForTimeout(rebalanceTimeoutMs));
 
         // Execute commit -> onPartitionsRevoked -> onPartitionsAssigned.
-        commitResult.whenComplete((commitReqResult, commitReqError) -> {
+        commitResult.whenComplete((__, commitReqError) -> {
             if (commitReqError != null) {
                 // The call to commit, that includes retry logic for retriable errors, failed to
                 // complete within the time boundaries (fatal error or retriable that did not
@@ -887,9 +885,6 @@ private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitio
             revocationResult.thenCompose(__ -> {
                 boolean memberHasRejoined = memberEpochOnReconciliationStart != memberEpoch;
                 if (state == MemberState.RECONCILING && !memberHasRejoined) {
-                    // Reschedule the auto commit starting from now that member has a new assignment.
-                    commitRequestManager.resetAutoCommitTimer();
-
                     // Apply assignment
                     return assignPartitions(assignedTopicIdPartitions, addedPartitions);
                 } else {
@@ -916,6 +911,9 @@ private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitio
             } else {
                 boolean memberHasRejoined = memberEpochOnReconciliationStart != memberEpoch;
                 if (state == MemberState.RECONCILING && !memberHasRejoined) {
+                    // Reschedule the auto commit starting from now that the member has a new assignment.
+                    commitRequestManager.resetAutoCommitTimer();
+
                     // Make assignment effective on the broker by transitioning to send acknowledge.
                     transitionTo(MemberState.ACKNOWLEDGING);
                 } else {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
index 4396df2785368..ac7ccc56c55f2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
@@ -24,7 +24,7 @@
 public abstract class ApplicationEvent {
 
     public enum Type {
-        COMMIT, POLL, FETCH_COMMITTED_OFFSETS, NEW_TOPICS_METADATA_UPDATE, ASSIGNMENT_CHANGE,
+        COMMIT_ASYNC, COMMIT_SYNC, POLL, FETCH_COMMITTED_OFFSETS, NEW_TOPICS_METADATA_UPDATE, ASSIGNMENT_CHANGE,
         LIST_OFFSETS, RESET_POSITIONS, VALIDATE_POSITIONS, TOPIC_METADATA, SUBSCRIPTION_CHANGE,
         UNSUBSCRIBE, CONSUMER_REBALANCE_LISTENER_CALLBACK_COMPLETED,
         COMMIT_ON_CLOSE, LEAVE_ON_CLOSE
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
index 24a7acf39a9c7..9e48b4de6daad 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
@@ -32,7 +32,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
@@ -69,8 +68,12 @@ public boolean process() {
     @Override
     public void process(ApplicationEvent event) {
         switch (event.type()) {
-            case COMMIT:
-                process((CommitApplicationEvent) event);
+            case COMMIT_ASYNC:
+                process((AsyncCommitApplicationEvent) event);
+                return;
+
+            case COMMIT_SYNC:
+                process((SyncCommitApplicationEvent) event);
                 return;
 
             case POLL:
@@ -139,18 +142,17 @@ private void process(final PollApplicationEvent event) {
         requestManagers.heartbeatRequestManager.ifPresent(hrm -> hrm.resetPollTimer(event.pollTimeMs()));
     }
 
-    private void process(final CommitApplicationEvent event) {
-        if (!requestManagers.commitRequestManager.isPresent()) {
-            // Leaving this error handling here, but it is a bit strange as the commit API should enforce the group.id
-            // upfront, so we should never get to this block.
-            Exception exception = new KafkaException("Unable to commit offset. Most likely because the group.id wasn't set");
-            event.future().completeExceptionally(exception);
-            return;
-        }
+    private void process(final AsyncCommitApplicationEvent event) {
+        CommitRequestManager manager = requestManagers.commitRequestManager.get();
+        CompletableFuture<Void> commitResult = manager.commitAsync(event.offsets());
+        event.chain(commitResult);
+    }
 
+    private void process(final SyncCommitApplicationEvent event) {
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        Optional<Long> expirationTimeMs = event.retryTimeoutMs().map(this::getExpirationTimeForTimeout);
-        event.chain(manager.addOffsetCommitRequest(event.offsets(), expirationTimeMs, false));
+        long expirationTimeoutMs = getExpirationTimeForTimeout(event.retryTimeoutMs());
+        CompletableFuture<Void> commitResult = manager.commitSync(event.offsets(), expirationTimeoutMs);
+        event.chain(commitResult);
     }
 
     private void process(final FetchCommittedOffsetsApplicationEvent event) {
@@ -161,7 +163,7 @@ private void process(final FetchCommittedOffsetsApplicationEvent event) {
         }
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
         long expirationTimeMs = getExpirationTimeForTimeout(event.timeout());
-        event.chain(manager.addOffsetFetchRequest(event.partitions(), expirationTimeMs));
+        event.chain(manager.fetchOffsets(event.partitions(), expirationTimeMs));
     }
 
     private void process(final NewTopicsMetadataUpdateRequestEvent ignored) {
@@ -179,7 +181,7 @@ private void process(final AssignmentChangeApplicationEvent event) {
         }
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
         manager.updateAutoCommitTimer(event.currentTimeMs());
-        manager.maybeAutoCommitAllConsumedAsync();
+        manager.maybeAutoCommitAsync();
     }
 
     private void process(final ListOffsetsApplicationEvent event) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java
new file mode 100644
index 0000000000000..7a939ce3cfd16
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+import java.util.Map;
+
+/**
+ * Event to commit offsets without waiting for a response, so the request won't be retried.
+ */
+public class AsyncCommitApplicationEvent extends CommitApplicationEvent {
+
+    public AsyncCommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets) {
+        super(offsets, Type.COMMIT_ASYNC);
+    }
+
+    @Override
+    public String toString() {
+        return "AsyncCommitApplicationEvent{" +
+            toStringBase() +
+            ", offsets=" + offsets() +
+            '}';
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java
index d2205227c4979..69d969d7b0f46 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java
@@ -21,32 +21,17 @@
 
 import java.util.Collections;
 import java.util.Map;
-import java.util.Optional;
 
-public class CommitApplicationEvent extends CompletableApplicationEvent<Void> {
+public abstract class CommitApplicationEvent extends CompletableApplicationEvent<Void> {
 
     /**
      * Offsets to commit per partition.
      */
     private final Map<TopicPartition, OffsetAndMetadata> offsets;
 
-    /**
-     * Time to wait for a response, retrying on retriable errors. If not present, the request is
-     * triggered without waiting for a response or being retried.
-     */
-    private final Optional<Long> retryTimeoutMs;
-
-    /**
-     * Create new event to commit offsets. If timer is present, the request will be retried on
-     * retriable errors until the timer expires (sync commit offsets request). If the timer is
-     * not present, the request will be sent without waiting for a response of retrying (async
-     * commit offsets request).
-     */
-    public CommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                  final Optional<Long> retryTimeoutMs) {
-        super(Type.COMMIT);
+    public CommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets, Type type) {
+        super(type);
         this.offsets = Collections.unmodifiableMap(offsets);
-        this.retryTimeoutMs = retryTimeoutMs;
 
         for (OffsetAndMetadata offsetAndMetadata : offsets.values()) {
             if (offsetAndMetadata.offset() < 0) {
@@ -59,10 +44,6 @@ public Map<TopicPartition, OffsetAndMetadata> offsets() {
         return offsets;
     }
 
-    public Optional<Long> retryTimeoutMs() {
-        return retryTimeoutMs;
-    }
-
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -80,13 +61,4 @@ public int hashCode() {
         result = 31 * result + offsets.hashCode();
         return result;
     }
-
-    @Override
-    public String toString() {
-        return "CommitApplicationEvent{" +
-                toStringBase() +
-                ", offsets=" + offsets +
-                ", retryTimeout=" + (retryTimeoutMs.map(t -> t + "ms").orElse("none")) +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java
new file mode 100644
index 0000000000000..43dfee6ab18b5
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+import java.util.Map;
+
+/**
+ * Event to commit offsets waiting for a response and retrying on expected retriable errors until
+ * the timer expires.
+ */
+public class SyncCommitApplicationEvent extends CommitApplicationEvent {
+
+    /**
+     * Time to wait for a response, retrying on retriable errors.
+     */
+    private final long retryTimeoutMs;
+
+    public SyncCommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
+                                      final long retryTimeoutMs) {
+        super(offsets, Type.COMMIT_SYNC);
+        this.retryTimeoutMs = retryTimeoutMs;
+    }
+
+    public Long retryTimeoutMs() {
+        return retryTimeoutMs;
+    }
+
+    @Override
+    public String toString() {
+        return "SyncCommitApplicationEvent{" +
+            toStringBase() +
+            ", offsets=" + offsets() +
+            ", retryTimeout=" + retryTimeoutMs + "ms" +
+            '}';
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 575469519823b..2db666e95e935 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -32,8 +32,8 @@
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
-import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
@@ -48,6 +48,7 @@
 import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
 import org.apache.kafka.common.KafkaException;
@@ -252,9 +253,9 @@ public void testCommitAsyncWithNullCallback() {
 
         consumer.commitAsync(offsets, null);
 
-        final ArgumentCaptor<CommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(CommitApplicationEvent.class);
+        final ArgumentCaptor<AsyncCommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitApplicationEvent.class);
         verify(applicationEventHandler).add(commitEventCaptor.capture());
-        final CommitApplicationEvent commitEvent = commitEventCaptor.getValue();
+        final AsyncCommitApplicationEvent commitEvent = commitEventCaptor.getValue();
         assertEquals(offsets, commitEvent.offsets());
         assertDoesNotThrow(() -> commitEvent.future().complete(null));
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, null));
@@ -266,7 +267,7 @@ public void testCommitAsyncUserSuppliedCallbackNoException() {
 
         Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
         offsets.put(new TopicPartition("my-topic", 1), new OffsetAndMetadata(200L));
-        completeCommitApplicationEventSuccessfully();
+        completeCommitAsyncApplicationEventSuccessfully();
 
         MockCommitCallback callback = new MockCommitCallback();
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
@@ -283,7 +284,7 @@ public void testCommitAsyncUserSuppliedCallbackWithException(Exception exception
 
         Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
         offsets.put(new TopicPartition("my-topic", 1), new OffsetAndMetadata(200L));
-        completeCommitApplicationEventExceptionally(exception);
+        completeCommitAsyncApplicationEventExceptionally(exception);
 
         MockCommitCallback callback = new MockCommitCallback();
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
@@ -306,9 +307,9 @@ public void testCommitAsyncWithFencedException() {
 
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
 
-        final ArgumentCaptor<CommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(CommitApplicationEvent.class);
+        final ArgumentCaptor<AsyncCommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitApplicationEvent.class);
         verify(applicationEventHandler).add(commitEventCaptor.capture());
-        final CommitApplicationEvent commitEvent = commitEventCaptor.getValue();
+        final AsyncCommitApplicationEvent commitEvent = commitEventCaptor.getValue();
         commitEvent.future().completeExceptionally(Errors.FENCED_INSTANCE_ID.exception());
 
         assertThrows(Errors.FENCED_INSTANCE_ID.exception().getClass(), () -> consumer.commitAsync());
@@ -433,7 +434,7 @@ public void testCommitInRebalanceCallback() {
         sortedPartitions.add(tp);
         CompletableBackgroundEvent<Void> e = new ConsumerRebalanceListenerCallbackNeededEvent(ON_PARTITIONS_REVOKED, sortedPartitions);
         backgroundEventQueue.add(e);
-        completeCommitApplicationEventSuccessfully();
+        completeCommitSyncApplicationEventSuccessfully();
 
         ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
             @Override
@@ -478,7 +479,7 @@ public void testEnsureCallbackExecutedByApplicationThread() {
         consumer = newConsumer();
         final String currentThread = Thread.currentThread().getName();
         MockCommitCallback callback = new MockCommitCallback();
-        completeCommitApplicationEventSuccessfully();
+        completeCommitAsyncApplicationEventSuccessfully();
 
         assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
         forceCommitCallbackInvocation();
@@ -515,7 +516,7 @@ public void testCommitSyncLeaderEpochUpdate() {
         HashMap<TopicPartition, OffsetAndMetadata> topicPartitionOffsets = new HashMap<>();
         topicPartitionOffsets.put(t0, new OffsetAndMetadata(10L, Optional.of(2), ""));
         topicPartitionOffsets.put(t1, new OffsetAndMetadata(20L, Optional.of(1), ""));
-        completeCommitApplicationEventSuccessfully();
+        completeCommitSyncApplicationEventSuccessfully();
 
         consumer.assign(Arrays.asList(t0, t1));
 
@@ -523,7 +524,7 @@ public void testCommitSyncLeaderEpochUpdate() {
 
         verify(metadata).updateLastSeenEpochIfNewer(t0, 2);
         verify(metadata).updateLastSeenEpochIfNewer(t1, 1);
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(CommitApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
     }
 
     @Test
@@ -557,14 +558,14 @@ public void testCommitAsyncLeaderEpochUpdate() {
 
         verify(metadata).updateLastSeenEpochIfNewer(t0, 2);
         verify(metadata).updateLastSeenEpochIfNewer(t1, 1);
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(CommitApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
     }
 
     @Test
     public void testEnsurePollExecutedCommitAsyncCallbacks() {
         consumer = newConsumer();
         MockCommitCallback callback = new MockCommitCallback();
-        completeCommitApplicationEventSuccessfully();
+        completeCommitAsyncApplicationEventSuccessfully();
         doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
         completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
 
@@ -579,7 +580,7 @@ public void testEnsurePollExecutedCommitAsyncCallbacks() {
     public void testEnsureShutdownExecutedCommitAsyncCallbacks() {
         consumer = newConsumer();
         MockCommitCallback callback = new MockCommitCallback();
-        completeCommitApplicationEventSuccessfully();
+        completeCommitAsyncApplicationEventSuccessfully();
         assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
         assertMockCommitCallbackInvoked(() -> consumer.close(),
             callback,
@@ -670,7 +671,7 @@ public void testAutoCommitSyncEnabled() {
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
         consumer.maybeAutoCommitSync(true, time.timer(100), null);
-        verify(applicationEventHandler).add(any(CommitApplicationEvent.class));
+        verify(applicationEventHandler).add(any(SyncCommitApplicationEvent.class));
     }
 
     @Test
@@ -688,7 +689,7 @@ public void testAutoCommitSyncDisabled() {
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
         consumer.maybeAutoCommitSync(false, time.timer(100), null);
-        verify(applicationEventHandler, never()).add(any(CommitApplicationEvent.class));
+        verify(applicationEventHandler, never()).add(any(SyncCommitApplicationEvent.class));
     }
 
     private void assertMockCommitCallbackInvoked(final Executable task,
@@ -893,7 +894,7 @@ public void testInterceptorAutoCommitOnClose() {
 
         consumer = newConsumer(props);
         assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
-        completeCommitApplicationEventSuccessfully();
+        completeCommitSyncApplicationEventSuccessfully();
 
         consumer.close(Duration.ZERO);
 
@@ -909,7 +910,7 @@ public void testInterceptorCommitSync() {
 
         consumer = newConsumer(props);
         assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
-        completeCommitApplicationEventSuccessfully();
+        completeCommitSyncApplicationEventSuccessfully();
 
         consumer.commitSync(mockTopicPartitionOffset());
 
@@ -925,7 +926,7 @@ public void testNoInterceptorCommitSyncFailed() {
         consumer = newConsumer(props);
         assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
         KafkaException expected = new KafkaException("Test exception");
-        completeCommitApplicationEventExceptionally(expected);
+        completeCommitSyncApplicationEventExceptionally(expected);
 
         KafkaException actual = assertThrows(KafkaException.class, () -> consumer.commitSync(mockTopicPartitionOffset()));
         assertEquals(expected, actual);
@@ -941,7 +942,7 @@ public void testInterceptorCommitAsync() {
         consumer = newConsumer(props);
         assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
 
-        completeCommitApplicationEventSuccessfully();
+        completeCommitAsyncApplicationEventSuccessfully();
         consumer.commitAsync(mockTopicPartitionOffset(), new MockCommitCallback());
         assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
 
@@ -957,7 +958,7 @@ public void testNoInterceptorCommitAsyncFailed() {
 
         consumer = newConsumer(props);
         assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
-        completeCommitApplicationEventExceptionally(new KafkaException("Test exception"));
+        completeCommitAsyncApplicationEventExceptionally(new KafkaException("Test exception"));
 
         consumer.commitAsync(mockTopicPartitionOffset(), new MockCommitCallback());
         assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
@@ -1542,20 +1543,36 @@ private HashMap<TopicPartition, Long> mockTimestampToSearch() {
         return timestampToSearch;
     }
 
-    private void completeCommitApplicationEventExceptionally(Exception ex) {
+    private void completeCommitAsyncApplicationEventExceptionally(Exception ex) {
         doAnswer(invocation -> {
-            CommitApplicationEvent event = invocation.getArgument(0);
+            AsyncCommitApplicationEvent event = invocation.getArgument(0);
             event.future().completeExceptionally(ex);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(CommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
     }
 
-    private void completeCommitApplicationEventSuccessfully() {
+    private void completeCommitSyncApplicationEventExceptionally(Exception ex) {
         doAnswer(invocation -> {
-            CommitApplicationEvent event = invocation.getArgument(0);
+            SyncCommitApplicationEvent event = invocation.getArgument(0);
+            event.future().completeExceptionally(ex);
+            return null;
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
+    }
+
+    private void completeCommitAsyncApplicationEventSuccessfully() {
+        doAnswer(invocation -> {
+            AsyncCommitApplicationEvent event = invocation.getArgument(0);
+            event.future().complete(null);
+            return null;
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
+    }
+
+    private void completeCommitSyncApplicationEventSuccessfully() {
+        doAnswer(invocation -> {
+            SyncCommitApplicationEvent event = invocation.getArgument(0);
             event.future().complete(null);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(CommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
     }
 
     private void completeFetchedCommittedOffsetApplicationEventSuccessfully(final Map<TopicPartition, OffsetAndMetadata> committedOffsets) {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
index 17745fa866093..c27494d69a038 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
@@ -63,6 +63,7 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -78,13 +79,16 @@
 import static org.apache.kafka.test.TestUtils.assertFutureThrows;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -126,7 +130,7 @@ public void testPollSkipIfCoordinatorUnknown() {
 
         Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
         offsets.put(new TopicPartition("t1", 0), new OffsetAndMetadata(0));
-        commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        commitRequestManger.commitAsync(offsets);
         assertPoll(false, 0, commitRequestManger);
     }
 
@@ -137,7 +141,7 @@ public void testPollEnsureManualCommitSent() {
 
         Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
         offsets.put(new TopicPartition("t1", 0), new OffsetAndMetadata(0));
-        commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        commitRequestManger.commitAsync(offsets);
         assertPoll(1, commitRequestManger);
     }
 
@@ -177,13 +181,11 @@ public void testPollEnsureCorrectInflightRequestBufferSize() {
         offsets2.put(new TopicPartition("test", 4), new OffsetAndMetadata(20L));
 
         // Add the requests to the CommitRequestManager and store their futures
-        ArrayList<CompletableFuture<Void>> commitFutures = new ArrayList<>();
-        ArrayList<CompletableFuture<Map<TopicPartition, OffsetAndMetadata>>> fetchFutures = new ArrayList<>();
         long expirationTimeMs  = time.milliseconds() + defaultApiTimeoutMs;
-        commitFutures.add(commitManager.addOffsetCommitRequest(offsets1, Optional.of(expirationTimeMs), false));
-        fetchFutures.add(commitManager.addOffsetFetchRequest(Collections.singleton(new TopicPartition("test", 0)), expirationTimeMs));
-        commitFutures.add(commitManager.addOffsetCommitRequest(offsets2, Optional.of(expirationTimeMs), false));
-        fetchFutures.add(commitManager.addOffsetFetchRequest(Collections.singleton(new TopicPartition("test", 1)), expirationTimeMs));
+        commitManager.commitSync(offsets1, expirationTimeMs);
+        commitManager.fetchOffsets(Collections.singleton(new TopicPartition("test", 0)), expirationTimeMs);
+        commitManager.commitSync(offsets2, expirationTimeMs);
+        commitManager.fetchOffsets(Collections.singleton(new TopicPartition("test", 1)), expirationTimeMs);
 
         // Poll the CommitRequestManager and verify that the inflightOffsetFetches size is correct
         NetworkClientDelegate.PollResult result = commitManager.poll(time.milliseconds());
@@ -195,9 +197,16 @@ public void testPollEnsureCorrectInflightRequestBufferSize() {
         assertFalse(commitManager.pendingRequests.hasUnsentRequests());
         assertEquals(2, commitManager.pendingRequests.inflightOffsetFetches.size());
 
+        // Complete requests with a response
+        result.unsentRequests.forEach(req -> {
+            if (req.requestBuilder() instanceof OffsetFetchRequest.Builder) {
+                req.handler().onComplete(buildOffsetFetchClientResponse(req, Collections.emptySet(), Errors.NONE));
+            } else {
+                req.handler().onComplete(buildOffsetCommitClientResponse(new OffsetCommitResponse(0, new HashMap<>())));
+            }
+        });
+
         // Verify that the inflight offset fetch requests have been removed from the pending request buffer
-        commitFutures.forEach(f -> f.complete(null));
-        fetchFutures.forEach(f -> f.complete(null));
         assertEquals(0, commitManager.pendingRequests.inflightOffsetFetches.size());
     }
 
@@ -208,7 +217,7 @@ public void testPollEnsureEmptyPendingRequestAfterPoll() {
         Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(
             new TopicPartition("topic", 1),
             new OffsetAndMetadata(0));
-        commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        commitRequestManger.commitAsync(offsets);
         assertEquals(1, commitRequestManger.unsentOffsetCommitRequests().size());
         assertEquals(1, commitRequestManger.poll(time.milliseconds()).unsentRequests.size());
         assertTrue(commitRequestManger.unsentOffsetCommitRequests().isEmpty());
@@ -255,24 +264,19 @@ public void testAsyncAutocommitNotRetriedAfterException() {
                 Errors.NONE));
     }
 
-    // This is the case of the sync auto commit sent when the consumer is being closed (sync commit
-    // that should be retried until it succeeds, fails, or timer expires).
+    // This is the case of the sync commit triggered from an API call to commitSync or when the
+    // consumer is being closed. It should be retried until it succeeds, fails, or timer expires.
     @ParameterizedTest
     @MethodSource("offsetCommitExceptionSupplier")
-    public void testAutoCommitSyncRetriedAfterExpectedRetriableException(Errors error) {
-        long commitInterval = retryBackoffMs * 2;
-        CommitRequestManager commitRequestManger = create(true, commitInterval);
-        TopicPartition tp = new TopicPartition("topic", 1);
-        subscriptionState.assignFromUser(Collections.singleton(tp));
-        subscriptionState.seek(tp, 100);
+    public void testCommitSyncRetriedAfterExpectedRetriableException(Errors error) {
+        CommitRequestManager commitRequestManger = create(false, 100);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
-        time.sleep(commitInterval);
-        commitRequestManger.updateAutoCommitTimer(time.milliseconds());
 
-        // Auto-commit all consume sync (ex. triggered when the consumer is closed).
+        Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(
+            new TopicPartition("topic", 1),
+            new OffsetAndMetadata(0));
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        CompletableFuture<Void> commitResult =
-            commitRequestManger.maybeAutoCommitAllConsumedNow(Optional.of(expirationTimeMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(offsets, expirationTimeMs);
         sendAndVerifyOffsetCommitRequestFailedAndMaybeRetried(commitRequestManger, error, commitResult);
 
         // We expect that request should have been retried on this sync commit.
@@ -292,7 +296,7 @@ public void testCommitSyncFailsWithExpectedException(Errors commitError,
 
         // Send sync offset commit that fails and verify it propagates the expected exception.
         Long expirationTimeMs = time.milliseconds() + retryBackoffMs;
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(offsets, Optional.of(expirationTimeMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(offsets, expirationTimeMs);
         completeOffsetCommitRequestWithError(commitRequestManger, commitError);
         assertFutureThrows(commitResult, expectedException);
     }
@@ -309,36 +313,25 @@ private static Stream<Arguments> commitSyncExpectedExceptions() {
     }
 
     @Test
-    public void testOffsetCommitFailsWithCommitFailedExceptionIfUnknownMemberId() {
-        long commitInterval = retryBackoffMs * 2;
-        CommitRequestManager commitRequestManger = create(true, commitInterval);
-        TopicPartition tp = new TopicPartition("topic", 1);
-        subscriptionState.assignFromUser(Collections.singleton(tp));
-        subscriptionState.seek(tp, 100);
+    public void testCommitSyncFailsWithCommitFailedExceptionIfUnknownMemberId() {
+        CommitRequestManager commitRequestManger = create(false, 100);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
-        time.sleep(commitInterval);
-        commitRequestManger.updateAutoCommitTimer(time.milliseconds());
 
-        // Auto-commit all consume sync (ex. triggered when the consumer is closed).
+        Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(
+            new TopicPartition("topic", 1),
+            new OffsetAndMetadata(0));
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        CompletableFuture<Void> commitResult =
-            commitRequestManger.maybeAutoCommitAllConsumedNow(Optional.of(expirationTimeMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(offsets, expirationTimeMs);
 
         completeOffsetCommitRequestWithError(commitRequestManger, Errors.UNKNOWN_MEMBER_ID);
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(0, res.unsentRequests.size());
-        // Commit should fail with CommitFailedException
         assertTrue(commitResult.isDone());
         assertFutureThrows(commitResult, CommitFailedException.class);
     }
 
-    /**
-     * This is the case where a request to commit offsets is performed without retrying on
-     * STALE_MEMBER_EPOCH (ex. commits triggered from the consumer API, auto-commits triggered on
-     * the interval). The expectation is that the request should fail with CommitFailedException.
-     */
     @Test
-    public void testOffsetCommitFailsWithCommitFailedExceptionIfStaleMemberEpochNotRetried() {
+    public void testCommitSyncFailsWithCommitFailedExceptionOnStaleMemberEpoch() {
         CommitRequestManager commitRequestManger = create(true, 100);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
 
@@ -347,8 +340,8 @@ public void testOffsetCommitFailsWithCommitFailedExceptionIfStaleMemberEpochNotR
             new OffsetAndMetadata(0));
 
         // Send commit request expected to be retried on retriable errors
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(
-            offsets, Optional.of(time.milliseconds() + defaultApiTimeoutMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(
+            offsets, time.milliseconds() + defaultApiTimeoutMs);
         completeOffsetCommitRequestWithError(commitRequestManger, Errors.STALE_MEMBER_EPOCH);
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(0, res.unsentRequests.size());
@@ -359,29 +352,35 @@ public void testOffsetCommitFailsWithCommitFailedExceptionIfStaleMemberEpochNotR
     }
 
     /**
-     * This is the case of the async auto commit request triggered on the interval. The
-     * expectation is that the request should fail with RetriableCommitFailedException, so that
-     * the timer is reset and the request is retried on the next interval.
+     * This is the case of the async auto commit request triggered on the interval. The request
+     * internally fails with the fatal stale epoch error, and the expectation is that it just
+     * resets the commit timer to the interval, to attempt again when the interval expires.
      */
     @Test
-    public void testCommitAsyncFailsWithRetriableOnStaleMemberEpoch() {
-        CommitRequestManager commitRequestManger = create(false, 100);
+    public void testAutoCommitAsyncFailsWithStaleMemberEpochContinuesToCommitOnTheInterval() {
+        CommitRequestManager commitRequestManger = create(true, 100);
+        time.sleep(100);
+        commitRequestManger.updateAutoCommitTimer(time.milliseconds());
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
+        TopicPartition t1p = new TopicPartition("topic1", 0);
+        subscriptionState.assignFromUser(singleton(t1p));
+        subscriptionState.seek(t1p, 10);
 
-        Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(
-            new TopicPartition("topic", 1),
-            new OffsetAndMetadata(0));
-
-        // Async commit that won't be retried.
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(
-            offsets, Optional.empty(), true);
+        // Async commit on the interval fails with fatal stale epoch and just resets the timer to
+        // the interval
+        commitRequestManger.maybeAutoCommitAsync();
         completeOffsetCommitRequestWithError(commitRequestManger, Errors.STALE_MEMBER_EPOCH);
+        verify(commitRequestManger).resetAutoCommitTimer();
+
+        // Async commit retried, only when the interval expires
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
-        assertEquals(0, res.unsentRequests.size());
+        assertEquals(0, res.unsentRequests.size(), "No request should be generated until the " +
+            "interval expires");
+        time.sleep(100);
+        commitRequestManger.updateAutoCommitTimer(time.milliseconds());
+        res = commitRequestManger.poll(time.milliseconds());
+        assertEquals(1, res.unsentRequests.size());
 
-        // Commit should fail with RetriableCommitFailedException.
-        assertTrue(commitResult.isDone());
-        assertFutureThrows(commitResult, RetriableCommitFailedException.class);
     }
 
     @Test
@@ -394,8 +393,7 @@ public void testCommitAsyncFailsWithRetriableOnCoordinatorDisconnected() {
             new OffsetAndMetadata(0));
 
         // Async commit that won't be retried.
-        CompletableFuture<Void> commitResult = commitRequestManager.addOffsetCommitRequest(
-            offsets, Optional.empty(), true);
+        CompletableFuture<Void> commitResult = commitRequestManager.commitAsync(offsets);
 
         NetworkClientDelegate.PollResult res = commitRequestManager.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -473,11 +471,10 @@ public void testAutoCommitEmptyOffsetsDoesNotGenerateRequest() {
         CommitRequestManager commitRequestManger = create(true, 100);
         time.sleep(100);
         commitRequestManger.updateAutoCommitTimer(time.milliseconds());
-        CompletableFuture<Void> result = commitRequestManger.maybeAutoCommitAllConsumedAsync();
-
+        // CompletableFuture<Void> result = commitRequestManger.maybeAutoCommitAllConsumedAsync();
+        commitRequestManger.maybeAutoCommitAsync();
         assertTrue(commitRequestManger.pendingRequests.unsentOffsetCommits.isEmpty());
-        assertTrue(result.isDone());
-        assertFalse(result.isCompletedExceptionally());
+        verify(commitRequestManger).resetAutoCommitTimer();
     }
 
     @Test
@@ -489,18 +486,17 @@ public void testAutoCommitEmptyDoesNotLeaveInflightRequestFlagOn() {
         // Auto-commit of empty offsets
         time.sleep(100);
         commitRequestManger.updateAutoCommitTimer(time.milliseconds());
-        CompletableFuture<Void> result = commitRequestManger.maybeAutoCommitAllConsumedAsync();
-        assertTrue(result.isDone());
-        assertFalse(result.isCompletedExceptionally());
+        commitRequestManger.maybeAutoCommitAsync();
 
         // Next auto-commit consumed offsets (not empty). Should generate a request, ensuring
         // that the previous auto-commit of empty did not leave the inflight request flag on
         subscriptionState.seek(t1p, 100);
         time.sleep(100);
         commitRequestManger.updateAutoCommitTimer(time.milliseconds());
-        result = commitRequestManger.maybeAutoCommitAllConsumedAsync();
-        assertFalse(commitRequestManger.pendingRequests.unsentOffsetCommits.isEmpty());
-        assertFalse(result.isDone());
+        commitRequestManger.maybeAutoCommitAsync();
+        assertEquals(1, commitRequestManger.pendingRequests.unsentOffsetCommits.size());
+
+        verify(commitRequestManger, times(2)).resetAutoCommitTimer();
     }
 
     @Test
@@ -534,7 +530,7 @@ public void testOffsetFetchRequestErroredRequests(final Errors error, final bool
         List<CompletableFuture<Map<TopicPartition, OffsetAndMetadata>>> futures = sendAndVerifyDuplicatedOffsetFetchRequests(
             commitRequestManger,
             partitions,
-            5,
+            1,
             error);
         // we only want to make sure to purge the outbound buffer for non-retriables, so retriable will be re-queued.
         if (isRetriable)
@@ -545,6 +541,49 @@ public void testOffsetFetchRequestErroredRequests(final Errors error, final bool
         }
     }
 
+    @Test
+    public void testSuccessfulOffsetFetch() {
+        CommitRequestManager commitManager = create(false, 100);
+        when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
+
+        long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> fetchResult =
+            commitManager.fetchOffsets(Collections.singleton(new TopicPartition("test", 0)),
+                expirationTimeMs);
+
+        // Send fetch request
+        NetworkClientDelegate.PollResult result = commitManager.poll(time.milliseconds());
+        assertEquals(1, result.unsentRequests.size());
+        assertEquals(1, commitManager.pendingRequests.inflightOffsetFetches.size());
+        assertFalse(fetchResult.isDone());
+
+        // Complete request with a response
+        TopicPartition tp = new TopicPartition("topic1", 0);
+        long expectedOffset = 100;
+        NetworkClientDelegate.UnsentRequest req = result.unsentRequests.get(0);
+        Map<TopicPartition, OffsetFetchResponse.PartitionData> topicPartitionData =
+            Collections.singletonMap(
+                tp,
+                new OffsetFetchResponse.PartitionData(expectedOffset, Optional.of(1), "", Errors.NONE));
+        req.handler().onComplete(buildOffsetFetchClientResponse(req, topicPartitionData, Errors.NONE, false));
+
+        // Validate request future completes with the response received
+        assertTrue(fetchResult.isDone());
+        assertFalse(fetchResult.isCompletedExceptionally());
+        Map<TopicPartition, OffsetAndMetadata> offsetsAndMetadata = null;
+        try {
+            offsetsAndMetadata = fetchResult.get();
+        } catch (InterruptedException | ExecutionException e) {
+            fail(e);
+        }
+        assertNotNull(offsetsAndMetadata);
+        assertEquals(1, offsetsAndMetadata.size());
+        assertTrue(offsetsAndMetadata.containsKey(tp));
+        assertEquals(expectedOffset, offsetsAndMetadata.get(tp).offset());
+        assertEquals(0, commitManager.pendingRequests.inflightOffsetFetches.size(), "Inflight " +
+            "request should be removed from the queue when a response is received.");
+    }
+
     @ParameterizedTest
     @MethodSource("offsetFetchRetriableCoordinatorErrors")
     public void testOffsetFetchMarksCoordinatorUnknownOnRetriableCoordinatorErrors(Errors error,
@@ -556,7 +595,7 @@ public void testOffsetFetchMarksCoordinatorUnknownOnRetriableCoordinatorErrors(E
         partitions.add(new TopicPartition("t1", 0));
 
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = commitRequestManager.addOffsetFetchRequest(partitions, expirationTimeMs);
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = commitRequestManager.fetchOffsets(partitions, expirationTimeMs);
 
         completeOffsetFetchRequestWithError(commitRequestManager, partitions, error);
 
@@ -583,7 +622,7 @@ public void testOffsetFetchMarksCoordinatorUnknownOnCoordinatorDisconnectedAndRe
         partitions.add(new TopicPartition("t1", 0));
 
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = commitRequestManger.addOffsetFetchRequest(partitions, expirationTimeMs);
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> result = commitRequestManger.fetchOffsets(partitions, expirationTimeMs);
 
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -595,8 +634,7 @@ public void testOffsetFetchMarksCoordinatorUnknownOnCoordinatorDisconnectedAndRe
         assertFalse(result.isDone());
         assertCoordinatorDisconnect();
 
-        // Request should be retried after backoff expires
-        time.sleep(100);
+        time.sleep(retryBackoffMs);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
         res = commitRequestManger.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -612,8 +650,8 @@ public void testOffsetCommitRequestErroredRequestsNotRetriedForAsyncCommit(final
         Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(new TopicPartition("topic", 1),
             new OffsetAndMetadata(0));
 
-        // Send commit request without expiration (async commit not expected to be retried).
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        // Send async commit (not expected to be retried).
+        CompletableFuture<Void> commitResult = commitRequestManger.commitAsync(offsets);
         completeOffsetCommitRequestWithError(commitRequestManger, error);
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(0, res.unsentRequests.size());
@@ -638,7 +676,7 @@ public void testOffsetCommitSyncTimeoutNotReturnedOnPollAndFails() {
 
         // Send sync offset commit request that fails with retriable error.
         Long expirationTimeMs = time.milliseconds() + retryBackoffMs * 2;
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(offsets, Optional.of(expirationTimeMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(offsets, expirationTimeMs);
         completeOffsetCommitRequestWithError(commitRequestManger, Errors.REQUEST_TIMED_OUT);
 
         // Request retried after backoff, and fails with retriable again. Should not complete yet
@@ -670,7 +708,7 @@ public void testOffsetCommitSyncFailedWithRetriableThrowsTimeoutWhenRetryTimeExp
 
         // Send offset commit request that fails with retriable error.
         long expirationTimeMs = time.milliseconds() + retryBackoffMs * 2;
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(offsets, Optional.of(expirationTimeMs), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitSync(offsets, expirationTimeMs);
         completeOffsetCommitRequestWithError(commitRequestManger, Errors.COORDINATOR_NOT_AVAILABLE);
 
         // Sleep to expire the request timeout. Request should fail on the next poll with a
@@ -696,7 +734,7 @@ public void testOffsetCommitAsyncFailedWithRetriableThrowsRetriableCommitExcepti
 
         // Send async commit request that fails with retriable error (not expected to be retried).
         Errors retriableError = Errors.COORDINATOR_NOT_AVAILABLE;
-        CompletableFuture<Void> commitResult = commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        CompletableFuture<Void> commitResult = commitRequestManger.commitAsync(offsets);
         completeOffsetCommitRequestWithError(commitRequestManger, retriableError);
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(0, res.unsentRequests.size());
@@ -710,7 +748,6 @@ public void testOffsetCommitAsyncFailedWithRetriableThrowsRetriableCommitExcepti
         assertFutureThrows(commitResult, RetriableCommitFailedException.class);
     }
 
-
     @Test
     public void testEnsureBackoffRetryOnOffsetCommitRequestTimeout() {
         CommitRequestManager commitRequestManger = create(true, 100);
@@ -720,7 +757,7 @@ public void testEnsureBackoffRetryOnOffsetCommitRequestTimeout() {
             new OffsetAndMetadata(0));
 
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        commitRequestManger.addOffsetCommitRequest(offsets, Optional.of(expirationTimeMs), false);
+        commitRequestManger.commitSync(offsets, expirationTimeMs);
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
         res.unsentRequests.get(0).handler().onFailure(time.milliseconds(), new TimeoutException());
@@ -802,7 +839,7 @@ public void testSyncOffsetFetchFailsWithStaleEpochAndRetriesWithNewEpoch() {
 
         // Send request that is expected to fail with invalid epoch.
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        commitRequestManager.addOffsetFetchRequest(partitions, expirationTimeMs);
+        commitRequestManager.fetchOffsets(partitions, expirationTimeMs);
 
         // Mock member has new a valid epoch.
         int newEpoch = 8;
@@ -842,7 +879,7 @@ public void testSyncOffsetFetchFailsWithStaleEpochAndNotRetriedIfMemberNotInGrou
         // Send request that is expected to fail with invalid epoch.
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
         CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> requestResult =
-            commitRequestManager.addOffsetFetchRequest(partitions, expirationTimeMs);
+            commitRequestManager.fetchOffsets(partitions, expirationTimeMs);
 
         // Mock member not having a valid epoch anymore (left/failed/fenced).
         commitRequestManager.onMemberEpochUpdated(Optional.empty(), Optional.empty());
@@ -862,40 +899,58 @@ public void testSyncOffsetFetchFailsWithStaleEpochAndNotRetriedIfMemberNotInGrou
     // STALE_MEMBER_EPOCH (ex. triggered from the reconciliation process), fails with invalid
     // member epoch. The expectation is that if the member already has a new epoch, the request
     // should be retried with the new epoch.
-    @Test
-    public void testOffsetCommitFailsWithStaleEpochAndRetriesWithNewEpoch() {
-        CommitRequestManager commitRequestManager = create(true, 100);
+    @ParameterizedTest
+    @MethodSource("offsetCommitExceptionSupplier")
+    public void testAutoCommitSyncBeforeRevocationRetriesOnRetriableAndStaleEpoch(Errors error) {
+        // Enable auto-commit but with very long interval to avoid triggering auto-commits on the
+        // interval and just test the auto-commits triggered before revocation
+        CommitRequestManager commitRequestManager = create(true, Integer.MAX_VALUE);
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(mockedNode));
 
-        Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(new TopicPartition("topic", 1),
-            new OffsetAndMetadata(0));
+        TopicPartition tp = new TopicPartition("topic", 1);
+        subscriptionState.assignFromUser(singleton(tp));
+        subscriptionState.seek(tp, 5);
+        long expirationTimeMs = time.milliseconds() + retryBackoffMs * 2;
 
-        // Send commit request expected to be retried on STALE_MEMBER_EPOCH error while it does not expire.
-        long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
-        commitRequestManager.addOffsetCommitRequest(offsets, Optional.of(expirationTimeMs), true);
+        // Send commit request expected to be retried on STALE_MEMBER_EPOCH error while it does not expire
+        commitRequestManager.maybeAutoCommitSyncNow(expirationTimeMs);
 
-        // Mock member has new a valid epoch.
         int newEpoch = 8;
         String memberId = "member1";
-        commitRequestManager.onMemberEpochUpdated(Optional.of(newEpoch), Optional.of(memberId));
-
-        completeOffsetCommitRequestWithError(commitRequestManager, Errors.STALE_MEMBER_EPOCH);
-
-        // Check that the request that failed was removed from the inflight requests buffer.
-        assertEquals(0, commitRequestManager.pendingRequests.inflightOffsetFetches.size());
-        assertEquals(1, commitRequestManager.pendingRequests.unsentOffsetCommits.size());
-
-        // Request should be retried with backoff.
-        NetworkClientDelegate.PollResult res = commitRequestManager.poll(time.milliseconds());
-        assertEquals(0, res.unsentRequests.size());
-        time.sleep(retryBackoffMs);
-        res = commitRequestManager.poll(time.milliseconds());
-        assertEquals(1, res.unsentRequests.size());
+        if (error == Errors.STALE_MEMBER_EPOCH) {
+            // Mock member has new a valid epoch
+            commitRequestManager.onMemberEpochUpdated(Optional.of(newEpoch), Optional.of(memberId));
+        }
 
-        // The retried request should include the latest member ID and epoch.
-        OffsetCommitRequestData reqData = (OffsetCommitRequestData) res.unsentRequests.get(0).requestBuilder().build().data();
-        assertEquals(newEpoch, reqData.generationIdOrMemberEpoch());
-        assertEquals(memberId, reqData.memberId());
+        completeOffsetCommitRequestWithError(commitRequestManager, error);
+
+        if (error.exception() instanceof RetriableException || error == Errors.STALE_MEMBER_EPOCH) {
+            assertEquals(1, commitRequestManager.pendingRequests.unsentOffsetCommits.size(),
+                "Request to be retried should be added to the outbound queue");
+
+            // Request should be retried with backoff
+            NetworkClientDelegate.PollResult res = commitRequestManager.poll(time.milliseconds());
+            assertEquals(0, res.unsentRequests.size());
+            time.sleep(retryBackoffMs);
+            res = commitRequestManager.poll(time.milliseconds());
+            assertEquals(1, res.unsentRequests.size());
+            if (error == Errors.STALE_MEMBER_EPOCH) {
+                // The retried request should include the latest member ID and epoch
+                OffsetCommitRequestData reqData = (OffsetCommitRequestData) res.unsentRequests.get(0).requestBuilder().build().data();
+                assertEquals(newEpoch, reqData.generationIdOrMemberEpoch());
+                assertEquals(memberId, reqData.memberId());
+            }
+        } else {
+            assertEquals(0, commitRequestManager.pendingRequests.unsentOffsetCommits.size(),
+                "Non-retriable failed request should be removed from the outbound queue");
+
+            // Request should not be retried, even after the backoff expires
+            NetworkClientDelegate.PollResult res = commitRequestManager.poll(time.milliseconds());
+            assertEquals(0, res.unsentRequests.size());
+            time.sleep(retryBackoffMs);
+            res = commitRequestManager.poll(time.milliseconds());
+            assertEquals(0, res.unsentRequests.size());
+        }
     }
 
     @Test
@@ -920,7 +975,7 @@ private void commitOffsetWithAssertedLatency(CommitRequestManager commitRequestM
                 new OffsetAndMetadata(0));
 
         long commitCreationTimeMs = time.milliseconds();
-        commitRequestManager.addOffsetCommitRequest(offsets, Optional.empty(), true);
+        commitRequestManager.commitAsync(offsets);
 
         NetworkClientDelegate.PollResult res = commitRequestManager.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -1029,7 +1084,7 @@ public void testOffsetFetchRequestPartitionDataError(final Errors error, final b
         partitions.add(tp2);
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
         CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future =
-                commitRequestManger.addOffsetFetchRequest(partitions, expirationTimeMs);
+                commitRequestManger.fetchOffsets(partitions, expirationTimeMs);
 
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -1058,7 +1113,7 @@ public void testSignalClose() {
         Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(new TopicPartition("topic", 1),
             new OffsetAndMetadata(0));
 
-        commitRequestManger.addOffsetCommitRequest(offsets, Optional.empty(), false);
+        commitRequestManger.commitAsync(offsets);
         commitRequestManger.signalClose();
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
         assertEquals(1, res.unsentRequests.size());
@@ -1089,7 +1144,7 @@ private List<CompletableFuture<Map<TopicPartition, OffsetAndMetadata>>> sendAndV
         List<CompletableFuture<Map<TopicPartition, OffsetAndMetadata>>> futures = new ArrayList<>();
         long expirationTimeMs = time.milliseconds() + defaultApiTimeoutMs;
         for (int i = 0; i < numRequest; i++) {
-            futures.add(commitRequestManger.addOffsetFetchRequest(partitions, expirationTimeMs));
+            futures.add(commitRequestManger.fetchOffsets(partitions, expirationTimeMs));
         }
 
         NetworkClientDelegate.PollResult res = commitRequestManger.poll(time.milliseconds());
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
index 4ff652ff49318..a491df417de45 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
@@ -21,11 +21,13 @@
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ListOffsetsApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
+import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.TopicMetadataApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
 import org.apache.kafka.common.Node;
@@ -135,7 +137,7 @@ public void testStartupAndTearDown() throws InterruptedException {
 
     @Test
     public void testApplicationEvent() {
-        ApplicationEvent e = new CommitApplicationEvent(new HashMap<>(), Optional.empty());
+        ApplicationEvent e = new PollApplicationEvent(100);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor, times(1)).process(e);
@@ -150,11 +152,19 @@ public void testMetadataUpdateEvent() {
     }
 
     @Test
-    public void testCommitEvent() {
-        ApplicationEvent e = new CommitApplicationEvent(new HashMap<>(), Optional.empty());
+    public void testAsyncCommitEvent() {
+        ApplicationEvent e = new AsyncCommitApplicationEvent(new HashMap<>());
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(CommitApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(AsyncCommitApplicationEvent.class));
+    }
+
+    @Test
+    public void testSyncCommitEvent() {
+        ApplicationEvent e = new SyncCommitApplicationEvent(new HashMap<>(), 100L);
+        applicationEventsQueue.add(e);
+        consumerNetworkThread.runOnce();
+        verify(applicationEventProcessor).process(any(SyncCommitApplicationEvent.class));
     }
 
     @Test
@@ -209,7 +219,7 @@ public void testAssignmentChangeEvent() {
         verify(networkClient, times(1)).poll(anyLong(), anyLong());
         verify(commitRequestManager, times(1)).updateAutoCommitTimer(currentTimeMs);
         // Assignment change should generate an async commit (not retried).
-        verify(commitRequestManager, times(1)).maybeAutoCommitAllConsumedAsync();
+        verify(commitRequestManager, times(1)).maybeAutoCommitAsync();
     }
 
     @Test
@@ -273,8 +283,8 @@ void testEnsureEventsAreCompleted() {
         coordinatorRequestManager.markCoordinatorUnknown("test", time.milliseconds());
         client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, "group-id", node));
         prepareOffsetCommitRequest(new HashMap<>(), Errors.NONE, false);
-        CompletableApplicationEvent<Void> event1 = spy(new CommitApplicationEvent(Collections.emptyMap(), Optional.empty()));
-        ApplicationEvent event2 = new CommitApplicationEvent(Collections.emptyMap(), Optional.empty());
+        CompletableApplicationEvent<Void> event1 = spy(new AsyncCommitApplicationEvent(Collections.emptyMap()));
+        ApplicationEvent event2 = new AsyncCommitApplicationEvent(Collections.emptyMap());
         CompletableFuture<Void> future = new CompletableFuture<>();
         when(event1.future()).thenReturn(future);
         applicationEventsQueue.add(event1);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 2d8ef9397f35e..f5c65d58bfc54 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -69,6 +69,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyCollection;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anySet;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
@@ -1922,6 +1923,7 @@ private void verifyReconciliationTriggeredAndCompleted(MembershipManagerImpl mem
         Map<Uuid, SortedSet<Integer>> assignmentByTopicId = assignmentByTopicId(expectedAssignment);
         assertEquals(assignmentByTopicId, membershipManager.currentAssignment());
 
+        // The auto-commit interval should be reset (only once), when the reconciliation completes
         verify(commitRequestManager).resetAutoCommitTimer();
     }
 
@@ -1947,7 +1949,7 @@ private CompletableFuture<Void> mockRevocationNoCallbacks(boolean withAutoCommit
         if (withAutoCommit) {
             when(commitRequestManager.autoCommitEnabled()).thenReturn(true);
             CompletableFuture<Void> commitResult = new CompletableFuture<>();
-            when(commitRequestManager.maybeAutoCommitAllConsumedNow(any(), anyBoolean())).thenReturn(commitResult);
+            when(commitRequestManager.maybeAutoCommitSyncNow(anyLong())).thenReturn(commitResult);
             return commitResult;
         } else {
             return CompletableFuture.completedFuture(null);

From cc49fc76564b15f2111801faf62341aaf01c4d8c Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <cadonna@apache.org>
Date: Wed, 21 Feb 2024 11:54:46 +0100
Subject: [PATCH 063/258] HOTFIX: Fix compilation error in
 TransactionManagerTest (#15405)

Variable metadataMock was removed by #15323 after the CI build of #15320 was run and before #15320 was merged.

Reviewers: Luke Chen <showuon@gmail.com>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../clients/producer/internals/TransactionManagerTest.java     | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index de02738a34ad6..b099ebfc41932 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -2550,8 +2550,7 @@ public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedExc
         transactionManager.maybeAddPartition(tp0);
         accumulator.beginFlush();
 
-        drainedBatches = accumulator.drain(metadataMock, nodes, Integer.MAX_VALUE,
-            time.milliseconds());
+        drainedBatches = accumulator.drain(metadataCache, nodes, Integer.MAX_VALUE, time.milliseconds());
 
         // We still shouldn't drain batches because the partition call didn't complete yet.
         assertTrue(drainedBatches.containsKey(node1.id()));

From 4c012c5c23b9b5031f4cfbf932ba399afd72e156 Mon Sep 17 00:00:00 2001
From: Anton Liauchuk <anton93lev@gmail.com>
Date: Wed, 21 Feb 2024 14:25:15 +0300
Subject: [PATCH 064/258] KAFKA-16278: Missing license for scala related
 dependencies (#15398)

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 LICENSE-binary | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/LICENSE-binary b/LICENSE-binary
index c3b370da06ea8..61eb5f379668c 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -258,10 +258,15 @@ plexus-utils-3.5.1
 reflections-0.10.2
 reload4j-1.2.25
 rocksdbjni-7.9.2
+scala-collection-compat_2.12-2.10.0
 scala-collection-compat_2.13-2.10.0
+scala-library-2.12.18
 scala-library-2.13.12
+scala-logging_2.12-3.9.4
 scala-logging_2.13-3.9.4
+scala-reflect-2.12.18
 scala-reflect-2.13.12
+scala-java8-compat_2.12-1.0.2
 scala-java8-compat_2.13-1.0.2
 snappy-java-1.1.10.5
 swagger-annotations-2.2.8

From a3528a316f4df969b9995f00af89161a5840a547 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Wed, 21 Feb 2024 08:01:11 -0800
Subject: [PATCH 065/258] MINOR: remove unnecessary logging (#15396)

We already record dropping record via metrics and logging at WARN level
is too noise. This PR removes the unnecessary logging.

Reviewers: Kalpesh Patel <kpatel@confluent.io>, Walker Carlson <wcarlson@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../state/internals/AbstractRocksDBSegmentedBytesStore.java      | 1 -
 1 file changed, 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
index 348d6bb18632a..c0749ce37a757 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStore.java
@@ -264,7 +264,6 @@ public void put(final Bytes key,
         final S segment = segments.getOrCreateSegmentIfLive(segmentId, context, observedStreamTime);
         if (segment == null) {
             expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
-            LOG.warn("Skipping record for expired segment.");
         } else {
             synchronized (position) {
                 StoreQueryUtils.updatePosition(position, stateStoreContext);

From 02ebfc6108c86d8f7c74396f440ae5336e0bcae2 Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <cadonna@apache.org>
Date: Wed, 21 Feb 2024 18:19:38 +0100
Subject: [PATCH 066/258] KAFKA-16194: Do not return records from poll if group
 metadata unknown (#15369)

Due to the asynchronous nature of the async consumer, it might happen that on the application thread the group metadata is not known after the first poll returns records. If the offsets of those records are then send to a transaction with

txnProducer.sendOffsetsToTransaction(offsetsToCommit, groupMetadata);

and then the transaction is committed, the group coordinator will raise an error saying that the member is not known since the member in groupMetadata is still -1 before the metadata is updated.

This commit avoids this error by not returning any records from poll() until the group metadata is updated, i.e., the member ID and the generation ID (a.k.a. member epoch) are known. This check is only done if group management is used.

Additionally, this commit resets the group metadata when the consumer unsubscribes.

Reviewer: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/AsyncKafkaConsumer.java         |  24 ++-
 .../internals/AsyncKafkaConsumerTest.java     | 148 ++++++++++++++++++
 2 files changed, 166 insertions(+), 6 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index d2e4788a1bc18..d4b461b0140d8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -712,14 +712,18 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
                 wakeupTrigger.maybeTriggerWakeup();
 
                 updateAssignmentMetadataIfNeeded(timer);
-                final Fetch<K, V> fetch = pollForFetches(timer);
-                if (!fetch.isEmpty()) {
-                    if (fetch.records().isEmpty()) {
-                        log.trace("Returning empty records from `poll()` "
+                if (isGenerationKnownOrPartitionsUserAssigned()) {
+                    final Fetch<K, V> fetch = pollForFetches(timer);
+                    if (!fetch.isEmpty()) {
+                        if (fetch.records().isEmpty()) {
+                            log.trace("Returning empty records from `poll()` "
                                 + "since the consumer's position has advanced for at least one topic partition");
-                    }
+                        }
 
-                    return interceptors.onConsume(new ConsumerRecords<>(fetch.records()));
+                        return interceptors.onConsume(new ConsumerRecords<>(fetch.records()));
+                    }
+                } else {
+                    timer.update();
                 }
                 // We will wait for retryBackoffMs
             } while (timer.notExpired());
@@ -731,6 +735,13 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
         }
     }
 
+    private boolean isGenerationKnownOrPartitionsUserAssigned() {
+        if (subscriptions.hasAutoAssignedPartitions()) {
+            return groupMetadata.filter(g -> g.generationId() != JoinGroupRequest.UNKNOWN_GENERATION_ID).isPresent();
+        }
+        return true;
+    }
+
     /**
      * Commit offsets returned on the last {@link #poll(Duration) poll()} for all the subscribed list of topics and
      * partitions.
@@ -1463,6 +1474,7 @@ public void unsubscribe() {
                 } catch (TimeoutException e) {
                     log.error("Failed while waiting for the unsubscribe event to complete");
                 }
+                groupMetadata = initializeGroupMetadata(groupMetadata.get().groupId(), Optional.empty());
             }
             subscriptions.unsubscribe();
         } finally {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 2db666e95e935..abd0f45c73f3c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -51,6 +51,7 @@
 import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
+import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.Node;
@@ -99,6 +100,7 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -1131,6 +1133,144 @@ public void testGroupMetadataUpdateSingleCall() {
         assertEquals(expectedGroupMetadata, secondActualGroupMetadataWithoutUpdate);
     }
 
+    @Test
+    public void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsedWithTopics() {
+        testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(() -> {
+            consumer.subscribe(singletonList("topic"));
+        });
+    }
+
+    @Test
+    public void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsedWithPattern() {
+        testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(() -> {
+            when(metadata.fetch()).thenReturn(Cluster.empty());
+            consumer.subscribe(Pattern.compile("topic"));
+        });
+    }
+
+    private void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(final Runnable subscription) {
+        final String groupId = "consumerGroupA";
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+        consumer = newConsumer(config);
+        subscription.run();
+
+        consumer.poll(Duration.ZERO);
+
+        verify(fetchCollector, never()).collectFetch(any(FetchBuffer.class));
+    }
+
+    @Test
+    public void testPollReturningRecordsIfGroupIdSetAndGroupManagementIsNotUsed() {
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId("consumerGroupA"));
+        testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(config);
+    }
+
+    @Test
+    public void testPollReturningRecordsIfGroupIdNotSetAndGroupManagementIsNotUsed() {
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerProperties());
+        testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(config);
+    }
+
+    private void testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(final ConsumerConfig config) {
+        final String topic = "topic";
+        final TopicPartition topicPartition = new TopicPartition(topic, 0);
+        consumer = newConsumer(config);
+        consumer.assign(singletonList(topicPartition));
+        final List<ConsumerRecord<String, String>> records = singletonList(
+            new ConsumerRecord<>(topic, 0, 2, "key1", "value1")
+        );
+        when(fetchCollector.collectFetch(any(FetchBuffer.class)))
+            .thenReturn(Fetch.forPartition(topicPartition, records, true));
+        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
+
+        consumer.poll(Duration.ZERO);
+
+        verify(fetchCollector).collectFetch(any(FetchBuffer.class));
+    }
+
+    @Test
+    public void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsedWithTopics() {
+        final String topic = "topic";
+        testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(
+            topic,
+            () -> {
+                consumer.subscribe(singletonList(topic));
+            });
+    }
+
+    @Test
+    public void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsedWithPattern() {
+        final String topic = "topic";
+        testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(
+            topic,
+            () -> {
+                when(metadata.fetch()).thenReturn(Cluster.empty());
+                consumer.subscribe(Pattern.compile(topic));
+            });
+    }
+
+    private void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(final String topic,
+                                                                                  final Runnable subscription) {
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId("consumerGroupA"));
+        final int generation = 1;
+        final String memberId = "newMemberId";
+        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
+            generation,
+            memberId
+        );
+        backgroundEventQueue.add(groupMetadataUpdateEvent);
+        final TopicPartition topicPartition = new TopicPartition(topic, 0);
+        final List<ConsumerRecord<String, String>> records = singletonList(
+            new ConsumerRecord<>(topic, 0, 2, "key1", "value1")
+        );
+        when(fetchCollector.collectFetch(any(FetchBuffer.class)))
+            .thenReturn(Fetch.forPartition(topicPartition, records, true));
+        consumer = newConsumer(config);
+        subscription.run();
+
+        consumer.poll(Duration.ZERO);
+
+        verify(fetchCollector).collectFetch(any(FetchBuffer.class));
+    }
+
+    @Test
+    public void testGroupMetadataIsResetAfterUnsubscribe() {
+        final String groupId = "consumerGroupA";
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+        consumer = newConsumer(config);
+        consumer.subscribe(singletonList("topic"));
+        final int generation = 1;
+        final String memberId = "newMemberId";
+        final ConsumerGroupMetadata groupMetadataAfterSubscription = new ConsumerGroupMetadata(
+            groupId,
+            generation,
+            memberId,
+            Optional.empty()
+        );
+        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
+            generation,
+            memberId
+        );
+        backgroundEventQueue.add(groupMetadataUpdateEvent);
+        when(fetchCollector.collectFetch(any(FetchBuffer.class))).thenReturn(Fetch.empty());
+        consumer.poll(Duration.ZERO);
+
+        assertEquals(groupMetadataAfterSubscription, consumer.groupMetadata());
+
+        completeUnsubscribeApplicationEventSuccessfully();
+
+        consumer.unsubscribe();
+
+        final ConsumerGroupMetadata groupMetadataAfterUnsubscription = new ConsumerGroupMetadata(
+            groupId,
+            JoinGroupRequest.UNKNOWN_GENERATION_ID,
+            JoinGroupRequest.UNKNOWN_MEMBER_ID,
+            Optional.empty()
+        );
+
+        assertEquals(groupMetadataAfterUnsubscription, consumer.groupMetadata());
+    }
+
     /**
      * Tests that the consumer correctly invokes the callbacks for {@link ConsumerRebalanceListener} that was
      * specified. We don't go through the full effort to emulate heartbeats and correct group management here. We're
@@ -1321,6 +1461,7 @@ public void testEnsurePollEventSentOnConsumerPoll() {
         final TopicPartition tp = new TopicPartition("topic", 0);
         final List<ConsumerRecord<String, String>> records = singletonList(
                 new ConsumerRecord<>("topic", 0, 2, "key1", "value1"));
+        backgroundEventQueue.add(new GroupMetadataUpdateEvent(1, "memberId"));
         doAnswer(invocation -> Fetch.forPartition(tp, records, true))
                 .when(fetchCollector)
                 .collectFetch(Mockito.any(FetchBuffer.class));
@@ -1416,6 +1557,13 @@ public void testLongPollWaitIsLimited() {
             new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
         );
 
+        final int generation = 1;
+        final String memberId = "newMemberId";
+        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
+            generation,
+            memberId
+        );
+        backgroundEventQueue.add(groupMetadataUpdateEvent);
         // On the first iteration, return no data; on the second, return two records
         doAnswer(invocation -> {
             // Mock the subscription being assigned as the first fetch is collected

From 9402d525abe2d8bf844fc94c0fa03b250d3da618 Mon Sep 17 00:00:00 2001
From: Stanislav Kozlovski <stanislav_kozlovski@outlook.com>
Date: Thu, 22 Feb 2024 10:39:58 +0100
Subject: [PATCH 067/258] MINOR: Reconcile upgrade.html with kafka-site/36's
 version (#15406)

The usual flow of updating the upgrade.html docs is to first do it in apache/kafka/trunk, then cherry-pick to the relative release branch and then copy into the kafka-site repo.

It seems like this was not done with a few commits updating the 3.6.1, 3.5.2 and 3.5.1, resulting in kafka-site's latest upgrade.html containing content that isn't here. This was caught while we were adding the 3.7 upgrade docs.

This patch reconciles both files by taking the extra changes from kafka-site and placing them here. This was done by simply comparing a diff of both changes and taking the ones that apply
---
 docs/upgrade.html | 39 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/docs/upgrade.html b/docs/upgrade.html
index 628396f361f54..7f13b6f76cdd2 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -36,9 +36,9 @@ <h5><a id="upgrade_370_notable" href="#upgrade_370_notable">Notable changes in 3
         </li>
     </ul>
 
-<h4><a id="upgrade_3_6_0" href="#upgrade_3_6_0">Upgrading to 3.6.0 from any version 0.8.x through 3.5.x</a></h4>
+<h4><a id="upgrade_3_6_1" href="#upgrade_3_6_1">Upgrading to 3.6.1 from any version 0.8.x through 3.5.x</a></h4>
 
-    <h5><a id="upgrade_360_zk" href="#upgrade_360_zk">Upgrading ZooKeeper-based clusters</a></h5>
+    <h5><a id="upgrade_361_zk" href="#upgrade_361_zk">Upgrading ZooKeeper-based clusters</a></h5>
     <p><b>If you are upgrading from a version prior to 2.1.x, please see the note in step 5 below about the change to the schema used to store consumer offsets.
         Once you have changed the inter.broker.protocol.version to the latest version, it will not be possible to downgrade to a version prior to 2.1.</b></p>
 
@@ -79,7 +79,7 @@ <h5><a id="upgrade_360_zk" href="#upgrade_360_zk">Upgrading ZooKeeper-based clus
         </li>
     </ol>
 
-    <h5><a id="upgrade_360_kraft" href="#upgrade_360_kraft">Upgrading KRaft-based clusters</a></h5>
+    <h5><a id="upgrade_361_kraft" href="#upgrade_361_kraft">Upgrading KRaft-based clusters</a></h5>
     <p><b>If you are upgrading from a version prior to 3.3.0, please see the note in step 3 below. Once you have changed the metadata.version to the latest version, it will not be possible to downgrade to a version prior to 3.3-IV0.</b></p>
 
     <p><b>For a rolling upgrade:</b></p>
@@ -140,6 +140,39 @@ <h5><a id="upgrade_360_notable" href="#upgrade_360_notable">Notable changes in 3
         </li>
     </ul>
 
+<h4><a id="upgrade_3_5_2" href="#upgrade_3_5_2">Upgrading to 3.5.2 from any version 0.8.x through 3.4.x</a></h4>
+    All upgrade steps remain same as <a href="#upgrade_3_5_0">upgrading to 3.5.0</a>
+    <h5><a id="upgrade_352_notable" href="#upgrade_352_notable">Notable changes in 3.5.2</a></h5>
+    <ul>
+    <li>
+        When migrating producer ID blocks from ZK to KRaft, there could be duplicate producer IDs being given to
+        transactional or idempotent producers. This can cause long term problems since the producer IDs are
+        persisted and reused for a long time.
+        See <a href="https://issues.apache.org/jira/browse/KAFKA-15552">KAFKA-15552</a> for more details.
+    </li>
+    <li>
+        In 3.5.0 and 3.5.1, there could be an issue that the empty ISR is returned from controller after AlterPartition request
+        during rolling upgrade. This issue will impact the availability of the topic partition.
+        See <a href="https://issues.apache.org/jira/browse/KAFKA-15353">KAFKA-15353</a> for more details.
+    </li>
+</ul>
+
+<h4><a id="upgrade_3_5_1" href="#upgrade_3_5_1">Upgrading to 3.5.1 from any version 0.8.x through 3.4.x</a></h4>
+    All upgrade steps remain same as <a href="#upgrade_3_5_0">upgrading to 3.5.0</a>
+    <h5><a id="upgrade_351_notable" href="#upgrade_351_notable">Notable changes in 3.5.1</a></h5>
+    <ul>
+    <li>
+        Upgraded the dependency, snappy-java, to a version which is not vulnerable to
+        <a href="https://nvd.nist.gov/vuln/detail/CVE-2023-34455">CVE-2023-34455.</a>
+        You can find more information about the CVE at <a href="https://kafka.apache.org/cve-list#CVE-2023-34455">Kafka CVE list.</a>
+    </li>
+    <li>
+        Fixed a regression introduced in 3.3.0, which caused <code>security.protocol</code> configuration values to be restricted to
+        upper case only. After the fix, <code>security.protocol</code> values are case insensitive.
+        See <a href="https://issues.apache.org/jira/browse/KAFKA-15053">KAFKA-15053</a> for details.
+    </li>
+</ul>
+
 <h4><a id="upgrade_3_5_0" href="#upgrade_3_5_0">Upgrading to 3.5.0 from any version 0.8.x through 3.4.x</a></h4>
 
     <h5><a id="upgrade_350_zk" href="#upgrade_350_zk">Upgrading ZooKeeper-based clusters</a></h5>

From ea9450767932eb3d63aeefd5af07dbc54cb92c31 Mon Sep 17 00:00:00 2001
From: Stanislav Kozlovski <stanislav_kozlovski@outlook.com>
Date: Thu, 22 Feb 2024 10:44:24 +0100
Subject: [PATCH 068/258] MINOR: Add 3.7 upgrade notes (#15407)

This patch adds the 3.7 upgrade notes.
---
 docs/upgrade.html | 72 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 71 insertions(+), 1 deletion(-)

diff --git a/docs/upgrade.html b/docs/upgrade.html
index 7f13b6f76cdd2..33009f062e27a 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -21,19 +21,89 @@
 
 <h4><a id="upgrade_3_7_0" href="#upgrade_3_7_0">Upgrading to 3.7.0 from any version 0.8.x through 3.6.x</a></h4>
 
+
+    <h5><a id="upgrade_370_zk" href="#upgrade_370_zk">Upgrading ZooKeeper-based clusters</a></h5>
+    <p><b>If you are upgrading from a version prior to 2.1.x, please see the note in step 5 below about the change to the schema used to store consumer offsets.
+        Once you have changed the inter.broker.protocol.version to the latest version, it will not be possible to downgrade to a version prior to 2.1.</b></p>
+
+    <p><b>For a rolling upgrade:</b></p>
+
+    <ol>
+        <li>Update server.properties on all brokers and add the following properties. CURRENT_KAFKA_VERSION refers to the version you
+            are upgrading from. CURRENT_MESSAGE_FORMAT_VERSION refers to the message format version currently in use. If you have previously
+            overridden the message format version, you should keep its current value. Alternatively, if you are upgrading from a version prior
+            to 0.11.0.x, then CURRENT_MESSAGE_FORMAT_VERSION should be set to match CURRENT_KAFKA_VERSION.
+            <ul>
+                <li>inter.broker.protocol.version=CURRENT_KAFKA_VERSION (e.g. <code>3.6</code>, <code>3.5</code>, etc.)</li>
+                <li>log.message.format.version=CURRENT_MESSAGE_FORMAT_VERSION  (See <a href="#upgrade_10_performance_impact">potential performance impact
+                    following the upgrade</a> for the details on what this configuration does.)</li>
+            </ul>
+            If you are upgrading from version 0.11.0.x or above, and you have not overridden the message format, then you only need to override
+            the inter-broker protocol version.
+            <ul>
+                <li>inter.broker.protocol.version=CURRENT_KAFKA_VERSION (e.g. <code>3.6</code>, <code>3.5</code>, etc.)</li>
+            </ul>
+        </li>
+        <li>Upgrade the brokers one at a time: shut down the broker, update the code, and restart it. Once you have done so, the
+            brokers will be running the latest version and you can verify that the cluster's behavior and performance meets expectations.
+            It is still possible to downgrade at this point if there are any problems.
+        </li>
+        <li>Once the cluster's behavior and performance has been verified, bump the protocol version by editing
+            <code>inter.broker.protocol.version</code> and setting it to <code>3.7</code>.
+        </li>
+        <li>Restart the brokers one by one for the new protocol version to take effect. Once the brokers begin using the latest
+            protocol version, it will no longer be possible to downgrade the cluster to an older version.
+        </li>
+        <li>If you have overridden the message format version as instructed above, then you need to do one more rolling restart to
+            upgrade it to its latest version. Once all (or most) consumers have been upgraded to 0.11.0 or later,
+            change log.message.format.version to 3.7 on each broker and restart them one by one. Note that the older Scala clients,
+            which are no longer maintained, do not support the message format introduced in 0.11, so to avoid conversion costs
+            (or to take advantage of <a href="#upgrade_11_exactly_once_semantics">exactly once semantics</a>),
+            the newer Java clients must be used.
+        </li>
+    </ol>
+
+    <h5><a id="upgrade_370_kraft" href="#upgrade_370_kraft">Upgrading KRaft-based clusters</a></h5>
+    <p><b>If you are upgrading from a version prior to 3.3.0, please see the note in step 3 below. Once you have changed the metadata.version to the latest version, it will not be possible to downgrade to a version prior to 3.3-IV0.</b></p>
+
+    <p><b>For a rolling upgrade:</b></p>
+
+    <ol>
+        <li>Upgrade the brokers one at a time: shut down the broker, update the code, and restart it. Once you have done so, the
+            brokers will be running the latest version and you can verify that the cluster's behavior and performance meets expectations.
+        </li>
+        <li>Once the cluster's behavior and performance has been verified, bump the metadata.version by running
+            <code>
+                ./bin/kafka-features.sh upgrade --metadata 3.7
+            </code>
+        </li>
+        <li>Note that cluster metadata downgrade is not supported in this version since it has metadata changes.
+            Every <a href="https://github.com/apache/kafka/blob/trunk/server-common/src/main/java/org/apache/kafka/server/common/MetadataVersion.java">MetadataVersion</a>
+            after 3.2.x has a boolean parameter that indicates if there are metadata changes (i.e. <code>IBP_3_3_IV3(7, "3.3", "IV3", true)</code> means this version has metadata changes).
+            Given your current and target versions, a downgrade is only possible if there are no metadata changes in the versions between.</li>
+    </ol>
+
     <h5><a id="upgrade_370_notable" href="#upgrade_370_notable">Notable changes in 3.7.0</a></h5>
     <ul>
         <li>Java 11 support for the broker and tools has been deprecated and will be removed in Apache Kafka 4.0. This complements
             the previous deprecation of Java 8 for all components. Please refer to
             <a href="https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=284789510">KIP-1013</a> for more details.
         </li>
+        <li>Client APIs released prior to Apache Kafka 2.1 are now marked deprecated in 3.7 and will be removed in Apache Kafka 4.0. See <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-896%3A+Remove+old+client+protocol+API+versions+in+Kafka+4.0">KIP-896</a> for details and RPC versions that are now deprecated.
+        </li>
+        <li>Early access of the new simplified Consumer Rebalance Protocol is available, and it is not recommended for use in production environments.
+            You are encouraged to test it and provide feedback!
+            For more information about the early access feature, please check <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-848%3A+The+Next+Generation+of+the+Consumer+Rebalance+Protocol">KIP-848</a> and the <a href="https://cwiki.apache.org/confluence/display/KAFKA/The+Next+Generation+of+the+Consumer+Rebalance+Protocol+%28KIP-848%29+-+Early+Access+Release+Notes">Early Access Release Notes</a>.
+        </li>
         <li>More metrics related to Tiered Storage have been introduced. They should improve the operational experience
             of running Tiered Storage in production.
             For more detailed information, please refer to <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-963%3A+Additional+metrics+in+Tiered+Storage">KIP-963</a>.
         </li>
-	<li>Kafka Streams ships multiple KIPs for IQv2 support.
+        <li>Kafka Streams ships multiple KIPs for IQv2 support.
             See the <a href="/{{version}}/documentation/streams/upgrade-guide#streams_api_changes_370">Kafka Streams upgrade section</a> for more details.
         </li>
+        <li>All the notable changes are present in the <a href="https://kafka.apache.org/blog#apache_kafka_370_release_announcement">blog post announcing the 3.7.0 release.</a>
+        </li>
     </ul>
 
 <h4><a id="upgrade_3_6_1" href="#upgrade_3_6_1">Upgrading to 3.6.1 from any version 0.8.x through 3.5.x</a></h4>

From 98a658f871fc2c533b16fb5fd567a5ceb1c340b7 Mon Sep 17 00:00:00 2001
From: Josep Prat <josep.prat@aiven.io>
Date: Thu, 22 Feb 2024 12:11:51 +0100
Subject: [PATCH 069/258] MINOR: Update dependencies (#15404)

* MINOR: Update dependencies

Updates minor versions for our dependencies and build tool

- Jackson from 2.16.0 to 2.16.1
- JUnit from 5.10.0 to 5.10.2
  https://junit.org/junit5/docs/5.10.2/release-notes/ and https://junit.org/junit5/docs/5.10.1/release-notes/
- Mockito from 5.8.0 to 5.10.0 (only if JDK 11 or higher)
  https://github.com/mockito/mockito/releases/tag/v5.10.0 and https://github.com/mockito/mockito/releases/tag/v5.9.0
- Gradle from 8.5 to 8.6 https://docs.gradle.org/8.6/release-notes.html

Reviewers: Divij Vaidya <diviv@amazon.com>


Signed-off-by: Josep Prat <josep.prat@aiven.io>
---
 LICENSE-binary                           | 22 +++++++++++-----------
 gradle/dependencies.gradle               |  8 ++++----
 gradle/wrapper/gradle-wrapper.properties |  4 ++--
 gradlew                                  | 14 +++++++-------
 4 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/LICENSE-binary b/LICENSE-binary
index 61eb5f379668c..69361f878e73a 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -216,17 +216,17 @@ commons-lang3-3.12.0
 commons-logging-1.2
 commons-validator-1.7
 error_prone_annotations-2.10.0
-jackson-annotations-2.16.0
-jackson-core-2.16.0
-jackson-databind-2.16.0
-jackson-dataformat-csv-2.16.0
-jackson-datatype-jdk8-2.16.0
-jackson-jaxrs-base-2.16.0
-jackson-jaxrs-json-provider-2.16.0
-jackson-module-afterburner-2.16.0
-jackson-module-jaxb-annotations-2.16.0
-jackson-module-scala_2.13-2.16.0
-jackson-module-scala_2.12-2.16.0
+jackson-annotations-2.16.1
+jackson-core-2.16.1
+jackson-databind-2.16.1
+jackson-dataformat-csv-2.16.1
+jackson-datatype-jdk8-2.16.1
+jackson-jaxrs-base-2.16.1
+jackson-jaxrs-json-provider-2.16.1
+jackson-module-afterburner-2.16.1
+jackson-module-jaxb-annotations-2.16.1
+jackson-module-scala_2.13-2.16.1
+jackson-module-scala_2.12-2.16.1
 jakarta.validation-api-2.0.2
 javassist-3.29.2-GA
 jetty-client-9.4.53.v20231009
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 6067d51f3674c..1194e75d60fbc 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -61,7 +61,7 @@ String mockitoVersion
 if (scalaVersion == "2.12")
   mockitoVersion = "4.9.0"
 else if (JavaVersion.current().isCompatibleWith(JavaVersion.VERSION_11))
-  mockitoVersion = "5.8.0"
+  mockitoVersion = "5.10.0"
 else
   mockitoVersion = "4.11.0"
 
@@ -100,10 +100,10 @@ versions += [
   commonsCli: "1.4",
   commonsValidator: "1.7",
   dropwizardMetrics: "4.1.12.1",
-  gradle: "8.5",
+  gradle: "8.6",
   grgit: "4.1.1",
   httpclient: "4.5.14",
-  jackson: "2.16.0",
+  jackson: "2.16.1",
   jacoco: "0.8.10",
   javassist: "3.29.2-GA",
   jetty: "9.4.53.v20231009",
@@ -118,7 +118,7 @@ versions += [
   jfreechart: "1.0.0",
   jopt: "5.0.4",
   jose4j: "0.9.4",
-  junit: "5.10.0",
+  junit: "5.10.2",
   jqwik: "1.7.4",
   kafka_0100: "0.10.0.1",
   kafka_0101: "0.10.1.1",
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index a7a990ab2a89e..865f1ba80d1e6 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,7 +1,7 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionSha256Sum=c16d517b50dd28b3f5838f0e844b7520b8f1eb610f2f29de7e4e04a1b7c9c79b
-distributionUrl=https\://services.gradle.org/distributions/gradle-8.5-all.zip
+distributionSha256Sum=85719317abd2112f021d4f41f09ec370534ba288432065f4b477b6a3b652910d
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.6-all.zip
 networkTimeout=10000
 validateDistributionUrl=true
 zipStoreBase=GRADLE_USER_HOME
diff --git a/gradlew b/gradlew
index 66582c5e7e0ae..756d7bf326d7b 100755
--- a/gradlew
+++ b/gradlew
@@ -158,7 +158,7 @@ if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then
     case $MAX_FD in #(
       max*)
         # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked.
-        # shellcheck disable=SC3045
+        # shellcheck disable=SC2039,SC3045
         MAX_FD=$( ulimit -H -n ) ||
             warn "Could not query maximum file descriptor limit"
     esac
@@ -166,7 +166,7 @@ if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then
       '' | soft) :;; #(
       *)
         # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked.
-        # shellcheck disable=SC3045
+        # shellcheck disable=SC2039,SC3045
         ulimit -n "$MAX_FD" ||
             warn "Could not set maximum file descriptor limit to $MAX_FD"
     esac
@@ -215,11 +215,11 @@ fi
 # Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
 DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"'
 
-# Collect all arguments for the java command;
-#   * $DEFAULT_JVM_OPTS, $JAVA_OPTS, and $GRADLE_OPTS can contain fragments of
-#     shell script including quotes and variable substitutions, so put them in
-#     double quotes to make sure that they get re-expanded; and
-#   * put everything else in single quotes, so that it's not re-expanded.
+# Collect all arguments for the java command:
+#   * DEFAULT_JVM_OPTS, JAVA_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments,
+#     and any embedded shellness will be escaped.
+#   * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be
+#     treated as '${Hostname}' itself on the command line.
 
 set -- \
         "-Dorg.gradle.appname=$APP_BASE_NAME" \

From a512ef1478ae0d42cbc1fa0db39d61376961063b Mon Sep 17 00:00:00 2001
From: Stanislav Kozlovski <stanislav_kozlovski@outlook.com>
Date: Thu, 22 Feb 2024 23:55:19 +0100
Subject: [PATCH 070/258] MINOR: Add 3.7 JBOD KRaft Early Acces disclaimer
 (#15418)

This patch adds a small disclaimer to the ops documentation page to mention that JBOD is in early access with 3.7
---
 docs/ops.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/ops.html b/docs/ops.html
index 7e9d9630e2574..92f676155570e 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -3774,7 +3774,7 @@ <h4 class="anchor-heading"><a id="kraft_missing" class="anchor-link"></a><a href
   <p>The following features are not fully implemented in KRaft mode:</p>
 
   <ul>
-    <li>Supporting JBOD configurations with multiple storage directories</li>
+    <li>Supporting JBOD configurations with multiple storage directories. Note that an Early Access release is supported in 3.7 as per <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-858%3A+Handle+JBOD+broker+disk+failure+in+KRaft">KIP-858</a>. Note that it is not yet recommended for use in production environments. Please refer to the <a href="https://cwiki.apache.org/confluence/display/KAFKA/Kafka+JBOD+in+KRaft+Early+Access+Release+Notes">release notes</a> to help us test it and provide feedback at <a href="https://issues.apache.org/jira/browse/KAFKA-16061">KAFKA-16061</a>.</li>
     <li>Modifying certain dynamic configurations on the standalone KRaft controller</li>
   </ul>
 

From 661e41c92f9cee4ea97835de40ecf328330465f5 Mon Sep 17 00:00:00 2001
From: Justine Olshan <jolshan@confluent.io>
Date: Thu, 22 Feb 2024 17:10:11 -0800
Subject: [PATCH 071/258] KAFKA-16302: Remove check for log message that is no
 longer present (fix builds) (#15422)

a3528a3 removed this log but not the test asserting it.

Builds are currently red because for some reason these tests can't retry. We should address that as a followup.

Reviewers:  Greg Harris <greg.harris@aiven.io>,  Matthias J. Sax <matthias@confluent.io>
---
 ...lSchemaRocksDBSegmentedBytesStoreTest.java | 21 +++++++------------
 ...bstractRocksDBSegmentedBytesStoreTest.java | 21 +++++++------------
 .../AbstractSessionBytesStoreTest.java        | 19 +++++++----------
 .../AbstractWindowBytesStoreTest.java         | 17 ++++++---------
 4 files changed, 27 insertions(+), 51 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStoreTest.java
index c81e57589a0ad..c7af25726a4f5 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStoreTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractDualSchemaRocksDBSegmentedBytesStoreTest.java
@@ -45,7 +45,6 @@
 import org.apache.kafka.streams.processor.internals.ProcessorRecordContext;
 import org.apache.kafka.streams.processor.internals.Task.TaskType;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
-import org.apache.kafka.common.utils.LogCaptureAppender;
 import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.state.KeyValueIterator;
 import org.apache.kafka.streams.state.StateSerdes;
@@ -83,7 +82,6 @@
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.apache.kafka.streams.state.internals.WindowKeySchema.timeWindowForSize;
 import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.hasItem;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -1579,7 +1577,7 @@ private List<ConsumerRecord<byte[], byte[]>> getChangelogRecordsWithoutHeaders()
 
 
     @Test
-    public void shouldLogAndMeasureExpiredRecords() {
+    public void shouldMeasureExpiredRecords() {
         final Properties streamsConfig = StreamsTestUtils.getStreamsConfig();
         final AbstractDualSchemaRocksDBSegmentedBytesStore<S> bytesStore = getBytesStore();
         final InternalMockProcessorContext context = new InternalMockProcessorContext(
@@ -1590,18 +1588,13 @@ public void shouldLogAndMeasureExpiredRecords() {
         context.setSystemTimeMs(time.milliseconds());
         bytesStore.init((StateStoreContext) context, bytesStore);
 
-        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
-            // write a record to advance stream time, with a high enough timestamp
-            // that the subsequent record in windows[0] will already be expired.
-            bytesStore.put(serializeKey(new Windowed<>("dummy", nextSegmentWindow)), serializeValue(0));
+        // write a record to advance stream time, with a high enough timestamp
+        // that the subsequent record in windows[0] will already be expired.
+        bytesStore.put(serializeKey(new Windowed<>("dummy", nextSegmentWindow)), serializeValue(0));
 
-            final Bytes key = serializeKey(new Windowed<>("a", windows[0]));
-            final byte[] value = serializeValue(5);
-            bytesStore.put(key, value);
-
-            final List<String> messages = appender.getMessages();
-            assertThat(messages, hasItem("Skipping record for expired segment."));
-        }
+        final Bytes key = serializeKey(new Windowed<>("a", windows[0]));
+        final byte[] value = serializeValue(5);
+        bytesStore.put(key, value);
 
         final Map<MetricName, ? extends Metric> metrics = context.metrics().metrics();
         final String threadId = Thread.currentThread().getName();
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStoreTest.java
index 852f6fd3cebff..a20bebd00814c 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStoreTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractRocksDBSegmentedBytesStoreTest.java
@@ -27,7 +27,6 @@
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.serialization.Serdes;
 import org.apache.kafka.common.utils.Bytes;
-import org.apache.kafka.common.utils.LogCaptureAppender;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.SystemTime;
@@ -81,7 +80,6 @@
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.apache.kafka.streams.state.internals.WindowKeySchema.timeWindowForSize;
 import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.hasItem;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.hasEntry;
@@ -784,7 +782,7 @@ private List<ConsumerRecord<byte[], byte[]>> getChangelogRecordsWithoutHeaders()
 
 
     @Test
-    public void shouldLogAndMeasureExpiredRecords() {
+    public void shouldMeasureExpiredRecords() {
         final Properties streamsConfig = StreamsTestUtils.getStreamsConfig();
         final AbstractRocksDBSegmentedBytesStore<S> bytesStore = getBytesStore();
         final InternalMockProcessorContext context = new InternalMockProcessorContext(
@@ -795,18 +793,13 @@ public void shouldLogAndMeasureExpiredRecords() {
         context.setSystemTimeMs(time.milliseconds());
         bytesStore.init((StateStoreContext) context, bytesStore);
 
-        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
-            // write a record to advance stream time, with a high enough timestamp
-            // that the subsequent record in windows[0] will already be expired.
-            bytesStore.put(serializeKey(new Windowed<>("dummy", nextSegmentWindow)), serializeValue(0));
+        // write a record to advance stream time, with a high enough timestamp
+        // that the subsequent record in windows[0] will already be expired.
+        bytesStore.put(serializeKey(new Windowed<>("dummy", nextSegmentWindow)), serializeValue(0));
 
-            final Bytes key = serializeKey(new Windowed<>("a", windows[0]));
-            final byte[] value = serializeValue(5);
-            bytesStore.put(key, value);
-
-            final List<String> messages = appender.getMessages();
-            assertThat(messages, hasItem("Skipping record for expired segment."));
-        }
+        final Bytes key = serializeKey(new Windowed<>("a", windows[0]));
+        final byte[] value = serializeValue(5);
+        bytesStore.put(key, value);
 
         final Map<MetricName, ? extends Metric> metrics = context.metrics().metrics();
         final String threadId = Thread.currentThread().getName();
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractSessionBytesStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractSessionBytesStoreTest.java
index 56c77ce4fe16b..2d75f7513444a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractSessionBytesStoreTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractSessionBytesStoreTest.java
@@ -793,7 +793,7 @@ public void shouldReturnSameResultsForSingleKeyFindSessionsAndEqualKeyRangeFindS
     }
 
     @Test
-    public void shouldLogAndMeasureExpiredRecords() {
+    public void shouldMeasureExpiredRecords() {
         final Properties streamsConfig = StreamsTestUtils.getStreamsConfig();
         final SessionStore<String, Long> sessionStore = buildSessionStore(RETENTION_PERIOD, Serdes.String(), Serdes.Long());
         final InternalMockProcessorContext context = new InternalMockProcessorContext(
@@ -806,18 +806,13 @@ public void shouldLogAndMeasureExpiredRecords() {
         context.setSystemTimeMs(time.milliseconds());
         sessionStore.init((StateStoreContext) context, sessionStore);
 
-        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
-            // Advance stream time by inserting record with large enough timestamp that records with timestamp 0 are expired
-            // Note that rocksdb will only expire segments at a time (where segment interval = 60,000 for this retention period)
-            sessionStore.put(new Windowed<>("initial record", new SessionWindow(0, 2 * SEGMENT_INTERVAL)), 0L);
+        // Advance stream time by inserting record with large enough timestamp that records with timestamp 0 are expired
+        // Note that rocksdb will only expire segments at a time (where segment interval = 60,000 for this retention period)
+        sessionStore.put(new Windowed<>("initial record", new SessionWindow(0, 2 * SEGMENT_INTERVAL)), 0L);
 
-            // Try inserting a record with timestamp 0 -- should be dropped
-            sessionStore.put(new Windowed<>("late record", new SessionWindow(0, 0)), 0L);
-            sessionStore.put(new Windowed<>("another on-time record", new SessionWindow(0, 2 * SEGMENT_INTERVAL)), 0L);
-
-            final List<String> messages = appender.getMessages();
-            assertThat(messages, hasItem("Skipping record for expired segment."));
-        }
+        // Try inserting a record with timestamp 0 -- should be dropped
+        sessionStore.put(new Windowed<>("late record", new SessionWindow(0, 0)), 0L);
+        sessionStore.put(new Windowed<>("another on-time record", new SessionWindow(0, 2 * SEGMENT_INTERVAL)), 0L);
 
         final Map<MetricName, ? extends Metric> metrics = context.metrics().metrics();
         final String threadId = Thread.currentThread().getName();
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractWindowBytesStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractWindowBytesStoreTest.java
index d29c6bf88d45b..688efab857d1b 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractWindowBytesStoreTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractWindowBytesStoreTest.java
@@ -984,7 +984,7 @@ public void shouldNotThrowInvalidRangeExceptionWithNegativeFromKey() {
     }
 
     @Test
-    public void shouldLogAndMeasureExpiredRecords() {
+    public void shouldMeasureExpiredRecords() {
         final Properties streamsConfig = StreamsTestUtils.getStreamsConfig();
         final WindowStore<Integer, String> windowStore =
             buildWindowStore(RETENTION_PERIOD, WINDOW_SIZE, false, Serdes.Integer(), Serdes.String());
@@ -998,17 +998,12 @@ public void shouldLogAndMeasureExpiredRecords() {
         context.setTime(1L);
         windowStore.init((StateStoreContext) context, windowStore);
 
-        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
-            // Advance stream time by inserting record with large enough timestamp that records with timestamp 0 are expired
-            windowStore.put(1, "initial record", 2 * RETENTION_PERIOD);
+        // Advance stream time by inserting record with large enough timestamp that records with timestamp 0 are expired
+        windowStore.put(1, "initial record", 2 * RETENTION_PERIOD);
 
-            // Try inserting a record with timestamp 0 -- should be dropped
-            windowStore.put(1, "late record", 0L);
-            windowStore.put(1, "another on-time record", RETENTION_PERIOD + 1);
-
-            final List<String> messages = appender.getMessages();
-            assertThat(messages, hasItem("Skipping record for expired segment."));
-        }
+        // Try inserting a record with timestamp 0 -- should be dropped
+        windowStore.put(1, "late record", 0L);
+        windowStore.put(1, "another on-time record", RETENTION_PERIOD + 1);
 
         final Map<MetricName, ? extends Metric> metrics = context.metrics().metrics();
 

From 06392f7ae272af33e6789ce48fc840005fa24899 Mon Sep 17 00:00:00 2001
From: Daan Gerits <daan.gerits@gmail.com>
Date: Fri, 23 Feb 2024 03:15:24 +0100
Subject: [PATCH 072/258] MINOR: Update of the PAPI testing classes to the
 latest implementation (#12740)

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../kstream/internals/KStreamBranchTest.java  |  4 +--
 .../internals/KStreamTransformTest.java       |  4 +--
 .../internals/KStreamTransformValuesTest.java |  2 +-
 .../internals/KTableTransformValuesTest.java  |  2 +-
 .../org/apache/kafka/test/MockProcessor.java  | 33 ++++++++++---------
 .../apache/kafka/test/MockProcessorNode.java  | 15 ++++-----
 .../kafka/test/MockProcessorSupplier.java     | 13 ++++----
 7 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamBranchTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamBranchTest.java
index 8a731f9e47453..d1b2477fe69ba 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamBranchTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamBranchTest.java
@@ -59,7 +59,7 @@ public void testKStreamBranch() {
 
         assertEquals(3, branches.length);
 
-        final MockProcessorSupplier<Integer, String> supplier = new MockProcessorSupplier<>();
+        final MockProcessorSupplier<Integer, String, Void, Void> supplier = new MockProcessorSupplier<>();
         for (final KStream<Integer, String> branch : branches) {
             branch.process(supplier);
         }
@@ -71,7 +71,7 @@ public void testKStreamBranch() {
             }
         }
 
-        final List<MockProcessor<Integer, String>> processors = supplier.capturedProcessors(3);
+        final List<MockProcessor<Integer, String, Void, Void>> processors = supplier.capturedProcessors(3);
         assertEquals(3, processors.get(0).processed().size());
         assertEquals(1, processors.get(1).processed().size());
         assertEquals(2, processors.get(2).processed().size());
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformTest.java
index 5aad9f07fd133..00b9d10fc984d 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformTest.java
@@ -74,7 +74,7 @@ public void close() { }
 
         final int[] expectedKeys = {1, 10, 100, 1000};
 
-        final MockProcessorSupplier<Integer, Integer> processor = new MockProcessorSupplier<>();
+        final MockProcessorSupplier<Integer, Integer, Void, Void> processor = new MockProcessorSupplier<>();
         final KStream<Integer, Integer> stream = builder.stream(TOPIC_NAME, Consumed.with(Serdes.Integer(), Serdes.Integer()));
         stream.transform(transformerSupplier).process(processor);
 
@@ -136,7 +136,7 @@ public void close() { }
 
         final int[] expectedKeys = {1, 10, 100, 1000};
 
-        final MockProcessorSupplier<Integer, Integer> processor = new MockProcessorSupplier<>();
+        final MockProcessorSupplier<Integer, Integer, Void, Void> processor = new MockProcessorSupplier<>();
         final KStream<Integer, Integer> stream = builder.stream(TOPIC_NAME, Consumed.with(Serdes.Integer(), Serdes.Integer()));
         stream.transform(transformerSupplier).process(processor);
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformValuesTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformValuesTest.java
index 05c07b00dbb9a..0063c6554494a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformValuesTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamTransformValuesTest.java
@@ -48,7 +48,7 @@
 @RunWith(MockitoJUnitRunner.StrictStubs.class)
 public class KStreamTransformValuesTest {
     private final String topicName = "topic";
-    private final MockProcessorSupplier<Integer, Integer> supplier = new MockProcessorSupplier<>();
+    private final MockProcessorSupplier<Integer, Integer, Void, Void> supplier = new MockProcessorSupplier<>();
     private final Properties props = StreamsTestUtils.getStreamsConfig(Serdes.Integer(), Serdes.Integer());
     @Mock
     private InternalProcessorContext context;
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableTransformValuesTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableTransformValuesTest.java
index 0bba2e406720d..3138c71cf8bfe 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableTransformValuesTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableTransformValuesTest.java
@@ -84,7 +84,7 @@ public class KTableTransformValuesTest {
     private static final Consumed<String, String> CONSUMED = Consumed.with(Serdes.String(), Serdes.String());
 
     private TopologyTestDriver driver;
-    private MockProcessorSupplier<String, String> capture;
+    private MockProcessorSupplier<String, String, Void, Void> capture;
     private StreamsBuilder builder;
     @Mock
     private KTableImpl<String, String, String> parent;
diff --git a/streams/src/test/java/org/apache/kafka/test/MockProcessor.java b/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
index dc2a38938a81a..9766d1c0fe85c 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
@@ -18,8 +18,9 @@
 
 import org.apache.kafka.streams.KeyValueTimestamp;
 import org.apache.kafka.streams.processor.Cancellable;
-import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.PunctuationType;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.processor.internals.InternalProcessorContext;
 import org.apache.kafka.streams.state.ValueAndTimestamp;
@@ -28,9 +29,9 @@
 import java.util.List;
 import java.util.Map;
 
-@SuppressWarnings("deprecation") // Old PAPI. Needs to be migrated.
-public class MockProcessor<K, V> extends org.apache.kafka.streams.processor.AbstractProcessor<K, V> {
-    private final MockApiProcessor<K, V, Object, Object> delegate;
+public class MockProcessor<KIn, VIn, KOut, VOut> implements Processor<KIn, VIn, KOut, VOut> {
+    private final MockApiProcessor<KIn, VIn, KOut, VOut> delegate;
+
 
     public MockProcessor(final PunctuationType punctuationType,
                          final long scheduleInterval) {
@@ -41,16 +42,14 @@ public MockProcessor() {
         delegate = new MockApiProcessor<>();
     }
 
-    @SuppressWarnings("unchecked")
     @Override
-    public void init(final ProcessorContext context) {
-        super.init(context);
-        delegate.init((org.apache.kafka.streams.processor.api.ProcessorContext<Object, Object>) context);
+    public void init(ProcessorContext<KOut, VOut> context) {
+        delegate.init(context);
     }
 
     @Override
-    public void process(final K key, final V value) {
-        delegate.process(new Record<>(key, value, context.timestamp(), context.headers()));
+    public void process(Record<KIn, VIn> record) {
+        delegate.process(record);
     }
 
     public void checkAndClearProcessResult(final KeyValueTimestamp<?, ?>... expected) {
@@ -69,7 +68,7 @@ public void checkAndClearPunctuateResult(final PunctuationType type, final long.
         delegate.checkAndClearPunctuateResult(type, expected);
     }
 
-    public Map<K, ValueAndTimestamp<V>> lastValueAndTimestampPerKey() {
+    public Map<KIn, ValueAndTimestamp<VIn>> lastValueAndTimestampPerKey() {
         return delegate.lastValueAndTimestampPerKey();
     }
 
@@ -81,14 +80,18 @@ public Cancellable scheduleCancellable() {
         return delegate.scheduleCancellable();
     }
 
-    public ArrayList<KeyValueTimestamp<K, V>> processed() {
+    public ArrayList<KeyValueTimestamp<KIn, VIn>> processed() {
         return delegate.processed();
     }
 
-    @SuppressWarnings("unchecked")
     public void addProcessorMetadata(final String key, final long value) {
-        if (context instanceof InternalProcessorContext) {
-            ((InternalProcessorContext<K, V>) context).addProcessorMetadataKeyValue(key, value);
+        if (delegate.context() instanceof InternalProcessorContext) {
+            ((InternalProcessorContext<KOut, VOut>) delegate.context()).addProcessorMetadataKeyValue(key, value);
         }
     }
+
+    @Override
+    public void close() {
+        delegate.close();
+    }
 }
diff --git a/streams/src/test/java/org/apache/kafka/test/MockProcessorNode.java b/streams/src/test/java/org/apache/kafka/test/MockProcessorNode.java
index 039c831e4dfcb..908641a0e6455 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockProcessorNode.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockProcessorNode.java
@@ -16,20 +16,20 @@
  */
 package org.apache.kafka.test;
 
-import java.util.Collections;
-import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.kafka.streams.processor.PunctuationType;
 import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.processor.internals.InternalProcessorContext;
-import org.apache.kafka.streams.processor.internals.ProcessorAdapter;
 import org.apache.kafka.streams.processor.internals.ProcessorNode;
 
+import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
+
 public class MockProcessorNode<KIn, VIn, KOut, VOut> extends ProcessorNode<KIn, VIn, KOut, VOut> {
 
     private static final String NAME = "MOCK-PROCESS-";
     private static final AtomicInteger INDEX = new AtomicInteger(1);
 
-    public final MockProcessor<KIn, VIn> mockProcessor;
+    public final MockProcessor<KIn, VIn, KOut, VOut> mockProcessor;
 
     public boolean closed;
     public boolean initialized;
@@ -46,9 +46,8 @@ public MockProcessorNode() {
         this(new MockProcessor<>());
     }
 
-    @SuppressWarnings("unchecked")
-    private MockProcessorNode(final MockProcessor<KIn, VIn> mockProcessor) {
-        super(NAME + INDEX.getAndIncrement(), ProcessorAdapter.adapt(mockProcessor), Collections.<String>emptySet());
+    private MockProcessorNode(final MockProcessor<KIn, VIn, KOut, VOut> mockProcessor) {
+        super(NAME + INDEX.getAndIncrement(), mockProcessor, Collections.<String>emptySet());
 
         this.mockProcessor = mockProcessor;
     }
@@ -61,7 +60,7 @@ public void init(final InternalProcessorContext<KOut, VOut> context) {
 
     @Override
     public void process(final Record<KIn, VIn> record) {
-        mockProcessor.process(record.key(), record.value());
+        mockProcessor.process(record);
     }
 
     @Override
diff --git a/streams/src/test/java/org/apache/kafka/test/MockProcessorSupplier.java b/streams/src/test/java/org/apache/kafka/test/MockProcessorSupplier.java
index c6b70f2763d82..deab0ac33c457 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockProcessorSupplier.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockProcessorSupplier.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.test;
 
 import org.apache.kafka.streams.processor.PunctuationType;
+import org.apache.kafka.streams.processor.api.Processor;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -24,11 +25,11 @@
 import static org.junit.Assert.assertEquals;
 
 @SuppressWarnings("deprecation") // Old PAPI. Needs to be migrated.
-public class MockProcessorSupplier<K, V> implements org.apache.kafka.streams.processor.ProcessorSupplier<K, V> {
+public class MockProcessorSupplier<KIn, VIn, KOut, VOut> implements org.apache.kafka.streams.processor.api.ProcessorSupplier<KIn, VIn, KOut, VOut> {
 
     private final long scheduleInterval;
     private final PunctuationType punctuationType;
-    private final List<MockProcessor<K, V>> processors = new ArrayList<>();
+    private final List<MockProcessor<KIn, VIn, KOut, VOut>> processors = new ArrayList<>();
 
     public MockProcessorSupplier() {
         this(-1L);
@@ -44,8 +45,8 @@ public MockProcessorSupplier(final long scheduleInterval, final PunctuationType
     }
 
     @Override
-    public org.apache.kafka.streams.processor.Processor<K, V> get() {
-        final MockProcessor<K, V> processor = new MockProcessor<>(punctuationType, scheduleInterval);
+    public Processor<KIn, VIn, KOut, VOut> get() {
+        final MockProcessor<KIn, VIn, KOut, VOut> processor = new MockProcessor<>(punctuationType, scheduleInterval);
 
         // to keep tests simple, ignore calls from ApiUtils.checkSupplier
         if (!StreamsTestUtils.isCheckSupplierCall()) {
@@ -56,7 +57,7 @@ public org.apache.kafka.streams.processor.Processor<K, V> get() {
     }
 
     // get the captured processor assuming that only one processor gets returned from this supplier
-    public MockProcessor<K, V> theCapturedProcessor() {
+    public MockProcessor<KIn, VIn, KOut, VOut> theCapturedProcessor() {
         return capturedProcessors(1).get(0);
     }
 
@@ -65,7 +66,7 @@ public int capturedProcessorsCount() {
     }
 
         // get the captured processors with the expected number
-    public List<MockProcessor<K, V>> capturedProcessors(final int expectedNumberOfProcessors) {
+    public List<MockProcessor<KIn, VIn, KOut, VOut>> capturedProcessors(final int expectedNumberOfProcessors) {
         assertEquals(expectedNumberOfProcessors, processors.size());
 
         return processors;

From 2185004083ebb8f0b3a443132b5a33908c459c65 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Fri, 23 Feb 2024 04:56:05 -0500
Subject: [PATCH 073/258] KAFKA-16251: Fix for not sending heartbeat while
 fenced (#15392)

Fix to ensure that a consumer that has been fenced by the coordinator stops sending heartbeats while it is on the FENCED state releasing its assignment (waiting for the onPartitionsLost callback to complete). Once the callback completes, the member transitions to JOINING and it's then when it should resume sending heartbeats again.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/MembershipManagerImpl.java      | 14 +++++-----
 .../HeartbeatRequestManagerTest.java          | 27 +++++++++++++++++++
 .../internals/MembershipManagerImplTest.java  |  2 ++
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index dd035506d4bea..81e65dfd86612 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -720,18 +720,16 @@ private boolean targetAssignmentReconciled() {
     }
 
     /**
-     * @return True if the member should not send heartbeats, which would be one of the following
-     * cases:
-     * <ul>
-     * <li>Member is not subscribed to any topics</li>
-     * <li>Member has received a fatal error in a previous heartbeat response</li>
-     * <li>Member is stale, meaning that it has left the group due to expired poll timer</li>
-     * </ul>
+     * @return True if the member should not send heartbeats, which is the case when it is in a
+     * state where it is not an active member of the group.
      */
     @Override
     public boolean shouldSkipHeartbeat() {
         MemberState state = state();
-        return state == MemberState.UNSUBSCRIBED || state == MemberState.FATAL || state == MemberState.STALE;
+        return state == MemberState.UNSUBSCRIBED ||
+            state == MemberState.FATAL ||
+            state == MemberState.STALE ||
+            state == MemberState.FENCED;
     }
 
     /**
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 4016b74b27bd6..72a5c0349d274 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -663,6 +663,33 @@ public void testHeartbeatMetrics() {
         assertEquals((double) randomSleepS, getMetric("last-heartbeat-seconds-ago").metricValue());
     }
 
+    @Test
+    public void testFencedMemberStopHeartbeatUntilItReleasesAssignmentToRejoin() {
+        mockStableMember();
+
+        time.sleep(DEFAULT_HEARTBEAT_INTERVAL_MS);
+        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
+        assertEquals(1, result.unsentRequests.size());
+
+        // Receive HB response fencing member
+        when(subscriptions.hasAutoAssignedPartitions()).thenReturn(true);
+        doNothing().when(membershipManager).transitionToFenced();
+        ClientResponse response = createHeartbeatResponse(result.unsentRequests.get(0), Errors.FENCED_MEMBER_EPOCH);
+        result.unsentRequests.get(0).handler().onComplete(response);
+
+        verify(membershipManager).transitionToFenced();
+        verify(heartbeatRequestState).onFailedAttempt(anyLong());
+        verify(heartbeatRequestState).reset();
+
+        when(membershipManager.state()).thenReturn(MemberState.FENCED);
+        result = heartbeatRequestManager.poll(time.milliseconds());
+        assertEquals(0, result.unsentRequests.size(), "Member should not send heartbeats while FENCED");
+
+        when(membershipManager.state()).thenReturn(MemberState.JOINING);
+        result = heartbeatRequestManager.poll(time.milliseconds());
+        assertEquals(1, result.unsentRequests.size(), "Fenced member should resume heartbeat after transitioning to JOINING");
+    }
+
     private void assertHeartbeat(HeartbeatRequestManager hrm, int nextPollMs) {
         NetworkClientDelegate.PollResult pollResult = hrm.poll(time.milliseconds());
         assertEquals(1, pollResult.unsentRequests.size());
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index f5c65d58bfc54..50f28bb523367 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -1705,6 +1705,8 @@ private void testOnPartitionsLost(Optional<RuntimeException> lostError) {
         assertEquals(0, listener.assignedCount());
         assertEquals(0, listener.lostCount());
 
+        assertTrue(membershipManager.shouldSkipHeartbeat(), "Member should not send heartbeat while fenced");
+
         // Step 3: invoke the callback
         performCallback(
                 membershipManager,

From 7ac50a86113c9a489667451eecf61b6ab80870db Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Fri, 23 Feb 2024 16:43:50 +0100
Subject: [PATCH 074/258] KAFKA-16152: Fix
 PlaintextConsumerTest.testStaticConsumerDetectsNewPartitionCreatedAfterRestart
 (#15419)

The group coordinator expects the instance ID to always be sent when
leaving the group in a static membership configuration, see

https://github.com/apache/kafka/blob/ea9450767932eb3d63aeefd5af07dbc54cb92c31/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java#L814

The failure was silent, because the group coordinator does not log
failed requests and the consumer doesn't wait for the heartbeat response
during close.

Reviewers: Matthias J. Sax <matthias@confluent.io>, Kirk True <ktrue@confluent.io>, Bruno Cadonna <cadonna@apache.org>
---
 .../clients/consumer/internals/HeartbeatRequestManager.java   | 3 ++-
 .../scala/integration/kafka/api/PlaintextConsumerTest.scala   | 4 +---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 7408f584cdb4f..2aabcfe20bc81 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -546,8 +546,9 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
             data.setMemberEpoch(membershipManager.memberEpoch());
 
             // InstanceId - only sent if has changed since the last heartbeat
+            // Always send when leaving the group as a static member
             membershipManager.groupInstanceId().ifPresent(groupInstanceId -> {
-                if (!groupInstanceId.equals(sentFields.instanceId)) {
+                if (!groupInstanceId.equals(sentFields.instanceId) || membershipManager.memberEpoch() == ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH) {
                     data.setInstanceId(groupInstanceId);
                     sentFields.instanceId = groupInstanceId;
                 }
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index c96288703fb99..c826b8fd0c437 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -2121,10 +2121,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertThrows(classOf[KafkaException], () => createConsumer(configOverrides = consumer1Config))
   }
 
-  // TODO: Enable this test for both protocols when the Jira tracking its failure (KAFKA-16152) is fixed. This
-  //       is done by setting the @MethodSource value to "getTestQuorumAndGroupProtocolParametersAll"
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testStaticConsumerDetectsNewPartitionCreatedAfterRestart(quorum:String, groupProtocol: String): Unit = {
     val foo = "foo"
     val foo0 = new TopicPartition(foo, 0)

From 7efdf40a77cf1c929e785ba04fa2363d4989c4c4 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Fri, 23 Feb 2024 17:21:48 +0100
Subject: [PATCH 075/258] MINOR: Improve test testCommitInRebalanceCallback
 (#15420)

Commit 1442862 introduced a test with an assertion inside a listener
callback. This improves the test by checking that the listener is
actually being executed, to avoid silent skipping of the assertion.

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../clients/consumer/internals/AsyncKafkaConsumerTest.java    | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index abd0f45c73f3c..976677dec8628 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -99,6 +99,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -437,11 +438,13 @@ public void testCommitInRebalanceCallback() {
         CompletableBackgroundEvent<Void> e = new ConsumerRebalanceListenerCallbackNeededEvent(ON_PARTITIONS_REVOKED, sortedPartitions);
         backgroundEventQueue.add(e);
         completeCommitSyncApplicationEventSuccessfully();
+        final AtomicBoolean callbackExecuted = new AtomicBoolean(false);
 
         ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
             @Override
             public void onPartitionsRevoked(final Collection<TopicPartition> partitions) {
                 assertDoesNotThrow(() -> consumer.commitSync(mkMap(mkEntry(tp, new OffsetAndMetadata(0)))));
+                callbackExecuted.set(true);
             }
 
             @Override
@@ -452,6 +455,7 @@ public void onPartitionsAssigned(final Collection<TopicPartition> partitions) {
 
         consumer.subscribe(Collections.singletonList(topicName), listener);
         consumer.poll(Duration.ZERO);
+        assertTrue(callbackExecuted.get());
     }
 
     @Test

From 474f8c1ad620d1f982ac15cd418ae8462dec4d9e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?=
 <jsancio@users.noreply.github.com>
Date: Fri, 23 Feb 2024 12:56:25 -0800
Subject: [PATCH 076/258] KAFKA-16286; Notify listener of latest leader and
 epoch (#15397)

KRaft was only notifying listeners of the latest leader and epoch when the replica transition to a new state. This can result in the listener never getting notified if the registration happened after it had become a follower.

This problem doesn't exists for the active leader because the KRaft implementation attempts to notified the listener of the latest leader and epoch when the replica is the active leader.

This issue is fixed by notifying the listeners of the latest leader and epoch after processing the listener registration request.

Reviewers: Colin P. McCabe <cmccabe@apache.org>
---
 .../apache/kafka/raft/KafkaRaftClient.java    |  8 +++
 .../kafka/raft/KafkaRaftClientTest.java       | 60 +++++++++++++++++++
 .../kafka/raft/RaftClientTestContext.java     |  6 ++
 3 files changed, 74 insertions(+)

diff --git a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
index fd98774355cd6..0f4c792132182 100644
--- a/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
+++ b/raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java
@@ -2202,6 +2202,14 @@ private void pollListeners() {
         quorum.highWatermark().ifPresent(highWatermarkMetadata -> {
             updateListenersProgress(highWatermarkMetadata.offset);
         });
+
+        // Notify the new listeners of the latest leader and epoch
+        Optional<LeaderState<T>> leaderState = quorum.maybeLeaderState();
+        if (leaderState.isPresent()) {
+            maybeFireLeaderChange(leaderState.get());
+        } else {
+            maybeFireLeaderChange();
+        }
     }
 
     private void processRegistration(Registration<T> registration) {
diff --git a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java
index 51a5938220f04..f18a0ff2e882e 100644
--- a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java
+++ b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java
@@ -2958,6 +2958,66 @@ public void testHandleCommitCallbackFiresInCandidateState() throws Exception {
         assertEquals(OptionalInt.empty(), secondListener.currentClaimedEpoch());
     }
 
+    @Test
+    public void testHandleLeaderChangeFiresAfterUnattachedRegistration() throws Exception {
+        // When registering a listener while the replica is unattached, it should get notified
+        // with the current epoch
+        // When transitioning to follower, expect another notification with the leader and epoch
+
+        int localId = 0;
+        int otherNodeId = 1;
+        int epoch = 7;
+        Set<Integer> voters = Utils.mkSet(localId, otherNodeId);
+
+        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
+            .withUnknownLeader(epoch)
+            .build();
+
+        // Register another listener and verify that it is notified of latest epoch
+        RaftClientTestContext.MockListener secondListener = new RaftClientTestContext.MockListener(
+            OptionalInt.of(localId)
+        );
+        context.client.register(secondListener);
+        context.client.poll();
+
+        // Expected leader change notification
+        LeaderAndEpoch expectedLeaderAndEpoch = new LeaderAndEpoch(OptionalInt.empty(), epoch);
+        assertEquals(expectedLeaderAndEpoch, secondListener.currentLeaderAndEpoch());
+
+        // Transition to follower and the expect a leader changed notification
+        context.deliverRequest(context.beginEpochRequest(epoch, otherNodeId));
+        context.pollUntilResponse();
+
+        // Expected leader change notification
+        expectedLeaderAndEpoch = new LeaderAndEpoch(OptionalInt.of(otherNodeId), epoch);
+        assertEquals(expectedLeaderAndEpoch, secondListener.currentLeaderAndEpoch());
+    }
+
+    @Test
+    public void testHandleLeaderChangeFiresAfterFollowerRegistration() throws Exception {
+        // When registering a listener while the replica is a follower, it should get notified with
+        // the current leader and epoch
+
+        int localId = 0;
+        int otherNodeId = 1;
+        int epoch = 7;
+        Set<Integer> voters = Utils.mkSet(localId, otherNodeId);
+
+        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
+            .withElectedLeader(epoch, otherNodeId)
+            .build();
+
+        // Register another listener and verify that it is notified of latest leader and epoch
+        RaftClientTestContext.MockListener secondListener = new RaftClientTestContext.MockListener(
+            OptionalInt.of(localId)
+        );
+        context.client.register(secondListener);
+        context.client.poll();
+
+        LeaderAndEpoch expectedLeaderAndEpoch = new LeaderAndEpoch(OptionalInt.of(otherNodeId), epoch);
+        assertEquals(expectedLeaderAndEpoch, secondListener.currentLeaderAndEpoch());
+    }
+
     @Test
     public void testObserverFetchWithNoLocalId() throws Exception {
         // When no `localId` is defined, the client will behave as an observer.
diff --git a/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java b/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
index 68241ae70b19f..ea4078b7c5663 100644
--- a/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
+++ b/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
@@ -83,6 +83,7 @@
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public final class RaftClientTestContext {
@@ -1195,6 +1196,11 @@ public void handleLeaderChange(LeaderAndEpoch leaderAndEpoch) {
             // We record the current committed offset as the claimed epoch's start
             // offset. This is useful to verify that the `handleLeaderChange` callback
             // was not received early on the leader.
+            assertTrue(
+                leaderAndEpoch.epoch() >= currentLeaderAndEpoch.epoch(),
+                String.format("new epoch (%d) not >= than old epoch (%d)", leaderAndEpoch.epoch(), currentLeaderAndEpoch.epoch())
+            );
+            assertNotEquals(currentLeaderAndEpoch, leaderAndEpoch);
             this.currentLeaderAndEpoch = leaderAndEpoch;
 
             currentClaimedEpoch().ifPresent(claimedEpoch -> {

From b4e96913cc6c827968e47a31261e0bd8fdf677b5 Mon Sep 17 00:00:00 2001
From: Yang Yu <yyu@confluent.io>
Date: Sat, 24 Feb 2024 08:09:23 -0600
Subject: [PATCH 077/258] KAFKA-16265: KIP-994 (Part 1) Minor Enhancements to
 ListTransactionsRequest (#15384)

Introduces a new filter in ListTransactionsRequest API. This enables caller to filter on transactions that have been running for longer than a certain duration of time.

This PR includes the following changes:

bumps version for ListTransactionsRequest API to 1. Set the durationFilter to -1L when communicating with an older broker that does not support version 1.
bumps version for ListTransactionsResponse to 1 without changing the response structure.
adds durationFilter option to kafka-transactions.sh --list
Tests:

Client side test to build request with correct combination of duration filter and API version: testBuildRequestWithDurationFilter
Server side test to filter transactions based on duration: testListTransactionsFiltering
Added test case for kafka-transactions.sh change in TransactionsCommandTest

Reviewers: Justine Olshan <jolshan@confluent.io>, Raman Verma <rverma@confluent.io>
---
 .../org/apache/kafka/clients/admin/Admin.java |  3 +-
 .../admin/ListTransactionsOptions.java        | 29 +++++++++++++++++--
 .../internals/ListTransactionsHandler.java    |  1 +
 .../requests/ListTransactionsRequest.java     |  5 ++++
 .../message/ListTransactionsRequest.json      |  6 +++-
 .../message/ListTransactionsResponse.json     |  3 +-
 .../ListTransactionsHandlerTest.java          | 26 +++++++++++++++++
 .../transaction/TransactionCoordinator.scala  |  5 ++--
 .../transaction/TransactionStateManager.scala |  6 +++-
 .../main/scala/kafka/server/KafkaApis.scala   |  3 +-
 .../TransactionStateManagerTest.scala         | 19 ++++++++----
 .../unit/kafka/server/KafkaApisTest.scala     |  4 +--
 .../kafka/tools/TransactionsCommand.java      | 14 +++++++--
 .../kafka/tools/TransactionsCommandTest.java  | 21 ++++++++++++--
 14 files changed, 124 insertions(+), 21 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
index ff7f4e661d692..ff0e60e766d51 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
@@ -1633,7 +1633,8 @@ default ListTransactionsResult listTransactions() {
      * coordinators in the cluster and collect the state of all transactions. Users
      * should typically attempt to reduce the size of the result set using
      * {@link ListTransactionsOptions#filterProducerIds(Collection)} or
-     * {@link ListTransactionsOptions#filterStates(Collection)}
+     * {@link ListTransactionsOptions#filterStates(Collection)} or
+     * {@link ListTransactionsOptions#durationFilter(Long)}
      *
      * @param options Options to control the method behavior (including filters)
      * @return The result
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListTransactionsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListTransactionsOptions.java
index c23d4441dfde1..49cf484f5d625 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/ListTransactionsOptions.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListTransactionsOptions.java
@@ -35,6 +35,7 @@ public class ListTransactionsOptions extends AbstractOptions<ListTransactionsOpt
     private Set<TransactionState> filteredStates = Collections.emptySet();
     private Set<Long> filteredProducerIds = Collections.emptySet();
 
+    private long filteredDuration = -1L;
     /**
      * Filter only the transactions that are in a specific set of states. If no filter
      * is specified or if the passed set of states is empty, then transactions in all
@@ -61,6 +62,19 @@ public ListTransactionsOptions filterProducerIds(Collection<Long> producerIdFilt
         return this;
     }
 
+    /**
+     * Filter only the transactions that are running longer than the specified duration.
+     * If no filter is specified or if the passed duration ms is less than 0,
+     * then the all transactions will be returned.
+     *
+     * @param durationMs the duration in milliseconds to filter by
+     * @return this object
+     */
+    public ListTransactionsOptions filterOnDuration(long durationMs) {
+        this.filteredDuration = durationMs;
+        return this;
+    }
+
     /**
      * Returns the set of states to be filtered or empty if no states have been specified.
      *
@@ -81,11 +95,21 @@ public Set<Long> filteredProducerIds() {
         return filteredProducerIds;
     }
 
+    /**
+     * Returns the duration ms value being filtered.
+     *
+     * @return the current duration filter value in ms (negative value means transactions are not filtered by duration)
+     */
+    public long filteredDuration() {
+        return filteredDuration;
+    }
+
     @Override
     public String toString() {
         return "ListTransactionsOptions(" +
             "filteredStates=" + filteredStates +
             ", filteredProducerIds=" + filteredProducerIds +
+            ", filteredDuration=" + filteredDuration +
             ", timeoutMs=" + timeoutMs +
             ')';
     }
@@ -96,11 +120,12 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) return false;
         ListTransactionsOptions that = (ListTransactionsOptions) o;
         return Objects.equals(filteredStates, that.filteredStates) &&
-            Objects.equals(filteredProducerIds, that.filteredProducerIds);
+            Objects.equals(filteredProducerIds, that.filteredProducerIds) &&
+            Objects.equals(filteredDuration, that.filteredDuration);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(filteredStates, filteredProducerIds);
+        return Objects.hash(filteredStates, filteredProducerIds, filteredDuration);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandler.java
index ca249bca7f4a4..b77cf762e1b1d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandler.java
@@ -73,6 +73,7 @@ public ListTransactionsRequest.Builder buildBatchedRequest(
         request.setStateFilters(options.filteredStates().stream()
             .map(TransactionState::toString)
             .collect(Collectors.toList()));
+        request.setDurationFilter(options.filteredDuration());
         return new ListTransactionsRequest.Builder(request);
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListTransactionsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ListTransactionsRequest.java
index 0651f1fe6e5cc..a5fef3ee7b29a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListTransactionsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListTransactionsRequest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.message.ListTransactionsRequestData;
 import org.apache.kafka.common.message.ListTransactionsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -35,6 +36,10 @@ public Builder(ListTransactionsRequestData data) {
 
         @Override
         public ListTransactionsRequest build(short version) {
+            if (data.durationFilter() >= 0 && version < 1) {
+                throw new UnsupportedVersionException("Duration filter can be set only when using API version 1 or higher." +
+                        " If client is connected to an older broker, do not specify duration filter or set duration filter to -1.");
+            }
             return new ListTransactionsRequest(data, version);
         }
 
diff --git a/clients/src/main/resources/common/message/ListTransactionsRequest.json b/clients/src/main/resources/common/message/ListTransactionsRequest.json
index 21f4552cd0365..2aeeaa62e28c8 100644
--- a/clients/src/main/resources/common/message/ListTransactionsRequest.json
+++ b/clients/src/main/resources/common/message/ListTransactionsRequest.json
@@ -18,7 +18,8 @@
   "type": "request",
   "listeners": ["zkBroker", "broker"],
   "name": "ListTransactionsRequest",
-  "validVersions": "0",
+  // Version 1: adds DurationFilter to list transactions older than specified duration
+  "validVersions": "0-1",
   "flexibleVersions": "0+",
   "fields": [
     { "name": "StateFilters", "type": "[]string", "versions": "0+",
@@ -26,6 +27,9 @@
     },
     { "name": "ProducerIdFilters", "type": "[]int64", "versions": "0+", "entityType": "producerId",
       "about": "The producerIds to filter by: if empty, all transactions will be returned; if non-empty, only transactions which match one of the filtered producerIds will be returned"
+    },
+    { "name": "DurationFilter", "type": "int64", "versions": "1+", "default": -1,
+      "about": "Duration (in millis) to filter by: if < 0, all transactions will be returned; otherwise, only transactions running longer than this duration will be returned"
     }
   ]
 }
diff --git a/clients/src/main/resources/common/message/ListTransactionsResponse.json b/clients/src/main/resources/common/message/ListTransactionsResponse.json
index 2f17873239143..e9924801cc059 100644
--- a/clients/src/main/resources/common/message/ListTransactionsResponse.json
+++ b/clients/src/main/resources/common/message/ListTransactionsResponse.json
@@ -17,7 +17,8 @@
   "apiKey": 66,
   "type": "response",
   "name": "ListTransactionsResponse",
-  "validVersions": "0",
+  // Version 1 is the same as version 0 (KIP-994).
+  "validVersions": "0-1",
   "flexibleVersions": "0+",
   "fields": [
       { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandlerTest.java
index 3c54fad65e8f9..89582529212da 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListTransactionsHandlerTest.java
@@ -22,6 +22,7 @@
 import org.apache.kafka.clients.admin.internals.AdminApiHandler.ApiResult;
 import org.apache.kafka.clients.admin.internals.AllBrokersStrategy.BrokerKey;
 import org.apache.kafka.common.Node;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.message.ListTransactionsResponseData;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.ListTransactionsRequest;
@@ -41,6 +42,7 @@
 import static java.util.Collections.singleton;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class ListTransactionsHandlerTest {
     private final LogContext logContext = new LogContext();
@@ -83,6 +85,30 @@ public void testBuildRequestWithFilteredState() {
         assertEquals(Collections.emptyList(), request.data().producerIdFilters());
     }
 
+    @Test
+    public void testBuildRequestWithDurationFilter() {
+        int brokerId = 1;
+        BrokerKey brokerKey = new BrokerKey(OptionalInt.of(brokerId));
+        ListTransactionsOptions options = new ListTransactionsOptions();
+        ListTransactionsHandler handler = new ListTransactionsHandler(options, logContext);
+        // case 1: check the default value (-1L) for durationFilter
+        ListTransactionsRequest request = handler.buildBatchedRequest(brokerId, singleton(brokerKey)).build((short) 1);
+        assertEquals(-1L, request.data().durationFilter());
+        request = handler.buildBatchedRequest(brokerId, singleton(brokerKey)).build((short) 0);
+        assertEquals(-1L, request.data().durationFilter());
+        // case 2: able to set a valid duration filter when using API version 1
+        options.filterOnDuration(10L);
+        request = handler.buildBatchedRequest(brokerId, singleton(brokerKey)).build((short) 1);
+        assertEquals(10L, request.data().durationFilter());
+        assertEquals(Collections.emptyList(), request.data().producerIdFilters());
+        // case 3: unable to set a valid duration filter when using API version 0
+        assertThrows(UnsupportedVersionException.class, () -> handler.buildBatchedRequest(brokerId, singleton(brokerKey)).build((short) 0));
+        // case 4: able to set duration filter to -1L when using API version 0
+        options.filterOnDuration(-1L);
+        ListTransactionsRequest request1 = handler.buildBatchedRequest(brokerId, singleton(brokerKey)).build((short) 0);
+        assertEquals(-1L, request1.data().durationFilter());
+    }
+
     @Test
     public void testHandleSuccessfulResponse() {
         int brokerId = 1;
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
index 74366317b40ea..c7f5a16b0662e 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
@@ -278,12 +278,13 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
 
   def handleListTransactions(
     filteredProducerIds: Set[Long],
-    filteredStates: Set[String]
+    filteredStates: Set[String],
+    filteredDuration: Long = -1L
   ): ListTransactionsResponseData = {
     if (!isActive.get()) {
       new ListTransactionsResponseData().setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code)
     } else {
-      txnManager.listTransactionStates(filteredProducerIds, filteredStates)
+      txnManager.listTransactionStates(filteredProducerIds, filteredStates, filteredDuration)
     }
   }
 
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
index c4e71ca76699e..355e848980746 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
@@ -301,7 +301,8 @@ class TransactionStateManager(brokerId: Int,
 
   def listTransactionStates(
     filterProducerIds: Set[Long],
-    filterStateNames: Set[String]
+    filterStateNames: Set[String],
+    filterDurationMs: Long
   ): ListTransactionsResponseData = {
     inReadLock(stateLock) {
       val response = new ListTransactionsResponseData()
@@ -316,6 +317,7 @@ class TransactionStateManager(brokerId: Int,
           }
         }
 
+        val now : Long = time.milliseconds()
         def shouldInclude(txnMetadata: TransactionMetadata): Boolean = {
           if (txnMetadata.state == Dead) {
             // We filter the `Dead` state since it is a transient state which
@@ -326,6 +328,8 @@ class TransactionStateManager(brokerId: Int,
             false
           } else if (filterStateNames.nonEmpty && !filterStates.contains(txnMetadata.state)) {
             false
+          } else if (filterDurationMs >= 0 && (now - txnMetadata.txnStartTimestamp) <= filterDurationMs) {
+            false
           } else {
             true
           }
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index c1367884272f0..e376b053e44f8 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -3747,7 +3747,8 @@ class KafkaApis(val requestChannel: RequestChannel,
     val listTransactionsRequest = request.body[ListTransactionsRequest]
     val filteredProducerIds = listTransactionsRequest.data.producerIdFilters.asScala.map(Long.unbox).toSet
     val filteredStates = listTransactionsRequest.data.stateFilters.asScala.toSet
-    val response = txnCoordinator.handleListTransactions(filteredProducerIds, filteredStates)
+    val durationFilter = listTransactionsRequest.data.durationFilter()
+    val response = txnCoordinator.handleListTransactions(filteredProducerIds, filteredStates, durationFilter)
 
     // The response should contain only transactionalIds that the principal
     // has `Describe` permission to access.
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
index 94ffe7a979557..907e4c9c17e15 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
@@ -489,7 +489,8 @@ class TransactionStateManagerTest {
     transactionManager.addLoadingPartition(partitionId = 0, coordinatorEpoch = 15)
     val listResponse = transactionManager.listTransactionStates(
       filterProducerIds = Set.empty,
-      filterStateNames = Set.empty
+      filterStateNames = Set.empty,
+      -1L
     )
     assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, Errors.forCode(listResponse.errorCode))
   }
@@ -513,12 +514,16 @@ class TransactionStateManagerTest {
 
     putTransaction(transactionalId = "t0", producerId = 0, state = Ongoing)
     putTransaction(transactionalId = "t1", producerId = 1, state = Ongoing)
+    // update time to create transactions with various durations
+    time.sleep(1000)
     putTransaction(transactionalId = "t2", producerId = 2, state = PrepareCommit)
     putTransaction(transactionalId = "t3", producerId = 3, state = PrepareAbort)
+    time.sleep(1000)
     putTransaction(transactionalId = "t4", producerId = 4, state = CompleteCommit)
     putTransaction(transactionalId = "t5", producerId = 5, state = CompleteAbort)
     putTransaction(transactionalId = "t6", producerId = 6, state = CompleteAbort)
     putTransaction(transactionalId = "t7", producerId = 7, state = PrepareEpochFence)
+    time.sleep(1000)
     // Note that `Dead` transactions are never returned. This is a transient state
     // which is used when the transaction state is in the process of being deleted
     // (whether though expiration or coordinator unloading).
@@ -527,16 +532,20 @@ class TransactionStateManagerTest {
     def assertListTransactions(
       expectedTransactionalIds: Set[String],
       filterProducerIds: Set[Long] = Set.empty,
-      filterStates: Set[String] = Set.empty
+      filterStates: Set[String] = Set.empty,
+      filterDuration: Long = -1L
     ): Unit = {
-      val listResponse = transactionManager.listTransactionStates(filterProducerIds, filterStates)
+      val listResponse = transactionManager.listTransactionStates(filterProducerIds, filterStates, filterDuration)
       assertEquals(Errors.NONE, Errors.forCode(listResponse.errorCode))
       assertEquals(expectedTransactionalIds, listResponse.transactionStates.asScala.map(_.transactionalId).toSet)
       val expectedUnknownStates = filterStates.filter(state => TransactionState.fromName(state).isEmpty)
       assertEquals(expectedUnknownStates, listResponse.unknownStateFilters.asScala.toSet)
     }
-
     assertListTransactions(Set("t0", "t1", "t2", "t3", "t4", "t5", "t6", "t7"))
+    assertListTransactions(Set("t0", "t1", "t2", "t3", "t4", "t5", "t6", "t7"), filterDuration = 0L)
+    assertListTransactions(Set("t0", "t1", "t2", "t3"), filterDuration = 1000L)
+    assertListTransactions(Set("t0", "t1"), filterDuration = 2000L)
+    assertListTransactions(Set(), filterDuration = 3000L)
     assertListTransactions(Set("t0", "t1"), filterStates = Set("Ongoing"))
     assertListTransactions(Set("t0", "t1"), filterStates = Set("Ongoing", "UnknownState"))
     assertListTransactions(Set("t2", "t4"), filterStates = Set("PrepareCommit", "CompleteCommit"))
@@ -843,7 +852,7 @@ class TransactionStateManagerTest {
   }
 
   private def listExpirableTransactionalIds(): Set[String] = {
-    val activeTransactionalIds = transactionManager.listTransactionStates(Set.empty, Set.empty)
+    val activeTransactionalIds = transactionManager.listTransactionStates(Set.empty, Set.empty, -1L)
       .transactionStates
       .asScala
       .map(_.transactionalId)
diff --git a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
index f272bd5a23b6d..90ff37ce32957 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
@@ -6499,7 +6499,7 @@ class KafkaApisTest extends Logging {
     when(clientRequestQuotaManager.maybeRecordAndGetThrottleTimeMs(any[RequestChannel.Request](),
       any[Long])).thenReturn(0)
 
-    when(txnCoordinator.handleListTransactions(Set.empty[Long], Set.empty[String]))
+    when(txnCoordinator.handleListTransactions(Set.empty[Long], Set.empty[String], -1L))
       .thenReturn(new ListTransactionsResponseData()
         .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code))
     kafkaApis = createKafkaApis()
@@ -6529,7 +6529,7 @@ class KafkaApisTest extends Logging {
       .setProducerId(98765)
       .setTransactionState("PrepareAbort"))
 
-    when(txnCoordinator.handleListTransactions(Set.empty[Long], Set.empty[String]))
+    when(txnCoordinator.handleListTransactions(Set.empty[Long], Set.empty[String], -1L))
       .thenReturn(new ListTransactionsResponseData()
         .setErrorCode(Errors.NONE.code)
         .setTransactionStates(transactionStates))
diff --git a/tools/src/main/java/org/apache/kafka/tools/TransactionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/TransactionsCommand.java
index 243f28035f4fd..479db63e18a5c 100644
--- a/tools/src/main/java/org/apache/kafka/tools/TransactionsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/TransactionsCommand.java
@@ -436,16 +436,26 @@ public String name() {
 
         @Override
         public void addSubparser(Subparsers subparsers) {
-            subparsers.addParser(name())
+            Subparser subparser = subparsers.addParser(name())
                 .help("list transactions");
+
+            subparser.addArgument("--duration-filter")
+                    .help("Duration (in millis) to filter by: if < 0, all transactions will be returned; " +
+                            "otherwise, only transactions running longer than this duration will be returned")
+                    .action(store())
+                    .type(Long.class)
+                    .required(false);
         }
 
         @Override
         public void execute(Admin admin, Namespace ns, PrintStream out) throws Exception {
+            ListTransactionsOptions options = new ListTransactionsOptions();
+            Optional.ofNullable(ns.getLong("duration_filter")).ifPresent(options::filterOnDuration);
+
             final Map<Integer, Collection<TransactionListing>> result;
 
             try {
-                result = admin.listTransactions(new ListTransactionsOptions())
+                result = admin.listTransactions(options)
                     .allByBrokerId()
                     .get();
             } catch (ExecutionException e) {
diff --git a/tools/src/test/java/org/apache/kafka/tools/TransactionsCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/TransactionsCommandTest.java
index ef7b5b2ab9930..fde6434cf0266 100644
--- a/tools/src/test/java/org/apache/kafka/tools/TransactionsCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/TransactionsCommandTest.java
@@ -187,14 +187,25 @@ private void testDescribeProducers(
         assertEquals(expectedRows, new HashSet<>(table.subList(1, table.size())));
     }
 
-    @Test
-    public void testListTransactions() throws Exception {
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    public void testListTransactions(boolean hasDurationFilter) throws Exception {
         String[] args = new String[] {
             "--bootstrap-server",
             "localhost:9092",
             "list"
         };
 
+        if (hasDurationFilter) {
+            args = new String[] {
+                "--bootstrap-server",
+                "localhost:9092",
+                "list",
+                "--duration-filter",
+                Long.toString(Long.MAX_VALUE)
+            };
+        }
+
         Map<Integer, Collection<TransactionListing>> transactions = new HashMap<>();
         transactions.put(0, asList(
             new TransactionListing("foo", 12345L, TransactionState.ONGOING),
@@ -204,7 +215,11 @@ public void testListTransactions() throws Exception {
             new TransactionListing("baz", 13579L, TransactionState.COMPLETE_COMMIT)
         ));
 
-        expectListTransactions(transactions);
+        if (hasDurationFilter) {
+            expectListTransactions(new ListTransactionsOptions().filterOnDuration(Long.MAX_VALUE), transactions);
+        } else {
+            expectListTransactions(transactions);
+        }
 
         execute(args);
         assertNormalExit();

From e8cd661becb08279702f1670bb8f1d816537adea Mon Sep 17 00:00:00 2001
From: Anton Agestam <git@antonagestam.se>
Date: Mon, 26 Feb 2024 07:49:59 +0100
Subject: [PATCH 078/258] MINOR: Document MetadataResponse invariants for name
 and ID (#15386)

Reviewers: Justine Olshan <jolshan@confluent.io>, Ziming Deng <dengziming1993@gmail.com>, Greg Harris <gharris1727@gmail.com>.
---
 .../src/main/resources/common/message/MetadataResponse.json  | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/clients/src/main/resources/common/message/MetadataResponse.json b/clients/src/main/resources/common/message/MetadataResponse.json
index 714b28b53d961..bb0b6a802c344 100644
--- a/clients/src/main/resources/common/message/MetadataResponse.json
+++ b/clients/src/main/resources/common/message/MetadataResponse.json
@@ -67,8 +67,9 @@
       { "name": "ErrorCode", "type": "int16", "versions": "0+",
         "about": "The topic error, or 0 if there was no error." },
       { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName", "nullableVersions": "12+",
-        "about": "The topic name." },
-      { "name": "TopicId", "type": "uuid", "versions": "10+", "ignorable": true, "about": "The topic id." },
+        "about": "The topic name. Null for non-existing topics queried by ID. This is never null when ErrorCode is zero. One of Name and TopicId is always populated." },
+      { "name": "TopicId", "type": "uuid", "versions": "10+", "ignorable": true,
+        "about": "The topic id. Zero for non-existing topics queried by name. This is never zero when ErrorCode is zero. One of Name and TopicId is always populated." },
       { "name": "IsInternal", "type": "bool", "versions": "1+", "default": "false", "ignorable": true,
         "about": "True if the topic is internal." },
       { "name": "Partitions", "type": "[]MetadataResponsePartition", "versions": "0+",

From 3e1e0203630640bacf2bf96f3c36b138f6abb667 Mon Sep 17 00:00:00 2001
From: Manikumar Reddy <manikumar.reddy@gmail.com>
Date: Mon, 26 Feb 2024 12:51:33 +0530
Subject: [PATCH 079/258] MINOR: Add validateOAuthMechanismAndNonNullJaasConfig
 to JaasOptionsUtils (#1024)

---
 .../internals/secured/JaasOptionsUtils.java   | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/secured/JaasOptionsUtils.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/secured/JaasOptionsUtils.java
index 83976de73ab65..d04172fe27ced 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/secured/JaasOptionsUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/secured/JaasOptionsUtils.java
@@ -50,13 +50,26 @@ public JaasOptionsUtils(Map<String, Object> options) {
     }
 
     public static Map<String, Object> getOptions(String saslMechanism, List<AppConfigurationEntry> jaasConfigEntries) {
-        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism))
-            throw new IllegalArgumentException(String.format("Unexpected SASL mechanism: %s", saslMechanism));
+        validateOAuthMechanismAndNonNullJaasConfig(saslMechanism, jaasConfigEntries);
+        return Collections.unmodifiableMap(jaasConfigEntries.get(0).getOptions());
+    }
 
-        if (Objects.requireNonNull(jaasConfigEntries).size() != 1 || jaasConfigEntries.get(0) == null)
-            throw new IllegalArgumentException(String.format("Must supply exactly 1 non-null JAAS mechanism configuration (size was %d)", jaasConfigEntries.size()));
+    public static void validateOAuthMechanismAndNonNullJaasConfig(
+            String saslMechanism,
+            List<AppConfigurationEntry> jaasConfigEntries
+    ) {
+        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism)) {
+            throw new IllegalArgumentException(String.format("Unexpected SASL mechanism: %s", saslMechanism));
+        }
 
-        return Collections.unmodifiableMap(jaasConfigEntries.get(0).getOptions());
+        if (Objects.requireNonNull(jaasConfigEntries).size() != 1 || jaasConfigEntries.get(0) == null) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Must supply exactly 1 non-null JAAS mechanism configuration (size was %d)",
+                            jaasConfigEntries.size()
+                    )
+            );
+        }
     }
 
     public boolean shouldCreateSSLSocketFactory(URL url) {

From 9bc9fae9425e4dac64ef078cd3a4e7e6e09cc45a Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Mon, 26 Feb 2024 05:39:33 -0500
Subject: [PATCH 080/258] KAFKA-16258: callback to release assignment when
 stale member leaves group (#15415)

Introduce call to onPartitionsLost callback to release assignment when a consumer pro-actively leaves the group due to poll timer expired.

When the poll timer expires, the member sends a leave group request (reusing same existing LEAVING state and logic), and then transitions to STALE to release it assignment and wait for the poll timer reset. Once both conditions are met, the consumer transitions out of the STALE state to rejoin the group. Note that while on this STALE state, the member is not part of the group so it does not send heartbeats.

This PR also includes the fix to ensure that while STALE or in any other state where the member is not in the group, heartbeat responses that may be received are ignored.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/HeartbeatRequestManager.java    |  15 +-
 .../consumer/internals/MemberState.java       |  17 +-
 .../consumer/internals/MembershipManager.java |  11 +-
 .../internals/MembershipManagerImpl.java      | 112 +++++++--
 .../HeartbeatRequestManagerTest.java          |  16 +-
 .../internals/MembershipManagerImplTest.java  | 217 ++++++++++++++----
 .../kafka/api/PlaintextConsumerTest.scala     |   4 +-
 7 files changed, 299 insertions(+), 93 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 2aabcfe20bc81..550e5b92ebf24 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -200,13 +200,13 @@ public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
                 "messages. You can address this either by increasing max.poll.interval.ms or by " +
                 "reducing the maximum size of batches returned in poll() with max.poll.records.");
 
-            // This should trigger a heartbeat with leave group epoch
-            membershipManager.transitionToStale();
-            NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(currentTimeMs, true);
+            membershipManager.transitionToSendingLeaveGroup(true);
+            NetworkClientDelegate.UnsentRequest leaveHeartbeat = makeHeartbeatRequest(currentTimeMs, true);
+
             // We can ignore the leave response because we can join before or after receiving the response.
             heartbeatRequestState.reset();
             heartbeatState.reset();
-            return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request));
+            return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(leaveHeartbeat));
         }
 
         boolean heartbeatNow = membershipManager.shouldHeartbeatNow() && !heartbeatRequestState.requestInFlight();
@@ -256,11 +256,12 @@ public long maximumTimeToWait(long currentTimeMs) {
      * member to {@link MemberState#JOINING}, so that it rejoins the group.
      */
     public void resetPollTimer(final long pollMs) {
+        if (pollTimer.isExpired()) {
+            logger.debug("Poll timer has been reset after it had expired");
+            membershipManager.maybeRejoinStaleMember();
+        }
         pollTimer.update(pollMs);
         pollTimer.reset(maxPollIntervalMs);
-        if (membershipManager.state() == MemberState.STALE) {
-            membershipManager.transitionToJoining();
-        }
     }
 
     private NetworkClientDelegate.UnsentRequest makeHeartbeatRequest(final long currentTimeMs,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
index 1df4d30e5945d..9f0c7d947ea7e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MemberState.java
@@ -79,11 +79,12 @@ public enum MemberState {
     FENCED,
 
     /**
-     * The member transitions to this state after a call to unsubscribe. While in this state, the
-     * member will stop sending heartbeats, will commit offsets if needed and release its
-     * assignment (calling user's callback for partitions revoked or lost). When all these
-     * actions complete, the member will transition out of this state into {@link #LEAVING} to
-     * effectively leave the group.
+     * The member transitions to this state before sending a heartbeat to leave the group,
+     * While in this state, the member will continue sending heartbeats while it release its
+     * assignment calling the user's callback. When callbacks complete, the member will transition
+     * out of this state into {@link #LEAVING} to send a heartbeat to leave the group. Note that
+     * if leaving due to expired poll timer, the member does not execute any callbacks while in
+     * this state and just transitions to {@link #LEAVING} and then {@link #STALE}
      */
     PREPARE_LEAVING,
 
@@ -130,13 +131,13 @@ public enum MemberState {
         JOINING.previousValidStates = Arrays.asList(FENCED, UNSUBSCRIBED, STALE);
 
         PREPARE_LEAVING.previousValidStates = Arrays.asList(JOINING, STABLE, RECONCILING,
-                ACKNOWLEDGING, UNSUBSCRIBED, FENCED);
+                ACKNOWLEDGING, UNSUBSCRIBED);
 
         LEAVING.previousValidStates = Arrays.asList(PREPARE_LEAVING);
 
-        UNSUBSCRIBED.previousValidStates = Arrays.asList(PREPARE_LEAVING, LEAVING);
+        UNSUBSCRIBED.previousValidStates = Arrays.asList(PREPARE_LEAVING, LEAVING, FENCED);
 
-        STALE.previousValidStates = Arrays.asList(JOINING, RECONCILING, ACKNOWLEDGING, STABLE);
+        STALE.previousValidStates = Arrays.asList(LEAVING);
     }
 
     private List<MemberState> previousValidStates;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
index f95552cc3a353..a9c23d7b4d5da 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
@@ -159,10 +159,9 @@ public interface MembershipManager extends RequestManager {
     void transitionToJoining();
 
     /**
-     * When the user stops polling the consumer and the <code>max.poll.interval.ms</code> timer expires, we transition
-     * the member to STALE.
+     * Transition to the {@link MemberState#LEAVING} state to send a heartbeat to leave the group.
      */
-    void transitionToStale();
+    void transitionToSendingLeaveGroup(boolean dueToPollTimerExpired);
 
     /**
      * Register a listener that will be called whenever the member state changes due to
@@ -175,4 +174,10 @@ public interface MembershipManager extends RequestManager {
      * leaving (sending last heartbeat).
      */
     boolean isLeavingGroup();
+
+    /**
+     * Transition a {@link MemberState#STALE} member to {@link MemberState#JOINING} when it completes
+     * releasing its assignment. This is expected to be used when the poll timer is reset.
+     */
+    void maybeRejoinStaleMember();
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 81e65dfd86612..a9b0f3b94d831 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -257,8 +257,23 @@ public class MembershipManagerImpl implements MembershipManager {
      */
     private final BackgroundEventHandler backgroundEventHandler;
 
+    /**
+     * Future that will complete when a stale member completes releasing its assignment after
+     * leaving the group due to poll timer expired. Used to make sure that the member rejoins
+     * when the timer is reset, only when it completes releasing its assignment.
+     */
+    private CompletableFuture<Void> staleMemberAssignmentRelease;
+
     private final Time time;
 
+    /**
+     * True if the poll timer has expired, signaled by a call to
+     * {@link #transitionToSendingLeaveGroup(boolean)} with dueToExpiredPollTimer param true. This
+     * will be used to determine that the member should transition to STALE after leaving the
+     * group, to release its assignment and wait for the timer to be reset.
+     */
+    private boolean isPollTimerExpired;
+
     public MembershipManagerImpl(String groupId,
                                  Optional<String> groupInstanceId,
                                  int rebalanceTimeoutMs,
@@ -347,11 +362,17 @@ public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData respo
             );
             throw new IllegalArgumentException(errorMessage);
         }
+        MemberState state = state();
         if (state == MemberState.LEAVING) {
             log.debug("Ignoring heartbeat response received from broker. Member {} with epoch {} is " +
                     "already leaving the group.", memberId, memberEpoch);
             return;
         }
+        if (isNotInGroup()) {
+            log.debug("Ignoring heartbeat response received from broker. Member {} is in {} state" +
+                " so it's not a member of the group. ", memberId, state);
+            return;
+        }
 
         // Update the group member id label in the client telemetry reporter if the member id has
         // changed. Initially the member id is empty, and it is updated when the member joins the
@@ -382,6 +403,16 @@ public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData respo
         }
     }
 
+    /**
+     * @return True if the consumer is not a member of the group.
+     */
+    private boolean isNotInGroup() {
+        return state == MemberState.UNSUBSCRIBED ||
+            state == MemberState.FENCED ||
+            state == MemberState.FATAL ||
+            state == MemberState.STALE;
+    }
+
     /**
      * This will process the assignment received if it is different from the member's current
      * assignment. If a new assignment is received, this will make sure reconciliation is attempted
@@ -462,7 +493,12 @@ public void transitionToFenced() {
                         " after member got fenced. Member will rejoin the group anyways.", error);
             }
             updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
-            transitionToJoining();
+            if (state == MemberState.FENCED) {
+                transitionToJoining();
+            } else {
+                log.debug("Fenced member onPartitionsLost callback completed but the state has " +
+                    "already changed to {}, so the member won't rejoin the group", state);
+            }
         });
     }
 
@@ -564,9 +600,11 @@ public void transitionToJoining() {
      */
     @Override
     public CompletableFuture<Void> leaveGroup() {
-        if (state == MemberState.UNSUBSCRIBED || state == MemberState.FATAL) {
-            // Member is not part of the group. No-op and return completed future to avoid
-            // unnecessary transitions.
+        if (isNotInGroup()) {
+            if (state == MemberState.FENCED) {
+                updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+                transitionTo(MemberState.UNSUBSCRIBED);
+            }
             return CompletableFuture.completedFuture(null);
         }
 
@@ -588,7 +626,7 @@ public CompletableFuture<Void> leaveGroup() {
             // Transition to ensure that a heartbeat request is sent out to effectively leave the
             // group (even in the case where the member had no assignment to release or when the
             // callback execution failed.)
-            transitionToSendingLeaveGroup();
+            transitionToSendingLeaveGroup(false);
         });
 
         // Return future to indicate that the leave group is done when the callbacks
@@ -633,11 +671,15 @@ private CompletableFuture<Void> invokeOnPartitionsRevokedOrLostToReleaseAssignme
 
     /**
      * Reset member epoch to the value required for the leave the group heartbeat request, and
-     * transition to the {@link MemberState#LEAVING} state so that a heartbeat
-     * request is sent out with it.
-     * Visible for testing.
+     * transition to the {@link MemberState#LEAVING} state so that a heartbeat request is sent
+     * out with it.
+     *
+     * @param dueToExpiredPollTimer True if the leave group is due to an expired poll timer. This
+     *                              will indicate that the member must remain STALE after leaving,
+     *                              until it releases its assignment and the timer is reset.
      */
-    void transitionToSendingLeaveGroup() {
+    @Override
+    public void transitionToSendingLeaveGroup(boolean dueToExpiredPollTimer) {
         if (state == MemberState.FATAL) {
             log.warn("Member {} with epoch {} won't send leave group request because it is in " +
                     "FATAL state", memberId, memberEpoch);
@@ -648,6 +690,14 @@ void transitionToSendingLeaveGroup() {
                 memberId);
             return;
         }
+
+        if (dueToExpiredPollTimer) {
+            this.isPollTimerExpired = true;
+            // Briefly transition through prepare leaving. The member does not have to release
+            // any assignment before sending the leave group given that is stale. It will invoke
+            // onPartitionsLost after sending the leave group on the STALE state.
+            transitionTo(MemberState.PREPARE_LEAVING);
+        }
         int leaveEpoch = groupInstanceId.isPresent() ?
                 ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH :
                 ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
@@ -690,7 +740,11 @@ public void onHeartbeatRequestSent() {
                 transitionTo(MemberState.RECONCILING);
             }
         } else if (state == MemberState.LEAVING) {
-            transitionToUnsubscribed();
+            if (isPollTimerExpired) {
+                transitionToStale();
+            } else {
+                transitionToUnsubscribed();
+            }
         }
     }
 
@@ -742,16 +796,37 @@ public boolean isLeavingGroup() {
         return state == MemberState.PREPARE_LEAVING || state == MemberState.LEAVING;
     }
 
+    @Override
+    public void maybeRejoinStaleMember() {
+        isPollTimerExpired = false;
+        if (state == MemberState.STALE) {
+            log.debug("Expired poll timer has been reset so stale member {} will rejoin the group" +
+                "when it completes releasing its previous assignment.", memberId);
+            staleMemberAssignmentRelease.whenComplete((__, error) -> transitionToJoining());
+        }
+    }
+
     /**
-     * Sets the epoch to the leave group epoch and clears the assignments. The member will rejoin with
-     * the existing subscriptions after the next application poll event.
+     * Transition to STALE to release assignments because the member has left the group due to
+     * expired poll timer. This will trigger the onPartitionsLost callback. Once the callback
+     * completes, the member will remain stale until the poll timer is reset by an application
+     * poll event. See {@link #maybeRejoinStaleMember()}.
      */
-    @Override
-    public void transitionToStale() {
-        memberEpoch = ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
-        // Clear the current assignment and subscribed partitions before member sending the leave group
-        updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+    private void transitionToStale() {
         transitionTo(MemberState.STALE);
+
+        // Release assignment
+        CompletableFuture<Void> callbackResult = invokeOnPartitionsLostCallback(subscriptions.assignedPartitions());
+        staleMemberAssignmentRelease = callbackResult.whenComplete((result, error) -> {
+            if (error != null) {
+                log.error("onPartitionsLost callback invocation failed while releasing assignment" +
+                    " after member left group due to expired poll timer.", error);
+            }
+            updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+            log.debug("Member {} sent leave group heartbeat and released its assignment. It will remain " +
+                "in {} state until the poll timer is reset, and it will then rejoin the group",
+                memberId, MemberState.STALE);
+        });
     }
 
     /**
@@ -1184,8 +1259,7 @@ private CompletableFuture<Void> invokeOnPartitionsAssignedCallback(Set<TopicPart
         }
     }
 
-    // Visible for testing
-    CompletableFuture<Void> invokeOnPartitionsLostCallback(Set<TopicPartition> partitionsLost) {
+    private CompletableFuture<Void> invokeOnPartitionsLostCallback(Set<TopicPartition> partitionsLost) {
         // This should not trigger the callback if partitionsLost is empty, to keep the current
         // behaviour.
         Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 72a5c0349d274..5cf5b9e2d92c5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -71,6 +71,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
@@ -575,20 +576,21 @@ public void testPollTimerExpiration() {
         when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(new Node(1, "localhost", 9999)));
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
 
-        // On poll timer expiration, the member should transition to stale and a last heartbeat 
-        // should be sent to leave the group
+        // On poll timer expiration, the member should send a last heartbeat to leave the group
+        // and notify the membership manager
         time.sleep(maxPollIntervalMs);
         assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
+        verify(membershipManager).transitionToSendingLeaveGroup(true);
         verify(heartbeatState).reset();
         verify(heartbeatRequestState).reset();
-        verify(membershipManager).transitionToStale();
+        verify(membershipManager).onHeartbeatRequestSent();
 
         when(membershipManager.state()).thenReturn(MemberState.STALE);
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(true);
         assertNoHeartbeat(heartbeatRequestManager);
         heartbeatRequestManager.resetPollTimer(time.milliseconds());
         assertTrue(pollTimer.notExpired());
-        verify(membershipManager).transitionToJoining();
+        verify(membershipManager).maybeRejoinStaleMember();
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
         assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
     }
@@ -604,13 +606,13 @@ public void testPollTimerExpiration() {
     public void testPollTimerExpirationShouldNotMarkMemberStaleIfMemberAlreadyLeaving() {
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
         when(membershipManager.isLeavingGroup()).thenReturn(true);
-        doNothing().when(membershipManager).transitionToStale();
 
         time.sleep(maxPollIntervalMs);
         NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
 
-        // No transition to STALE should be triggered, because the member is already leaving the group
-        verify(membershipManager, never()).transitionToStale();
+        // No transition to leave due to stale member should be triggered, because the member is
+        // already leaving the group
+        verify(membershipManager, never()).transitionToSendingLeaveGroup(anyBoolean());
 
         assertEquals(1, result.unsentRequests.size(), "A heartbeat request should be generated to" +
             " complete the ongoing leaving operation that was triggered before the poll timer expired.");
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 50f28bb523367..3294068b07487 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -36,6 +36,9 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -52,6 +55,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer.invokeRebalanceCallbacks;
 import static org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
@@ -59,6 +63,7 @@
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.apache.kafka.common.utils.Utils.mkSortedSet;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
@@ -117,23 +122,22 @@ public void tearDown() {
     }
 
     private MembershipManagerImpl createMembershipManagerJoiningGroup() {
-        MembershipManagerImpl manager = spy(new MembershipManagerImpl(
-                GROUP_ID, Optional.empty(), REBALANCE_TIMEOUT, Optional.empty(),
-                subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
-                backgroundEventHandler, time));
-        manager.transitionToJoining();
-        return manager;
+        return createMembershipManagerJoiningGroup(null);
     }
 
     private MembershipManagerImpl createMembershipManagerJoiningGroup(String groupInstanceId) {
-        MembershipManagerImpl manager = spy(new MembershipManagerImpl(
-                GROUP_ID, Optional.ofNullable(groupInstanceId), REBALANCE_TIMEOUT, Optional.empty(),
-                subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
-                backgroundEventHandler, time));
+        MembershipManagerImpl manager = createMembershipManager(groupInstanceId);
         manager.transitionToJoining();
         return manager;
     }
 
+    private MembershipManagerImpl createMembershipManager(String groupInstanceId) {
+        return spy(new MembershipManagerImpl(
+            GROUP_ID, Optional.ofNullable(groupInstanceId), REBALANCE_TIMEOUT, Optional.empty(),
+            subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
+            backgroundEventHandler, time));
+    }
+
     private MembershipManagerImpl createMembershipManagerJoiningGroup(String groupInstanceId,
                                                                       String serverAssignor) {
         MembershipManagerImpl manager = new MembershipManagerImpl(
@@ -715,6 +719,21 @@ public void testIgnoreHeartbeatWhenLeavingGroup() {
             "heartbeat request to leave is sent out.");
     }
 
+    @ParameterizedTest
+    @MethodSource("notInGroupStates")
+    public void testIgnoreHeartbeatResponseWhenNotInGroup(MemberState state) {
+        MembershipManagerImpl membershipManager = createMembershipManager(null);
+        when(membershipManager.state()).thenReturn(state);
+        ConsumerGroupHeartbeatResponseData responseData = mock(ConsumerGroupHeartbeatResponseData.class);
+
+        membershipManager.onHeartbeatResponseReceived(responseData);
+
+        assertEquals(state, membershipManager.state());
+        verify(responseData, never()).memberId();
+        verify(responseData, never()).memberEpoch();
+        verify(responseData, never()).assignment();
+    }
+
     @Test
     public void testLeaveGroupWhenStateIsReconciling() {
         MembershipManager membershipManager = mockJoinAndReceiveAssignment(false);
@@ -803,6 +822,33 @@ public void testLeaveGroupWhenMemberAlreadyLeft() {
         verify(subscriptionState, never()).assignFromSubscribed(Collections.emptySet());
     }
 
+    @Test
+    public void testLeaveGroupWhenMemberFenced() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+        ConsumerRebalanceListenerCallbackCompletedEvent callbackEvent = mockFencedMemberStuckOnUserCallback(membershipManager, invoker);
+        assertEquals(MemberState.FENCED, membershipManager.state());
+
+        mockLeaveGroup();
+        membershipManager.leaveGroup();
+        assertEquals(MemberState.UNSUBSCRIBED, membershipManager.state());
+        verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
+
+        completeCallback(callbackEvent, membershipManager);
+        assertEquals(MemberState.UNSUBSCRIBED, membershipManager.state());
+    }
+
+    @Test
+    public void testLeaveGroupWhenMemberIsStale() {
+        MembershipManagerImpl membershipManager = mockStaleMember();
+        assertEquals(MemberState.STALE, membershipManager.state());
+
+        mockLeaveGroup();
+        CompletableFuture<Void> leaveResult1 = membershipManager.leaveGroup();
+        assertTrue(leaveResult1.isDone());
+        assertEquals(MemberState.STALE, membershipManager.state());
+    }
+
     @Test
     public void testFatalFailureWhenStateIsUnjoined() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
@@ -1612,61 +1658,118 @@ public void testOnPartitionsLostError() {
         testOnPartitionsLost(Optional.of(new KafkaException("Intentional error for test")));
     }
 
+    private void assertLeaveGroupDueToExpiredPollAndTransitionToStale(MembershipManagerImpl membershipManager) {
+        assertDoesNotThrow(() -> membershipManager.transitionToSendingLeaveGroup(true));
+        assertEquals(LEAVE_GROUP_MEMBER_EPOCH, membershipManager.memberEpoch());
+        membershipManager.onHeartbeatRequestSent();
+        assertStaleMemberLeavesGroupAndClearsAssignment(membershipManager);
+    }
+
     @Test
-    public void testTransitionToStaleWhileReconciling() {
+    public void testTransitionToLeavingWhileReconcilingDueToStaleMember() {
         MembershipManagerImpl membershipManager = memberJoinWithAssignment();
         clearInvocations(subscriptionState);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
-
-        membershipManager.transitionToStale();
-        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+        assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
 
     @Test
-    public void testTransitionToStaleWhileJoining() {
+    public void testTransitionToLeavingWhileJoiningDueToStaleMember() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         assertEquals(MemberState.JOINING, membershipManager.state());
-
-        membershipManager.transitionToStale();
-        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+        assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
 
     @Test
-    public void testTransitionToStaleWhileStable() {
+    public void testTransitionToLeavingWhileStableDueToStaleMember() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
         membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         assertEquals(MemberState.STABLE, membershipManager.state());
-
-        membershipManager.transitionToStale();
-        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+        assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
 
     @Test
-    public void testTransitionToStaleWhileAcknowledging() {
+    public void testTransitionToLeavingWhileAcknowledgingDueToStaleMember() {
         MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true);
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         clearInvocations(subscriptionState);
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
-
-        membershipManager.transitionToStale();
-        assertStaleMemberClearsAssignmentsAndLeaves(membershipManager);
+        assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
 
     @Test
     public void testStaleMemberDoesNotSendHeartbeatAndAllowsTransitionToJoiningToRecover() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         doNothing().when(subscriptionState).assignFromSubscribed(any());
-        membershipManager.transitionToStale();
-        assertTrue(membershipManager.shouldSkipHeartbeat(), "Stale member should not send " +
-            "heartbeats");
+        membershipManager.transitionToSendingLeaveGroup(true);
+        membershipManager.onHeartbeatRequestSent();
+        assertEquals(MemberState.STALE, membershipManager.state());
+        assertTrue(membershipManager.shouldSkipHeartbeat(), "Stale member should not send heartbeats");
         // Check that a transition to joining is allowed, which is what is expected to happen
         // when the poll timer is reset.
-        membershipManager.transitionToJoining();
+        assertDoesNotThrow(membershipManager::maybeRejoinStaleMember);
+    }
+
+    @Test
+    public void testStaleMemberRejoinsWhenTimerResetsNoCallbacks() {
+        MembershipManagerImpl membershipManager = mockStaleMember();
+        assertStaleMemberLeavesGroupAndClearsAssignment(membershipManager);
+
+        membershipManager.maybeRejoinStaleMember();
+        assertEquals(MemberState.JOINING, membershipManager.state());
     }
 
+    @Test
+    public void testStaleMemberWaitsForCallbackToRejoinWhenTimerReset() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        Uuid topicId = Uuid.randomUuid();
+        String topicName = "topic1";
+        int ownedPartition = 0;
+        TopicPartition tp = new TopicPartition(topicName, ownedPartition);
+        mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId, topicName,
+            Collections.singletonList(new TopicIdPartition(topicId, tp)));
+        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+
+        membershipManager.transitionToSendingLeaveGroup(true);
+        membershipManager.onHeartbeatRequestSent();
+
+        assertEquals(MemberState.STALE, membershipManager.state());
+        verify(backgroundEventHandler).add(any(ConsumerRebalanceListenerCallbackNeededEvent.class));
+
+        // Stale member triggers onPartitionLost callback that will not complete just yet
+        ConsumerRebalanceListenerCallbackCompletedEvent callbackEvent = performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST,
+            topicPartitions(topicName, ownedPartition),
+            false
+        );
+
+        // Timer reset while callback hasn't completed. Member should stay in STALE while it
+        // completes releasing its assignment, and then transition to joining.
+        membershipManager.maybeRejoinStaleMember();
+        assertEquals(MemberState.STALE, membershipManager.state(), "Member should not transition " +
+            "out of the STALE state when the timer is reset if the callback has not completed.");
+        // Member should not clear its assignment to rejoin until the callback completes
+        verify(subscriptionState, never()).assignFromSubscribed(any());
+
+        completeCallback(callbackEvent, membershipManager);
+        assertEquals(MemberState.JOINING, membershipManager.state());
+        verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
+    }
+
+    private MembershipManagerImpl mockStaleMember() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        doNothing().when(subscriptionState).assignFromSubscribed(any());
+        membershipManager.transitionToSendingLeaveGroup(true);
+        membershipManager.onHeartbeatRequestSent();
+        return membershipManager;
+    }
     private void mockPartitionOwnedAndNewPartitionAdded(String topicName,
                                                         int partitionOwned,
                                                         int partitionAdded,
@@ -1814,30 +1917,17 @@ private void testFenceIsNoOp(MembershipManagerImpl membershipManager) {
         verify(subscriptionState, never()).rebalanceListener();
     }
 
-    private void assertStaleMemberClearsAssignmentsAndLeaves(MembershipManagerImpl membershipManager) {
+    private void assertStaleMemberLeavesGroupAndClearsAssignment(MembershipManagerImpl membershipManager) {
         assertEquals(MemberState.STALE, membershipManager.state());
 
-        // Should clear subscriptions, current assignments, and reset epoch to leave the group
+        // Should reset epoch to leave the group and release the assignment (right away because
+        // there is no onPartitionsLost callback defined)
         verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
         assertTrue(membershipManager.currentAssignment().isEmpty());
         assertTrue(membershipManager.topicsAwaitingReconciliation().isEmpty());
         assertEquals(LEAVE_GROUP_MEMBER_EPOCH, membershipManager.memberEpoch());
     }
 
-    @Test
-    public void testHeartbeatSentOnStaleMember() {
-        MembershipManagerImpl membershipManager = createMemberInStableState();
-        subscriptionState.subscribe(Collections.singleton("topic"), Optional.empty());
-        subscriptionState.assignFromSubscribed(Collections.singleton(new TopicPartition("topic", 0)));
-        membershipManager.transitionToStale();
-        membershipManager.onHeartbeatRequestSent();
-        // Member should remain in STALE state. Only when the poll timer is reset the member will
-        // transition to JOINING.
-        assertEquals(MemberState.STALE, membershipManager.state());
-        assertTrue(membershipManager.currentAssignment().isEmpty());
-        assertTrue(subscriptionState.assignedPartitions().isEmpty());
-    }
-
     @Test
     public void testMemberJoiningTransitionsToStableWhenReceivingEmptyAssignment() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(null);
@@ -2160,6 +2250,29 @@ private ConsumerRebalanceListenerCallbackCompletedEvent mockPrepareLeavingStuckO
         );
     }
 
+    private ConsumerRebalanceListenerCallbackCompletedEvent mockFencedMemberStuckOnUserCallback(
+        MembershipManagerImpl membershipManager,
+        ConsumerRebalanceListenerInvoker invoker) {
+        String topicName = "topic1";
+        TopicPartition ownedPartition = new TopicPartition(topicName, 0);
+
+        // Fence member and block waiting for onPartitionsLost callback to complete
+        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
+        when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(ownedPartition));
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+        // doNothing().when(subscriptionState).markPendingRevocation(anySet());
+        when(commitRequestManager.autoCommitEnabled()).thenReturn(false);
+        membershipManager.transitionToFenced();
+        return performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST,
+            topicPartitions(ownedPartition.topic(), ownedPartition.partition()),
+            false
+        );
+    }
+
     private void testStateUpdateOnFatalFailure(MembershipManagerImpl membershipManager) {
         String memberId = membershipManager.memberId();
         int lastEpoch = membershipManager.memberEpoch();
@@ -2231,4 +2344,16 @@ private MembershipManagerImpl memberJoinWithAssignment() {
         assertFalse(membershipManager.currentAssignment().isEmpty());
         return membershipManager;
     }
+
+    /**
+     * @return States where the member is not part of the group.
+     */
+    private static Stream<Arguments> notInGroupStates() {
+        return Stream.of(
+            Arguments.of(MemberState.UNSUBSCRIBED),
+            Arguments.of(MemberState.FENCED),
+            Arguments.of(MemberState.FATAL),
+            Arguments.of(MemberState.STALE));
+    }
+
 }
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index c826b8fd0c437..01dbb206a07cb 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -169,10 +169,8 @@ class PlaintextConsumerTest extends BaseConsumerTest {
       startingTimestamp = startingTimestamp)
   }
 
-  // TODO: Enable this test for both protocols when the Jira tracking its failure (KAFKA-16008) is fixed. This
-  //       is done by setting the @MethodSource value to "getTestQuorumAndGroupProtocolParametersAll"
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testMaxPollIntervalMs(quorum: String, groupProtocol: String): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)

From 027fad4b2a0d173df075f82e93501afb645ecc12 Mon Sep 17 00:00:00 2001
From: Cameron Redpath <80655142+credpath-seek@users.noreply.github.com>
Date: Tue, 27 Feb 2024 08:33:16 +1100
Subject: [PATCH 081/258] KAFKA-16277: AbstractStickyAssignor - Sort owned
 TopicPartitions by partition when reassigning (#15416)

Treats KAFKA-16277 - CooperativeStickyAssignor does not spread topics evenly among consumer group

Reviewers: Anna Sophie Blee-Goldman <ableegoldman@apache.org>
---
 .../internals/AbstractStickyAssignor.java     |  1 +
 .../internals/AbstractStickyAssignorTest.java | 65 +++++++++++++++++--
 2 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
index ee48bed6df0ec..727dac231acf9 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
@@ -663,6 +663,7 @@ private void assignOwnedPartitions() {
                 String consumer = consumerEntry.getKey();
                 List<TopicPartition> ownedPartitions = consumerEntry.getValue().stream()
                         .filter(tp -> !rackInfo.racksMismatch(consumer, tp))
+                        .sorted(Comparator.comparing(TopicPartition::partition).thenComparing(TopicPartition::topic))
                         .collect(Collectors.toList());
 
                 List<TopicPartition> consumerAssignment = assignment.get(consumer);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
index 71d188f3cd769..1b771929c4d18 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
@@ -414,6 +414,59 @@ public void testAddRemoveConsumerOneTopic(RackConfig rackConfig) {
         assertTrue(isFullyBalanced(assignment));
     }
 
+    @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
+    @EnumSource(RackConfig.class)
+    public void testTopicBalanceAfterReassignment(RackConfig rackConfig) {
+        initializeRacks(rackConfig);
+        List<String> allTopics = topics(topic1, topic2);
+        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, partitionInfos(topic1, 12));
+        partitionsPerTopic.put(topic2, partitionInfos(topic2, 12));
+        subscriptions.put(consumer1, subscription(allTopics, 0));
+        Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assignment.forEach((consumer, tps) -> assertEquals(12, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
+        assignment.forEach((consumer, tps) -> assertEquals(12, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
+        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
+        assertTrue(isFullyBalanced(assignment));
+
+        // Add another consumer
+        subscriptions.put(consumer1, buildSubscriptionV2Above(allTopics, assignment.get(consumer1), generationId, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 1));
+        assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
+        assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
+        assertTrue(isFullyBalanced(assignment));
+
+        // Add two more consumers
+        subscriptions.put(consumer1, buildSubscriptionV2Above(allTopics, assignment.get(consumer1), generationId, 0));
+        subscriptions.put(consumer2, buildSubscriptionV2Above(allTopics, assignment.get(consumer2), generationId, 1));
+        subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 2));
+        subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 3));
+        assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assignment.forEach((consumer, tps) -> assertEquals(3, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
+        assignment.forEach((consumer, tps) -> assertEquals(3, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
+        assertTrue(isFullyBalanced(assignment));
+
+        // remove 2 consumers
+        subscriptions.remove(consumer1);
+        subscriptions.remove(consumer2);
+        subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, assignment.get(consumer3), generationId, 2));
+        subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, assignment.get(consumer4), generationId, 3));
+        assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
+
+        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
+        assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
+        assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
+        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
+        assertTrue(isFullyBalanced(assignment));
+    }
+
     @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
     @EnumSource(RackConfig.class)
     public void testAddRemoveTwoConsumersTwoTopics(RackConfig rackConfig) {
@@ -442,15 +495,15 @@ public void testAddRemoveTwoConsumersTwoTopics(RackConfig rackConfig) {
         assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
 
         Map<TopicPartition, String> expectedPartitionsTransferringOwnership = new HashMap<>();
-        expectedPartitionsTransferringOwnership.put(tp(topic2, 1), consumer3);
+        expectedPartitionsTransferringOwnership.put(tp(topic1, 2), consumer3);
         expectedPartitionsTransferringOwnership.put(tp(topic2, 3), consumer3);
         expectedPartitionsTransferringOwnership.put(tp(topic2, 2), consumer4);
         assertEquals(expectedPartitionsTransferringOwnership, assignor.partitionsTransferringOwnership);
 
         verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
-        assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
-        assertEquals(partitions(tp(topic1, 1), tp(topic2, 0)), assignment.get(consumer2));
-        assertEquals(partitions(tp(topic2, 1), tp(topic2, 3)), assignment.get(consumer3));
+        assertEquals(partitions(tp(topic1, 0), tp(topic2, 1)), assignment.get(consumer1));
+        assertEquals(partitions(tp(topic2, 0), tp(topic1, 1)), assignment.get(consumer2));
+        assertEquals(partitions(tp(topic1, 2), tp(topic2, 3)), assignment.get(consumer3));
         assertEquals(partitions(tp(topic2, 2)), assignment.get(consumer4));
         assertTrue(isFullyBalanced(assignment));
 
@@ -460,8 +513,8 @@ public void testAddRemoveTwoConsumersTwoTopics(RackConfig rackConfig) {
         subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, assignment.get(consumer3), generationId, 2));
         subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, assignment.get(consumer4), generationId, 3));
         assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
-        assertEquals(partitions(tp(topic2, 1), tp(topic2, 3), tp(topic1, 0), tp(topic2, 0)), assignment.get(consumer3));
-        assertEquals(partitions(tp(topic2, 2), tp(topic1, 1), tp(topic1, 2)), assignment.get(consumer4));
+        assertEquals(partitions(tp(topic1, 2), tp(topic2, 3), tp(topic1, 0), tp(topic2, 1)), assignment.get(consumer3));
+        assertEquals(partitions(tp(topic2, 2), tp(topic1, 1), tp(topic2, 0)), assignment.get(consumer4));
 
         assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
 

From ddfcc333f8d1c2f4ab0a88abf6d664153d6e82fd Mon Sep 17 00:00:00 2001
From: Mayank Shekhar Narula <42991652+msn-tldr@users.noreply.github.com>
Date: Mon, 26 Feb 2024 23:57:11 +0000
Subject: [PATCH 082/258] KAFKA-16226 Add test for concurrently
 updatingMetadata and fetching snapshot/cluster (#15385)

Add test for concurrently updatingMetadata and fetching snapshot/cluster

Reviewers: Jason Gustafson <jason@confluent.io>

Co-authored-by: Zhifeng Chen <ericzhifengchen@gmail.com>
---
 .../apache/kafka/clients/MetadataTest.java    | 106 ++++++++++++++++++
 1 file changed, 106 insertions(+)

diff --git a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
index 600fc23ecb988..e6db9685eb52e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
@@ -16,6 +16,11 @@
  */
 package org.apache.kafka.clients;
 
+import java.util.OptionalInt;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.ClusterResourceListener;
 import org.apache.kafka.common.Node;
@@ -1260,6 +1265,107 @@ else if (partition.equals(internalPart))
         Mockito.reset(mockListener);
     }
 
+    /**
+     * Test that concurrently updating Metadata, and fetching the corresponding MetadataSnapshot & Cluster work as expected, i.e.
+     * snapshot & cluster contain the relevant updates.
+     */
+    @Test
+    public void testConcurrentUpdateAndFetchForSnapshotAndCluster() throws InterruptedException {
+        Time time = new MockTime();
+        metadata = new Metadata(refreshBackoffMs, refreshBackoffMaxMs, metadataExpireMs, new LogContext(), new ClusterResourceListeners());
+
+        // Setup metadata with 10 nodes, 2 topics, topic1 & 2, both to be retained in the update. Both will have leader-epoch 100.
+        int oldNodeCount = 10;
+        String topic1 = "test_topic1";
+        String topic2 = "test_topic2";
+        TopicPartition topic1Part0 = new TopicPartition(topic1, 0);
+        Map<String, Integer> topicPartitionCounts = new HashMap<>();
+        int oldPartitionCount = 1;
+        topicPartitionCounts.put(topic1, oldPartitionCount);
+        topicPartitionCounts.put(topic2, oldPartitionCount);
+        Map<String, Uuid> topicIds = new HashMap<>();
+        topicIds.put(topic1, Uuid.randomUuid());
+        topicIds.put(topic2, Uuid.randomUuid());
+        int oldLeaderEpoch = 100;
+        MetadataResponse metadataResponse =
+            RequestTestUtils.metadataUpdateWithIds("cluster", oldNodeCount, Collections.emptyMap(), topicPartitionCounts, _tp -> oldLeaderEpoch, topicIds);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, true, time.milliseconds());
+        MetadataSnapshot snapshot = metadata.fetchMetadataSnapshot();
+        Cluster cluster = metadata.fetch();
+        // Validate metadata snapshot & cluster are setup as expected.
+        assertEquals(cluster, snapshot.cluster());
+        assertEquals(oldNodeCount, snapshot.cluster().nodes().size());
+        assertEquals(oldPartitionCount, snapshot.cluster().partitionCountForTopic(topic1));
+        assertEquals(oldPartitionCount, snapshot.cluster().partitionCountForTopic(topic2));
+        assertEquals(OptionalInt.of(oldLeaderEpoch), snapshot.leaderEpochFor(topic1Part0));
+
+        // Setup 6 threads, where 3 are updating metadata & 3 are reading snapshot/cluster.
+        // Metadata will be updated with higher # of nodes, partition-counts, leader-epoch.
+        int numThreads = 6;
+        ExecutorService service = Executors.newFixedThreadPool(numThreads);
+        CountDownLatch allThreadsDoneLatch = new CountDownLatch(numThreads);
+        CountDownLatch atleastMetadataUpdatedOnceLatch = new CountDownLatch(1);
+        AtomicReference<MetadataSnapshot> newSnapshot = new AtomicReference<>();
+        AtomicReference<Cluster> newCluster = new AtomicReference<>();
+        for (int i = 0; i < numThreads; i++) {
+            final int id = i + 1;
+            service.execute(() -> {
+                if (id % 2 == 0) { // Thread to update metadata.
+                    String oldClusterId = "clusterId";
+                    int nNodes = oldNodeCount + id;
+                    Map<String, Integer> newTopicPartitionCounts = new HashMap<>();
+                    newTopicPartitionCounts.put(topic1, oldPartitionCount + id);
+                    newTopicPartitionCounts.put(topic2, oldPartitionCount + id);
+                    MetadataResponse newMetadataResponse =
+                        RequestTestUtils.metadataUpdateWithIds(oldClusterId, nNodes, Collections.emptyMap(), newTopicPartitionCounts, _tp -> oldLeaderEpoch + id, topicIds);
+                    metadata.updateWithCurrentRequestVersion(newMetadataResponse, true, time.milliseconds());
+                    atleastMetadataUpdatedOnceLatch.countDown();
+                } else { // Thread to read metadata snapshot, once its updated
+                    try {
+                        if (!atleastMetadataUpdatedOnceLatch.await(5, TimeUnit.MINUTES)) {
+                            assertFalse(true, "Test had to wait more than 5 minutes, something went wrong.");
+                        }
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                    newSnapshot.set(metadata.fetchMetadataSnapshot());
+                    newCluster.set(metadata.fetch());
+                }
+                allThreadsDoneLatch.countDown();
+            });
+        }
+        if (!allThreadsDoneLatch.await(5, TimeUnit.MINUTES)) {
+            assertFalse(true, "Test had to wait more than 5 minutes, something went wrong.");
+        }
+
+        // Validate new snapshot is upto-date. And has higher partition counts, nodes & leader epoch than earlier.
+        {
+            int newNodeCount = newSnapshot.get().cluster().nodes().size();
+            assertTrue(oldNodeCount < newNodeCount, "Unexpected value " + newNodeCount);
+            int newPartitionCountTopic1 = newSnapshot.get().cluster().partitionCountForTopic(topic1);
+            assertTrue(oldPartitionCount < newPartitionCountTopic1, "Unexpected value " + newPartitionCountTopic1);
+            int newPartitionCountTopic2 = newSnapshot.get().cluster().partitionCountForTopic(topic2);
+            assertTrue(oldPartitionCount < newPartitionCountTopic2, "Unexpected value " + newPartitionCountTopic2);
+            int newLeaderEpoch = newSnapshot.get().leaderEpochFor(topic1Part0).getAsInt();
+            assertTrue(oldLeaderEpoch < newLeaderEpoch, "Unexpected value " + newLeaderEpoch);
+        }
+
+        // Validate new cluster is upto-date. And has higher partition counts, nodes than earlier.
+        {
+            int newNodeCount = newCluster.get().nodes().size();
+            assertTrue(oldNodeCount < newNodeCount, "Unexpected value " + newNodeCount);
+            int newPartitionCountTopic1 = newCluster.get().partitionCountForTopic(topic1);
+            assertTrue(oldPartitionCount < newPartitionCountTopic1, "Unexpected value " + newPartitionCountTopic1);
+            int newPartitionCountTopic2 = newCluster.get()
+                .partitionCountForTopic(topic2);
+            assertTrue(oldPartitionCount < newPartitionCountTopic2, "Unexpected value " + newPartitionCountTopic2);
+        }
+
+        service.shutdown();
+        // Executor service should down much quickly, as all tasks are finished at this point.
+        assertTrue(service.awaitTermination(60, TimeUnit.SECONDS));
+    }
+
     /**
      * For testUpdatePartially, validates that updatedMetadata is matching expected part1Metadata, part2Metadata, interalPartMetadata, nodes & more.
      */

From 1c9f360f4af00c8cfcc18df3a10fb66ae1cf6b8e Mon Sep 17 00:00:00 2001
From: Almog Gavra <almog.gavra@gmail.com>
Date: Mon, 26 Feb 2024 18:52:25 -0800
Subject: [PATCH 083/258] KAFKA-15215: migrate StreamedJoinTest to Mockito
 (#15424)

Migrate StreamedJoinTest to Mockito

Reviewers: Anna Sophie Blee-Goldman <ableegoldman@apache.org>,
Divij Vaidya <diviv@amazon.com>
---
 build.gradle                                   |  2 ++
 .../scala/kstream/StreamJoinedTest.scala       | 18 ++++++++++--------
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/build.gradle b/build.gradle
index e715144d54fb1..b362c2f673a5b 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2333,6 +2333,8 @@ project(':streams:streams-scala') {
 
     testImplementation libs.junitJupiter
     testImplementation libs.easymock
+    testImplementation libs.mockitoCore
+    testImplementation libs.mockitoJunitJupiter // supports MockitoExtension
     testImplementation libs.hamcrest
     testRuntimeOnly libs.slf4jlog4j
   }
diff --git a/streams/streams-scala/src/test/scala/org/apache/kafka/streams/scala/kstream/StreamJoinedTest.scala b/streams/streams-scala/src/test/scala/org/apache/kafka/streams/scala/kstream/StreamJoinedTest.scala
index 206d77d86b93a..7a5a48ccb91fc 100644
--- a/streams/streams-scala/src/test/scala/org/apache/kafka/streams/scala/kstream/StreamJoinedTest.scala
+++ b/streams/streams-scala/src/test/scala/org/apache/kafka/streams/scala/kstream/StreamJoinedTest.scala
@@ -21,24 +21,26 @@ import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder
 import org.apache.kafka.streams.scala.serialization.Serdes
 import org.apache.kafka.streams.scala.serialization.Serdes._
 import org.apache.kafka.streams.state.Stores
-import org.easymock.EasyMock
-import org.easymock.EasyMock.{createMock, replay}
 import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.{BeforeEach, Test}
+import org.mockito.Mockito.{mock, when}
+import org.mockito.junit.jupiter.{MockitoExtension, MockitoSettings}
+import org.mockito.quality.Strictness
 
 import java.time.Duration
 
+@ExtendWith(Array(classOf[MockitoExtension]))
+@MockitoSettings(strictness = Strictness.STRICT_STUBS)
 class StreamJoinedTest {
 
-  val builder: InternalStreamsBuilder = createMock(classOf[InternalStreamsBuilder])
-  val topoBuilder: InternalTopologyBuilder = createMock(classOf[InternalTopologyBuilder])
+  val builder: InternalStreamsBuilder = mock(classOf[InternalStreamsBuilder])
+  val topoBuilder: InternalTopologyBuilder = mock(classOf[InternalTopologyBuilder])
 
   @BeforeEach
   def before(): Unit = {
-    EasyMock.expect(builder.internalTopologyBuilder()).andReturn(topoBuilder);
-    EasyMock.expect(topoBuilder.topologyConfigs()).andReturn(null)
-    replay(topoBuilder)
-    replay(builder)
+    when(builder.internalTopologyBuilder()).thenReturn(topoBuilder)
+    when(topoBuilder.topologyConfigs()).thenReturn(null)
   }
 
   @Test

From 3075a3f144c0583f570a888294087ef65ed379c9 Mon Sep 17 00:00:00 2001
From: Stanislav Kozlovski <stanislav_kozlovski@outlook.com>
Date: Tue, 27 Feb 2024 09:58:48 +0100
Subject: [PATCH 084/258] MINOR: Add 3.7.0 to system tests (#15436)

As per the release instructions, this patch bumps the versions in the associated files here as part of the 3.7 release
---
 gradle/dependencies.gradle | 2 ++
 tests/docker/Dockerfile    | 2 ++
 vagrant/base.sh            | 2 ++
 3 files changed, 6 insertions(+)

diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 1194e75d60fbc..9cd1515a79c0c 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -142,6 +142,7 @@ versions += [
   kafka_34: "3.4.1",
   kafka_35: "3.5.2",
   kafka_36: "3.6.1",
+  kafka_37: "3.7.0",
   lz4: "1.8.0",
   mavenArtifact: "3.9.6",
   metrics: "2.2.0",
@@ -234,6 +235,7 @@ libs += [
   kafkaStreams_34: "org.apache.kafka:kafka-streams:$versions.kafka_34",
   kafkaStreams_35: "org.apache.kafka:kafka-streams:$versions.kafka_35",
   kafkaStreams_36: "org.apache.kafka:kafka-streams:$versions.kafka_36",
+  kafkaStreams_37: "org.apache.kafka:kafka-streams:$versions.kafka_37",
   log4j: "ch.qos.reload4j:reload4j:$versions.reload4j",
   lz4: "org.lz4:lz4-java:$versions.lz4",
   metrics: "com.yammer.metrics:metrics-core:$versions.metrics",
diff --git a/tests/docker/Dockerfile b/tests/docker/Dockerfile
index 60da7e6b8b126..5701b21730fff 100644
--- a/tests/docker/Dockerfile
+++ b/tests/docker/Dockerfile
@@ -70,6 +70,7 @@ RUN mkdir -p "/opt/kafka-3.3.2" && chmod a+rw /opt/kafka-3.3.2 && curl -s "$KAFK
 RUN mkdir -p "/opt/kafka-3.4.1" && chmod a+rw /opt/kafka-3.4.1 && curl -s "$KAFKA_MIRROR/kafka_2.12-3.4.1.tgz" | tar xz --strip-components=1 -C "/opt/kafka-3.4.1"
 RUN mkdir -p "/opt/kafka-3.5.2" && chmod a+rw /opt/kafka-3.5.2 && curl -s "$KAFKA_MIRROR/kafka_2.12-3.5.2.tgz" | tar xz --strip-components=1 -C "/opt/kafka-3.5.2"
 RUN mkdir -p "/opt/kafka-3.6.1" && chmod a+rw /opt/kafka-3.6.1 && curl -s "$KAFKA_MIRROR/kafka_2.12-3.6.1.tgz" | tar xz --strip-components=1 -C "/opt/kafka-3.6.1"
+RUN mkdir -p "/opt/kafka-3.7.0" && chmod a+rw /opt/kafka-3.7.0 && curl -s "$KAFKA_MIRROR/kafka_2.12-3.7.0.tgz" | tar xz --strip-components=1 -C "/opt/kafka-3.7.0"
 
 
 # Streams test dependencies
@@ -95,6 +96,7 @@ RUN curl -s "$KAFKA_MIRROR/kafka-streams-3.3.2-test.jar" -o /opt/kafka-3.3.2/lib
 RUN curl -s "$KAFKA_MIRROR/kafka-streams-3.4.1-test.jar" -o /opt/kafka-3.4.1/libs/kafka-streams-3.4.1-test.jar
 RUN curl -s "$KAFKA_MIRROR/kafka-streams-3.5.2-test.jar" -o /opt/kafka-3.5.2/libs/kafka-streams-3.5.2-test.jar
 RUN curl -s "$KAFKA_MIRROR/kafka-streams-3.6.1-test.jar" -o /opt/kafka-3.6.1/libs/kafka-streams-3.6.1-test.jar
+RUN curl -s "$KAFKA_MIRROR/kafka-streams-3.7.0-test.jar" -o /opt/kafka-3.7.0/libs/kafka-streams-3.7.0-test.jar
 
 # The version of Kibosh to use for testing.
 # If you update this, also update vagrant/base.sh
diff --git a/vagrant/base.sh b/vagrant/base.sh
index 59bd86fda4fb4..42cfde40e28cf 100755
--- a/vagrant/base.sh
+++ b/vagrant/base.sh
@@ -162,6 +162,8 @@ get_kafka 3.5.2 2.12
 chmod a+rw /opt/kafka-3.5.2
 get_kafka 3.6.1 2.12
 chmod a+rw /opt/kafka-3.6.1
+get_kafka 3.7.0 2.12
+chmod a+rw /opt/kafka-3.7.0
 
 
 # For EC2 nodes, we want to use /mnt, which should have the local disk. On local

From 097932752087be671f3411d531c50a623dc41b87 Mon Sep 17 00:00:00 2001
From: Jeff Kim <kimkb2011@gmail.com>
Date: Tue, 27 Feb 2024 08:45:55 -0500
Subject: [PATCH 085/258] KAFKA-16306: fix GroupCoordinatorService logger
 (#15433)

This patch corrects the logger used for GroupCoordinatorService.

Reviewers: Anton Liauchuk <anton93lev@gmail.com>, David Jacot <djacot@confluent.io>
---
 .../kafka/coordinator/group/GroupCoordinatorService.java | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
index f635b705e9d30..a2363b4822274 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
@@ -230,9 +230,10 @@ public GroupCoordinatorService build() {
 
     /**
      *
-     * @param logContext
-     * @param config
-     * @param runtime
+     * @param logContext                The log context.
+     * @param config                    The group coordinator config.
+     * @param runtime                   The runtime.
+     * @param groupCoordinatorMetrics   The group coordinator metrics.
      */
     GroupCoordinatorService(
         LogContext logContext,
@@ -240,7 +241,7 @@ public GroupCoordinatorService build() {
         CoordinatorRuntime<GroupCoordinatorShard, Record> runtime,
         GroupCoordinatorMetrics groupCoordinatorMetrics
     ) {
-        this.log = logContext.logger(CoordinatorLoader.class);
+        this.log = logContext.logger(GroupCoordinatorService.class);
         this.config = config;
         this.runtime = runtime;
         this.groupCoordinatorMetrics = groupCoordinatorMetrics;

From 5d6936a4992b77ef68da216a7c2dbf1f8c9f909e Mon Sep 17 00:00:00 2001
From: Gaurav Narula <gaurav_narula2@apple.com>
Date: Wed, 28 Feb 2024 09:37:58 +0000
Subject: [PATCH 086/258]  KAFKA-16305: Avoid optimisation in handshakeUnwrap
 (#15434)

Performs additional unwrap during handshake after data from client is processed to support openssl, which needs the extra unwrap to complete handshake.

Reviewers: Ismael Juma <ismael@juma.me.uk>, Rajini Sivaram <rajinisivaram@googlemail.com>
---
 .../common/network/SslTransportLayer.java     |  7 ++-
 .../common/network/SslTransportLayerTest.java | 60 +++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/network/SslTransportLayer.java b/clients/src/main/java/org/apache/kafka/common/network/SslTransportLayer.java
index 904c5216a40d3..da80e363a95c4 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/SslTransportLayer.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/SslTransportLayer.java
@@ -498,13 +498,14 @@ private SSLEngineResult handshakeWrap(boolean doWrite) throws IOException {
     }
 
     /**
-     * Perform handshake unwrap
+     * Perform handshake unwrap.
+     * Visible for testing.
      * @param doRead boolean If true, read more from the socket channel
      * @param ignoreHandshakeStatus If true, continue to unwrap if data available regardless of handshake status
      * @return SSLEngineResult
      * @throws IOException
      */
-    private SSLEngineResult handshakeUnwrap(boolean doRead, boolean ignoreHandshakeStatus) throws IOException {
+    SSLEngineResult handshakeUnwrap(boolean doRead, boolean ignoreHandshakeStatus) throws IOException {
         log.trace("SSLHandshake handshakeUnwrap {}", channelId);
         SSLEngineResult result;
         int read = 0;
@@ -526,7 +527,7 @@ private SSLEngineResult handshakeUnwrap(boolean doRead, boolean ignoreHandshakeS
                     handshakeStatus == HandshakeStatus.NEED_UNWRAP) ||
                     (ignoreHandshakeStatus && netReadBuffer.position() != position);
             log.trace("SSLHandshake handshakeUnwrap: handshakeStatus {} status {}", handshakeStatus, result.getStatus());
-        } while (netReadBuffer.position() != 0 && cont);
+        } while (cont);
 
         // Throw EOF exception for failed read after processing already received data
         // so that handshake failures are reported correctly
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java b/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
index d92f4facb3c57..8b00bcdb955b4 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
@@ -36,6 +36,7 @@
 import org.apache.kafka.test.TestSslUtils;
 import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -51,6 +52,7 @@
 import java.nio.channels.Channels;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -65,13 +67,20 @@
 
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLEngineResult;
 import javax.net.ssl.SSLParameters;
+import javax.net.ssl.SSLSession;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /**
  * Tests for the SSL transport layer. These use a test harness that runs a simple socket server that echos back responses.
@@ -1467,4 +1476,55 @@ int updateAndGet(int actualSize, boolean update) {
             }
         }
     }
+
+    /**
+     * SSLEngine implementations may transition from NEED_UNWRAP to NEED_UNWRAP
+     * even after reading all the data from the socket. This test ensures we
+     * continue unwrapping and not break early.
+     * Please refer <a href="https://issues.apache.org/jira/browse/KAFKA-16305">KAFKA-16305</a>
+     * for more information.
+     */
+    @Test
+    public void testHandshakeUnwrapContinuesUnwrappingOnNeedUnwrapAfterAllBytesRead() throws IOException {
+        // Given
+        byte[] data = "ClientHello?".getBytes(StandardCharsets.UTF_8);
+
+        SSLEngine sslEngine = mock(SSLEngine.class);
+        SocketChannel socketChannel = mock(SocketChannel.class);
+        SelectionKey selectionKey = mock(SelectionKey.class);
+        when(selectionKey.channel()).thenReturn(socketChannel);
+        SSLSession sslSession = mock(SSLSession.class);
+        SslTransportLayer sslTransportLayer = new SslTransportLayer(
+                "test-channel",
+                selectionKey,
+                sslEngine,
+                mock(ChannelMetadataRegistry.class)
+        );
+
+        when(sslEngine.getSession()).thenReturn(sslSession);
+        when(sslSession.getPacketBufferSize()).thenReturn(data.length * 2);
+        sslTransportLayer.startHandshake(); // to initialize the buffers
+
+        ByteBuffer netReadBuffer = sslTransportLayer.netReadBuffer();
+        netReadBuffer.clear();
+        ByteBuffer appReadBuffer = sslTransportLayer.appReadBuffer();
+        when(socketChannel.read(any(ByteBuffer.class))).then(invocation -> {
+            ((ByteBuffer) invocation.getArgument(0)).put(data);
+            return data.length;
+        });
+
+        when(sslEngine.unwrap(netReadBuffer, appReadBuffer))
+                .thenAnswer(invocation -> {
+                    netReadBuffer.flip();
+                    return new SSLEngineResult(SSLEngineResult.Status.OK, SSLEngineResult.HandshakeStatus.NEED_UNWRAP, data.length, 0);
+                }).thenReturn(new SSLEngineResult(SSLEngineResult.Status.OK, SSLEngineResult.HandshakeStatus.NEED_WRAP, 0, 0));
+
+        // When
+        SSLEngineResult result = sslTransportLayer.handshakeUnwrap(true, false);
+
+        // Then
+        verify(sslEngine, times(2)).unwrap(netReadBuffer, appReadBuffer);
+        assertEquals(SSLEngineResult.Status.OK, result.getStatus());
+        assertEquals(SSLEngineResult.HandshakeStatus.NEED_WRAP, result.getHandshakeStatus());
+    }
 }

From 53c41aca7ba9469a0145023112f5fad254da4fa8 Mon Sep 17 00:00:00 2001
From: Philip Nee <pnee@confluent.io>
Date: Wed, 28 Feb 2024 01:52:01 -0800
Subject: [PATCH 087/258] KAFKA-16116: Rebalance Metrics for AsyncKafkaConsumer
 (#15339)

Adding the following rebalance metrics to the consumer:

rebalance-latency-avg
rebalance-latency-max
rebalance-latency-total
rebalance-rate-per-hour
rebalance-total
failed-rebalance-rate-per-hour
failed-rebalance-total

Due to the difference in protocol, we need to redefine when rebalance starts and ends.
Start of Rebalance:
Current: Right before sending out JoinGroup
ConsumerGroup: When the client receives assignments from the HB

End of Rebalance - Successful Case:
Current: Receiving SyncGroup request after transitioning to "COMPLETING_REBALANCE"
ConsumerGroup: After completing reconciliation and right before sending out "Ack" heartbeat

End of Rebalance - Failed Case:
Current: Any failure in the JoinGroup/SyncGroup response
ConsumerGroup: Failure in the heartbeat

Note: Afterall, we try to be consistent with the current protocol. Rebalances start and end with sending and receiving network requests. Failures in network requests signify the user failures in rebalance. And it is entirely possible to have multiple failures before having a successful one.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/HeartbeatRequestManager.java    |   3 +-
 .../consumer/internals/MembershipManager.java |   7 +-
 .../internals/MembershipManagerImpl.java      |  62 ++++-
 .../consumer/internals/RequestManagers.java   |   3 +-
 .../metrics/RebalanceMetricsManager.java      | 114 +++++++++
 .../internals/ConsumerTestBuilder.java        |  26 +-
 .../HeartbeatRequestManagerTest.java          |   8 +-
 .../internals/MembershipManagerImplTest.java  | 240 +++++++++++++++---
 8 files changed, 408 insertions(+), 55 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/RebalanceMetricsManager.java

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 550e5b92ebf24..d551dbe2508c0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -327,7 +327,7 @@ private void onResponse(final ConsumerGroupHeartbeatResponse response, long curr
             heartbeatRequestState.updateHeartbeatIntervalMs(response.data().heartbeatIntervalMs());
             heartbeatRequestState.onSuccessfulAttempt(currentTimeMs);
             heartbeatRequestState.resetTimer();
-            membershipManager.onHeartbeatResponseReceived(response.data());
+            membershipManager.onHeartbeatSuccess(response.data());
             maybeSendGroupMetadataUpdateEvent();
             return;
         }
@@ -357,6 +357,7 @@ private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
 
         this.heartbeatState.reset();
         this.heartbeatRequestState.onFailedAttempt(currentTimeMs);
+        membershipManager.onHeartbeatFailure();
 
         switch (error) {
             case NOT_COORDINATOR:
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
index a9c23d7b4d5da..f0a641d140fb2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
@@ -68,7 +68,12 @@ public interface MembershipManager extends RequestManager {
      *
      * @param response Heartbeat response to extract member info and errors from.
      */
-    void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData response);
+    void onHeartbeatSuccess(ConsumerGroupHeartbeatResponseData response);
+
+    /**
+     * Notify the member that an error heartbeat response was received.
+     */
+    void onHeartbeatFailure();
 
     /**
      * Update state when a heartbeat is sent out. This will transition out of the states that end
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index a9b0f3b94d831..6f3947eea46e5 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -27,11 +27,13 @@
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
 import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
+import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest;
 import org.apache.kafka.common.telemetry.internals.ClientTelemetryProvider;
@@ -263,6 +265,11 @@ public class MembershipManagerImpl implements MembershipManager {
      * when the timer is reset, only when it completes releasing its assignment.
      */
     private CompletableFuture<Void> staleMemberAssignmentRelease;
+    
+    /*
+     * Measures successful rebalance latency and number of failed rebalances.
+     */
+    private final RebalanceMetricsManager metricsManager;
 
     private final Time time;
 
@@ -284,7 +291,35 @@ public MembershipManagerImpl(String groupId,
                                  LogContext logContext,
                                  Optional<ClientTelemetryReporter> clientTelemetryReporter,
                                  BackgroundEventHandler backgroundEventHandler,
-                                 Time time) {
+                                 Time time,
+                                 Metrics metrics) {
+        this(groupId,
+            groupInstanceId,
+            rebalanceTimeoutMs,
+            serverAssignor,
+            subscriptions,
+            commitRequestManager,
+            metadata,
+            logContext,
+            clientTelemetryReporter,
+            backgroundEventHandler,
+            time,
+            new RebalanceMetricsManager(metrics));
+    }
+
+    // Visible for testing
+    MembershipManagerImpl(String groupId,
+                          Optional<String> groupInstanceId,
+                          int rebalanceTimeoutMs,
+                          Optional<String> serverAssignor,
+                          SubscriptionState subscriptions,
+                          CommitRequestManager commitRequestManager,
+                          ConsumerMetadata metadata,
+                          LogContext logContext,
+                          Optional<ClientTelemetryReporter> clientTelemetryReporter,
+                          BackgroundEventHandler backgroundEventHandler,
+                          Time time,
+                          RebalanceMetricsManager metricsManager) {
         this.groupId = groupId;
         this.state = MemberState.UNSUBSCRIBED;
         this.serverAssignor = serverAssignor;
@@ -301,6 +336,7 @@ public MembershipManagerImpl(String groupId,
         this.rebalanceTimeoutMs = rebalanceTimeoutMs;
         this.backgroundEventHandler = backgroundEventHandler;
         this.time = time;
+        this.metricsManager = metricsManager;
     }
 
     /**
@@ -314,10 +350,27 @@ private void transitionTo(MemberState nextState) {
             throw new IllegalStateException(String.format("Invalid state transition from %s to %s",
                     state, nextState));
         }
+
+        if (isCompletingRebalance(state, nextState)) {
+            metricsManager.recordRebalanceEnded(time.milliseconds());
+        }
+        if (isStartingRebalance(state, nextState)) {
+            metricsManager.recordRebalanceStarted(time.milliseconds());
+        }
+
         log.trace("Member {} with epoch {} transitioned from {} to {}.", memberId, memberEpoch, state, nextState);
         this.state = nextState;
     }
 
+    private static boolean isCompletingRebalance(MemberState currentState, MemberState nextState) {
+        return currentState == MemberState.RECONCILING &&
+            (nextState == MemberState.STABLE || nextState == MemberState.ACKNOWLEDGING);
+    }
+
+    private static boolean isStartingRebalance(MemberState currentState, MemberState nextState) {
+        return currentState != MemberState.RECONCILING && nextState == MemberState.RECONCILING;
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -354,7 +407,7 @@ public int memberEpoch() {
      * {@inheritDoc}
      */
     @Override
-    public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData response) {
+    public void onHeartbeatSuccess(ConsumerGroupHeartbeatResponseData response) {
         if (response.errorCode() != Errors.NONE.code()) {
             String errorMessage = String.format(
                     "Unexpected error in Heartbeat response. Expected no error, but received: %s",
@@ -403,6 +456,11 @@ public void onHeartbeatResponseReceived(ConsumerGroupHeartbeatResponseData respo
         }
     }
 
+    @Override
+    public void onHeartbeatFailure() {
+        metricsManager.maybeRecordRebalanceFailed();
+    }
+
     /**
      * @return True if the consumer is not a member of the group.
      */
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
index 2d90a3ad7082e..0b4c043d4a4d6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
@@ -189,7 +189,8 @@ protected RequestManagers create() {
                             logContext,
                             clientTelemetryReporter,
                             backgroundEventHandler,
-                            time);
+                            time,
+                            metrics);
                     membershipManager.registerStateListener(commit);
                     heartbeatRequestManager = new HeartbeatRequestManager(
                             logContext,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/RebalanceMetricsManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/RebalanceMetricsManager.java
new file mode 100644
index 0000000000000..a255487f37a08
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/RebalanceMetricsManager.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.metrics;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.Measurable;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.metrics.stats.CumulativeCount;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
+import org.apache.kafka.common.metrics.stats.Max;
+import org.apache.kafka.common.metrics.stats.Rate;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
+
+public class RebalanceMetricsManager {
+    private final Sensor successfulRebalanceSensor;
+    private final Sensor failedRebalanceSensor;
+    private final String metricGroupName;
+
+    public final MetricName rebalanceLatencyAvg;
+    public final MetricName rebalanceLatencyMax;
+    public final MetricName rebalanceLatencyTotal;
+    public final MetricName rebalanceTotal;
+    public final MetricName rebalanceRatePerHour;
+    public final MetricName lastRebalanceSecondsAgo;
+    public final MetricName failedRebalanceTotal;
+    public final MetricName failedRebalanceRate;
+    private long lastRebalanceEndMs = -1L;
+    private long lastRebalanceStartMs = -1L;
+
+    public RebalanceMetricsManager(Metrics metrics) {
+        metricGroupName = CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX;
+
+        rebalanceLatencyAvg = createMetric(metrics, "rebalance-latency-avg",
+            "The average time taken for a group to complete a rebalance");
+        rebalanceLatencyMax = createMetric(metrics, "rebalance-latency-max",
+            "The max time taken for a group to complete a rebalance");
+        rebalanceLatencyTotal = createMetric(metrics, "rebalance-latency-total",
+            "The total number of milliseconds spent in rebalances");
+        rebalanceTotal = createMetric(metrics, "rebalance-total",
+            "The total number of rebalance events");
+        rebalanceRatePerHour = createMetric(metrics, "rebalance-rate-per-hour",
+            "The number of rebalance events per hour");
+        failedRebalanceTotal = createMetric(metrics, "failed-rebalance-total",
+            "The total number of failed rebalance events");
+        failedRebalanceRate = createMetric(metrics, "failed-rebalance-rate-per-hour",
+            "The number of failed rebalance events per hour");
+
+        successfulRebalanceSensor = metrics.sensor("rebalance-latency");
+        successfulRebalanceSensor.add(rebalanceLatencyAvg, new Avg());
+        successfulRebalanceSensor.add(rebalanceLatencyMax, new Max());
+        successfulRebalanceSensor.add(rebalanceLatencyTotal, new CumulativeSum());
+        successfulRebalanceSensor.add(rebalanceTotal, new CumulativeCount());
+        successfulRebalanceSensor.add(rebalanceRatePerHour, new Rate(TimeUnit.HOURS, new WindowedCount()));
+
+        failedRebalanceSensor = metrics.sensor("failed-rebalance");
+        failedRebalanceSensor.add(failedRebalanceTotal, new CumulativeSum());
+        failedRebalanceSensor.add(failedRebalanceRate, new Rate(TimeUnit.HOURS, new WindowedCount()));
+
+        Measurable lastRebalance = (config, now) -> {
+            if (lastRebalanceEndMs == -1L)
+                return -1d;
+            else
+                return TimeUnit.SECONDS.convert(now - lastRebalanceEndMs, TimeUnit.MILLISECONDS);
+        };
+        lastRebalanceSecondsAgo = createMetric(metrics,
+            "last-rebalance-seconds-ago",
+            "The number of seconds since the last rebalance event");
+        metrics.addMetric(lastRebalanceSecondsAgo, lastRebalance);
+    }
+
+    private MetricName createMetric(Metrics metrics, String name, String description) {
+        return metrics.metricName(name, metricGroupName, description);
+    }
+
+    public void recordRebalanceStarted(long nowMs) {
+        lastRebalanceStartMs = nowMs;
+    }
+
+    public void recordRebalanceEnded(long nowMs) {
+        lastRebalanceEndMs = nowMs;
+        successfulRebalanceSensor.record(nowMs - lastRebalanceStartMs);
+    }
+
+    public void maybeRecordRebalanceFailed() {
+        if (lastRebalanceStartMs <= lastRebalanceEndMs)
+            return;
+        failedRebalanceSensor.record();
+    }
+
+    public boolean rebalanceStarted() {
+        return lastRebalanceStartMs > lastRebalanceEndMs;
+    }
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
index d6ae62950608c..d62d3d8a35ec4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerTestBuilder.java
@@ -26,6 +26,7 @@
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.metrics.RebalanceCallbackMetricsManager;
+import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.requests.MetadataResponse;
@@ -196,18 +197,19 @@ public ConsumerTestBuilder(Optional<GroupInformation> groupInfo, boolean enableA
                     gi.groupInstanceId,
                     metrics));
             MembershipManager mm = spy(
-                    new MembershipManagerImpl(
-                        gi.groupId,
-                        gi.groupInstanceId,
-                        groupRebalanceConfig.rebalanceTimeoutMs,
-                        gi.serverAssignor,
-                        subscriptions,
-                        commit,
-                        metadata,
-                        logContext,
-                        Optional.empty(),
-                        backgroundEventHandler,
-                        time
+                new MembershipManagerImpl(
+                    gi.groupId,
+                    gi.groupInstanceId,
+                    groupRebalanceConfig.rebalanceTimeoutMs,
+                    gi.serverAssignor,
+                    subscriptions,
+                    commit,
+                    metadata,
+                    logContext,
+                    Optional.empty(),
+                    backgroundEventHandler,
+                    time,
+                    mock(RebalanceMetricsManager.class)
                 )
             );
             HeartbeatRequestManager.HeartbeatState heartbeatState = spy(new HeartbeatRequestManager.HeartbeatState(
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 5cf5b9e2d92c5..4d4492bcb47e4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -306,7 +306,7 @@ public void testValidateConsumerGroupHeartbeatRequest(final short version) {
             new ConsumerGroupHeartbeatResponse(new ConsumerGroupHeartbeatResponseData()
             .setMemberId(memberId)
             .setMemberEpoch(memberEpoch));
-        membershipManager.onHeartbeatResponseReceived(result.data());
+        membershipManager.onHeartbeatSuccess(result.data());
 
         // Create a ConsumerHeartbeatRequest and verify the payload
         NetworkClientDelegate.PollResult pollResult = heartbeatRequestManager.poll(time.milliseconds());
@@ -441,7 +441,7 @@ public void testHeartbeatResponseOnErrorHandling(final Errors error, final boole
         switch (error) {
             case NONE:
                 verify(backgroundEventHandler).add(any(GroupMetadataUpdateEvent.class));
-                verify(membershipManager, times(2)).onHeartbeatResponseReceived(mockResponse.data());
+                verify(membershipManager, times(2)).onHeartbeatSuccess(mockResponse.data());
                 assertEquals(DEFAULT_HEARTBEAT_INTERVAL_MS, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()));
                 break;
 
@@ -547,7 +547,7 @@ public void testHeartbeatState() {
                 .setMemberEpoch(1)
                 .setAssignment(assignmentTopic1));
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, "topic1"));
-        membershipManager.onHeartbeatResponseReceived(rs1.data());
+        membershipManager.onHeartbeatSuccess(rs1.data());
 
         // We remain in RECONCILING state, as the assignment will be reconciled on the next poll
         assertEquals(MemberState.RECONCILING, membershipManager.state());
@@ -712,7 +712,7 @@ private void mockStableMember() {
                 .setHeartbeatIntervalMs(DEFAULT_HEARTBEAT_INTERVAL_MS)
                 .setMemberId(memberId)
                 .setMemberEpoch(memberEpoch));
-        membershipManager.onHeartbeatResponseReceived(rs1.data());
+        membershipManager.onHeartbeatSuccess(rs1.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 3294068b07487..8a0fcf85758e1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -16,12 +16,15 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.RebalanceCallbackMetricsManager;
+import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.Uuid;
@@ -102,6 +105,8 @@ public class MembershipManagerImplTest {
     private BlockingQueue<BackgroundEvent> backgroundEventQueue;
     private BackgroundEventHandler backgroundEventHandler;
     private Time time;
+    private Metrics metrics;
+    private RebalanceMetricsManager rebalanceMetricsManager;
 
     @BeforeEach
     public void setup() {
@@ -111,7 +116,9 @@ public void setup() {
         commitRequestManager = testBuilder.commitRequestManager.orElseThrow(IllegalStateException::new);
         backgroundEventQueue = testBuilder.backgroundEventQueue;
         backgroundEventHandler = testBuilder.backgroundEventHandler;
-        time = testBuilder.time;
+        time = new MockTime(0);
+        metrics = new Metrics(time);
+        rebalanceMetricsManager = new RebalanceMetricsManager(metrics);
     }
 
     @AfterEach
@@ -135,15 +142,16 @@ private MembershipManagerImpl createMembershipManager(String groupInstanceId) {
         return spy(new MembershipManagerImpl(
             GROUP_ID, Optional.ofNullable(groupInstanceId), REBALANCE_TIMEOUT, Optional.empty(),
             subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
-            backgroundEventHandler, time));
+            backgroundEventHandler, time, rebalanceMetricsManager));
     }
 
     private MembershipManagerImpl createMembershipManagerJoiningGroup(String groupInstanceId,
                                                                       String serverAssignor) {
-        MembershipManagerImpl manager = new MembershipManagerImpl(
+        MembershipManagerImpl manager = spy(new MembershipManagerImpl(
                 GROUP_ID, Optional.ofNullable(groupInstanceId), REBALANCE_TIMEOUT,
                 Optional.ofNullable(serverAssignor), subscriptionState, commitRequestManager,
-                metadata, logContext, Optional.empty(), backgroundEventHandler, time);
+                metadata, logContext, Optional.empty(), backgroundEventHandler, time,
+                rebalanceMetricsManager));
         manager.transitionToJoining();
         return manager;
     }
@@ -160,7 +168,6 @@ public void testMembershipManagerServerAssignor() {
     @Test
     public void testMembershipManagerInitSupportsEmptyGroupInstanceId() {
         createMembershipManagerJoiningGroup();
-        createMembershipManagerJoiningGroup(null, null);
     }
 
     @Test
@@ -169,7 +176,7 @@ public void testMembershipManagerRegistersForClusterMetadataUpdatesOnFirstJoin()
         MembershipManagerImpl manager = new MembershipManagerImpl(
                 GROUP_ID, Optional.empty(), REBALANCE_TIMEOUT, Optional.empty(),
                 subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
-                backgroundEventHandler, time);
+                backgroundEventHandler, time, rebalanceMetricsManager);
         manager.transitionToJoining();
         clearInvocations(metadata);
 
@@ -200,12 +207,12 @@ public void testTransitionToReconcilingOnlyIfAssignmentReceived() {
 
         ConsumerGroupHeartbeatResponse responseWithoutAssignment =
                 createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(responseWithoutAssignment.data());
+        membershipManager.onHeartbeatSuccess(responseWithoutAssignment.data());
         assertNotEquals(MemberState.RECONCILING, membershipManager.state());
 
         ConsumerGroupHeartbeatResponse responseWithAssignment =
                 createConsumerGroupHeartbeatResponse(createAssignment(true));
-        membershipManager.onHeartbeatResponseReceived(responseWithAssignment.data());
+        membershipManager.onHeartbeatSuccess(responseWithAssignment.data());
         assertEquals(MemberState.RECONCILING, membershipManager.state());
     }
 
@@ -213,7 +220,7 @@ public void testTransitionToReconcilingOnlyIfAssignmentReceived() {
     public void testMemberIdAndEpochResetOnFencedMembers() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
@@ -230,7 +237,7 @@ public void testTransitionToFatal() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse =
                 createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
@@ -246,7 +253,7 @@ public void testTransitionToFailedWhenTryingToJoin() {
         MembershipManagerImpl membershipManager = new MembershipManagerImpl(
                 GROUP_ID, Optional.empty(), REBALANCE_TIMEOUT, Optional.empty(),
                 subscriptionState, commitRequestManager, metadata, logContext, Optional.empty(),
-                backgroundEventHandler, time);
+                backgroundEventHandler, time, rebalanceMetricsManager);
         assertEquals(MemberState.UNSUBSCRIBED, membershipManager.state());
         membershipManager.transitionToJoining();
 
@@ -299,7 +306,7 @@ public void testListenersGetNotifiedOfMemberEpochUpdatesOnlyIfItChanges() {
         membershipManager.registerStateListener(listener);
         int epoch = 5;
 
-        membershipManager.onHeartbeatResponseReceived(new ConsumerGroupHeartbeatResponseData()
+        membershipManager.onHeartbeatSuccess(new ConsumerGroupHeartbeatResponseData()
                 .setErrorCode(Errors.NONE.code())
                 .setMemberId(MEMBER_ID)
                 .setMemberEpoch(epoch));
@@ -307,7 +314,7 @@ public void testListenersGetNotifiedOfMemberEpochUpdatesOnlyIfItChanges() {
         verify(listener).onMemberEpochUpdated(Optional.of(epoch), Optional.of(MEMBER_ID));
         clearInvocations(listener);
 
-        membershipManager.onHeartbeatResponseReceived(new ConsumerGroupHeartbeatResponseData()
+        membershipManager.onHeartbeatSuccess(new ConsumerGroupHeartbeatResponseData()
                 .setErrorCode(Errors.NONE.code())
                 .setMemberId(MEMBER_ID)
                 .setMemberEpoch(epoch));
@@ -316,7 +323,7 @@ public void testListenersGetNotifiedOfMemberEpochUpdatesOnlyIfItChanges() {
 
     private void mockStableMember(MembershipManagerImpl membershipManager) {
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
@@ -709,7 +716,7 @@ public void testIgnoreHeartbeatWhenLeavingGroup() {
 
         CompletableFuture<Void> leaveResult = membershipManager.leaveGroup();
 
-        membershipManager.onHeartbeatResponseReceived(createConsumerGroupHeartbeatResponse(createAssignment(true)).data());
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(createAssignment(true)).data());
 
         assertEquals(MemberState.LEAVING, membershipManager.state());
         assertEquals(-1, membershipManager.memberEpoch());
@@ -726,7 +733,7 @@ public void testIgnoreHeartbeatResponseWhenNotInGroup(MemberState state) {
         when(membershipManager.state()).thenReturn(state);
         ConsumerGroupHeartbeatResponseData responseData = mock(ConsumerGroupHeartbeatResponseData.class);
 
-        membershipManager.onHeartbeatResponseReceived(responseData);
+        membershipManager.onHeartbeatSuccess(responseData);
 
         assertEquals(state, membershipManager.state());
         verify(responseData, never()).memberId();
@@ -861,7 +868,7 @@ public void testFatalFailureWhenStateIsUnjoined() {
     public void testFatalFailureWhenStateIsStable() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
 
         testStateUpdateOnFatalFailure(membershipManager);
@@ -930,9 +937,9 @@ public void testUpdateStateFailsOnResponsesWithErrors() {
         // Updating state with a heartbeat response containing errors cannot be performed and
         // should fail.
         ConsumerGroupHeartbeatResponse unknownMemberResponse =
-                createConsumerGroupHeartbeatResponseWithError();
+                createConsumerGroupHeartbeatResponseWithError(Errors.UNKNOWN_MEMBER_ID);
         assertThrows(IllegalArgumentException.class,
-                () -> membershipManager.onHeartbeatResponseReceived(unknownMemberResponse.data()));
+                () -> membershipManager.onHeartbeatSuccess(unknownMemberResponse.data()));
     }
 
     /**
@@ -1099,7 +1106,7 @@ public void testMemberKeepsUnresolvedAssignmentWaitingForMetadataUntilResolved()
         // Target assignment received again with the same unresolved topic. Client should keep it
         // as unresolved.
         clearInvocations(subscriptionState);
-        membershipManager.onHeartbeatResponseReceived(createConsumerGroupHeartbeatResponse(assignment).data());
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment).data());
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertEquals(Collections.singleton(topic2), membershipManager.topicsAwaitingReconciliation());
         verify(subscriptionState, never()).assignFromSubscribed(anyCollection());
@@ -1173,6 +1180,9 @@ public void testReconciliationSkippedWhenSameAssignmentReceived() {
         verify(subscriptionState, never()).assignFromSubscribed(anyCollection());
 
         assertEquals(MemberState.STABLE, membershipManager.state());
+
+        assertEquals(1.0d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        assertEquals(0.0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
     }
 
     @Test
@@ -1376,7 +1386,6 @@ public void testOnSubscriptionUpdatedTransitionsToJoiningOnlyIfNotInGroup() {
 
     @Test
     public void testListenerCallbacksBasic() {
-        // Step 1: set up mocks
         MembershipManagerImpl membershipManager = createMemberInStableState();
         CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
         ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
@@ -1685,7 +1694,7 @@ public void testTransitionToLeavingWhileJoiningDueToStaleMember() {
     public void testTransitionToLeavingWhileStableDueToStaleMember() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
@@ -1832,8 +1841,8 @@ private ConsumerRebalanceListenerInvoker consumerRebalanceListenerInvoker() {
         return new ConsumerRebalanceListenerInvoker(
                 new LogContext(),
                 subscriptionState,
-                new MockTime(1),
-                new RebalanceCallbackMetricsManager(new Metrics())
+                time,
+                new RebalanceCallbackMetricsManager(new Metrics(time))
         );
     }
 
@@ -1930,7 +1939,7 @@ private void assertStaleMemberLeavesGroupAndClearsAssignment(MembershipManagerIm
 
     @Test
     public void testMemberJoiningTransitionsToStableWhenReceivingEmptyAssignment() {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(null);
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         assertEquals(MemberState.JOINING, membershipManager.state());
         receiveEmptyAssignment(membershipManager);
 
@@ -1940,6 +1949,144 @@ public void testMemberJoiningTransitionsToStableWhenReceivingEmptyAssignment() {
         assertEquals(MemberState.STABLE, membershipManager.state());
     }
 
+    @Test
+    public void testMetricsWhenHeartbeatFailed() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        membershipManager.onHeartbeatFailure();
+
+        // Not expecting rebalance failures without assignments being reconciled
+        assertEquals(0.0d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        assertEquals(0.0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
+    }
+
+    @Test
+    public void testRebalanceMetricsOnSuccessfulRebalance() {
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+        mockOwnedPartition(membershipManager, Uuid.randomUuid(), "topic1");
+
+        CompletableFuture<Void> commitResult = mockRevocationNoCallbacks(true);
+
+        receiveEmptyAssignment(membershipManager);
+        long reconciliationDurationMs = 1234;
+        time.sleep(reconciliationDurationMs);
+
+        membershipManager.poll(time.milliseconds());
+        // Complete commit request to complete the callback invocation
+        commitResult.complete(null);
+
+        assertEquals((double) reconciliationDurationMs, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyTotal));
+        assertEquals((double) reconciliationDurationMs, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyAvg));
+        assertEquals((double) reconciliationDurationMs, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyMax));
+        assertEquals(1d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        assertEquals(120d, 1d, (double) getMetricValue(metrics, rebalanceMetricsManager.rebalanceRatePerHour));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceRate));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.lastRebalanceSecondsAgo));
+    }
+
+    @Test
+    public void testRebalanceMetricsForMultipleReconcilations() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+
+        String topicName = "topic1";
+        Uuid topicId = Uuid.randomUuid();
+
+        SleepyRebalanceListener listener = new SleepyRebalanceListener(1453, time);
+        when(subscriptionState.assignedPartitions()).thenReturn(Collections.emptySet());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+        doNothing().when(subscriptionState).markPendingRevocation(anySet());
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
+
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
+        receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
+
+        membershipManager.poll(time.milliseconds());
+
+        // assign partitions
+        performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
+            topicPartitions(topicName, 0, 1),
+            true
+        );
+
+        long firstRebalanaceTimesMs = listener.sleepMs;
+        listener.reset();
+
+        // ack
+        membershipManager.onHeartbeatRequestSent();
+
+        // revoke all
+        when(subscriptionState.assignedPartitions()).thenReturn(topicPartitions(topicName, 0, 1));
+        receiveAssignment(topicId, Collections.singletonList(2), membershipManager);
+
+        membershipManager.poll(time.milliseconds());
+
+        performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED,
+            topicPartitions(topicName, 0, 1),
+            true
+        );
+
+        // assign new partition 2
+        performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
+            topicPartitions(topicName, 2),
+            true
+        );
+        membershipManager.onHeartbeatRequestSent();
+
+        long secondRebalanceMs = listener.sleepMs;
+        long total = firstRebalanaceTimesMs + secondRebalanceMs;
+        double avg = total / 2.0d;
+        long max = Math.max(firstRebalanaceTimesMs, secondRebalanceMs);
+        assertEquals((double) total, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyTotal));
+        assertEquals(avg, (double) getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyAvg), 1d);
+        assertEquals((double) max, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyMax));
+        assertEquals(2d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        // rate is not tested because it is subject to Rate implementation
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceRate));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.lastRebalanceSecondsAgo));
+
+    }
+
+    @Test
+    public void testRebalanceMetricsOnFailedRebalance() {
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+
+        Uuid topicId = Uuid.randomUuid();
+
+        receiveAssignment(topicId, Arrays.asList(0, 1), membershipManager);
+
+        // sleep for an arbitrary amount
+        time.sleep(2300);
+
+        assertTrue(rebalanceMetricsManager.rebalanceStarted());
+        membershipManager.onHeartbeatFailure();
+
+        assertEquals((double) 0, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyTotal));
+        assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        assertEquals(120d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceRate));
+        assertEquals(1d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
+        assertEquals(-1d, getMetricValue(metrics, rebalanceMetricsManager.lastRebalanceSecondsAgo));
+    }
+
+    private Object getMetricValue(Metrics metrics, MetricName name) {
+        return metrics.metrics().get(name).metricValue();
+    }
+
     private MembershipManagerImpl mockMemberSuccessfullyReceivesAndAcksAssignment(
             Uuid topicId, String topicName, List<Integer> partitions) {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
@@ -2119,7 +2266,7 @@ private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean expectSubscri
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty()).thenReturn(Optional.empty());
 
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         membershipManager.poll(time.milliseconds());
 
         if (expectSubscriptionUpdated) {
@@ -2136,9 +2283,9 @@ private MembershipManagerImpl createMemberInStableState() {
     }
 
     private MembershipManagerImpl createMemberInStableState(String groupInstanceId) {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(groupInstanceId);
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(groupInstanceId, null);
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         assertEquals(MemberState.STABLE, membershipManager.state());
         return membershipManager;
     }
@@ -2150,7 +2297,7 @@ private void receiveAssignment(Map<Uuid, SortedSet<Integer>> topicIdPartitionLis
                     .setTopicId(tp.getKey())
                     .setPartitions(new ArrayList<>(tp.getValue()))).collect(Collectors.toList()));
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(targetAssignment);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
     }
 
     private void receiveAssignment(Uuid topicId, List<Integer> partitions, MembershipManager membershipManager) {
@@ -2160,7 +2307,7 @@ private void receiveAssignment(Uuid topicId, List<Integer> partitions, Membershi
                                 .setTopicId(topicId)
                                 .setPartitions(partitions)));
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(targetAssignment);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
     }
 
     private void receiveAssignmentAfterRejoin(Uuid topicId, List<Integer> partitions, MembershipManager membershipManager) {
@@ -2171,7 +2318,7 @@ private void receiveAssignmentAfterRejoin(Uuid topicId, List<Integer> partitions
                                 .setPartitions(partitions)));
         ConsumerGroupHeartbeatResponse heartbeatResponse =
                 createConsumerGroupHeartbeatResponseWithBumpedEpoch(targetAssignment);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
     }
 
     private void receiveEmptyAssignment(MembershipManager membershipManager) {
@@ -2179,7 +2326,7 @@ private void receiveEmptyAssignment(MembershipManager membershipManager) {
         ConsumerGroupHeartbeatResponseData.Assignment targetAssignment = new ConsumerGroupHeartbeatResponseData.Assignment()
                 .setTopicPartitions(Collections.emptyList());
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(targetAssignment);
-        membershipManager.onHeartbeatResponseReceived(heartbeatResponse.data());
+        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
     }
 
     /**
@@ -2307,9 +2454,9 @@ private ConsumerGroupHeartbeatResponse createConsumerGroupHeartbeatResponseWithB
                 .setAssignment(assignment));
     }
 
-    private ConsumerGroupHeartbeatResponse createConsumerGroupHeartbeatResponseWithError() {
+    private ConsumerGroupHeartbeatResponse createConsumerGroupHeartbeatResponseWithError(Errors error) {
         return new ConsumerGroupHeartbeatResponse(new ConsumerGroupHeartbeatResponseData()
-                .setErrorCode(Errors.UNKNOWN_MEMBER_ID.code())
+                .setErrorCode(error.code())
                 .setMemberId(MEMBER_ID)
                 .setMemberEpoch(5));
     }
@@ -2356,4 +2503,29 @@ private static Stream<Arguments> notInGroupStates() {
             Arguments.of(MemberState.STALE));
     }
 
+    private static class SleepyRebalanceListener implements ConsumerRebalanceListener {
+        private long sleepMs;
+        private final long sleepDurationMs;
+        private final Time time;
+        SleepyRebalanceListener(long sleepDurationMs, Time time) {
+            this.sleepDurationMs = sleepDurationMs;
+            this.time = time;
+        }
+
+        @Override
+        public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+            sleepMs += sleepDurationMs;
+            time.sleep(sleepDurationMs);
+        }
+
+        @Override
+        public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+            sleepMs += sleepDurationMs;
+            time.sleep(sleepDurationMs);
+        }
+
+        public void reset() {
+            sleepMs = 0;
+        }
+    }
 }

From 52289c92be45ba3758d07376d9c64ddadbecb544 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Wed, 28 Feb 2024 05:38:02 -0800
Subject: [PATCH 088/258] MINOR: Optimize EventAccumulator (#15430)

`poll(long timeout, TimeUnit unit)` is either used with `Long.MAX_VALUE` or `0`. This patch replaces it with `poll` and `take`. It removes the `awaitNanos` usage.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../group/runtime/EventAccumulator.java       | 37 +++++++----
 .../runtime/MultiThreadedEventProcessor.java  |  9 +--
 .../group/runtime/EventAccumulatorTest.java   | 30 ++++-----
 .../MultiThreadedEventProcessorTest.java      | 61 +++++--------------
 4 files changed, 56 insertions(+), 81 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
index f46e8b8a8bfe3..16b61f8e99150 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
@@ -27,7 +27,6 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -137,31 +136,43 @@ public void add(T event) throws RejectedExecutionException {
     }
 
     /**
-     * Returns the next {{@link Event}} available. This method block indefinitely until
-     * one event is ready or the accumulator is closed.
+     * Returns the next {{@link Event}} available or null if no event is
+     * available.
      *
-     * @return The next event.
+     * @return The next event available or null.
      */
     public T poll() {
-        return poll(Long.MAX_VALUE, TimeUnit.SECONDS);
+        lock.lock();
+        try {
+            K key = randomKey();
+            if (key == null) return null;
+
+            Queue<T> queue = queues.get(key);
+            T event = queue.poll();
+
+            if (queue.isEmpty()) queues.remove(key);
+            inflightKeys.add(key);
+            size--;
+
+            return event;
+        } finally {
+            lock.unlock();
+        }
     }
 
     /**
-     * Returns the next {{@link Event}} available. This method blocks for the provided
-     * time and returns null of not event is available.
+     * Returns the next {{@link Event}} available. This method blocks until an
+     * event is available or accumulator is closed.
      *
-     * @param timeout   The timeout.
-     * @param unit      The timeout unit.
      * @return The next event available or null.
      */
-    public T poll(long timeout, TimeUnit unit) {
+    public T take() {
         lock.lock();
         try {
             K key = randomKey();
-            long nanos = unit.toNanos(timeout);
-            while (key == null && !closed && nanos > 0) {
+            while (key == null && !closed) {
                 try {
-                    nanos = condition.awaitNanos(nanos);
+                    condition.await();
                 } catch (InterruptedException e) {
                     // Ignore.
                 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
index e4adc18e95793..0e3d563861c86 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
@@ -25,7 +25,6 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -128,7 +127,7 @@ private class EventProcessorThread extends Thread {
         private void handleEvents() {
             while (!shuttingDown) {
                 recordPollStartTime(time.milliseconds());
-                CoordinatorEvent event = accumulator.poll();
+                CoordinatorEvent event = accumulator.take();
                 recordPollEndTime(time.milliseconds());
                 if (event != null) {
                     try {
@@ -148,8 +147,8 @@ private void handleEvents() {
         }
 
         private void drainEvents() {
-            CoordinatorEvent event = accumulator.poll(0, TimeUnit.MILLISECONDS);
-            while (event != null) {
+            CoordinatorEvent event;
+            while ((event = accumulator.poll()) != null) {
                 try {
                     log.debug("Draining event: {}.", event);
                     metrics.recordEventQueueTime(time.milliseconds() - event.createdTimeMs());
@@ -159,8 +158,6 @@ private void drainEvents() {
                 } finally {
                     accumulator.done(event);
                 }
-
-                event = accumulator.poll(0, TimeUnit.MILLISECONDS);
             }
         }
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
index 147cf08121c47..e077fb5e0226c 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
@@ -78,7 +78,7 @@ public void testBasicOperations() {
         EventAccumulator<Integer, MockEvent> accumulator = new EventAccumulator<>();
 
         assertEquals(0, accumulator.size());
-        assertNull(accumulator.poll(0, TimeUnit.MICROSECONDS));
+        assertNull(accumulator.poll());
 
         List<MockEvent> events = Arrays.asList(
             new MockEvent(1, 0),
@@ -97,14 +97,14 @@ public void testBasicOperations() {
 
         Set<MockEvent> polledEvents = new HashSet<>();
         for (int i = 0; i < events.size(); i++) {
-            MockEvent event = accumulator.poll(0, TimeUnit.MICROSECONDS);
+            MockEvent event = accumulator.poll();
             assertNotNull(event);
             polledEvents.add(event);
             assertEquals(events.size() - 1 - i, accumulator.size());
             accumulator.done(event);
         }
 
-        assertNull(accumulator.poll(0, TimeUnit.MICROSECONDS));
+        assertNull(accumulator.poll());
         assertEquals(new HashSet<>(events), polledEvents);
         assertEquals(0, accumulator.size());
 
@@ -126,27 +126,27 @@ public void testKeyConcurrentAndOrderingGuarantees() {
         MockEvent event = null;
 
         // Poll event0.
-        event = accumulator.poll(0, TimeUnit.MICROSECONDS);
+        event = accumulator.poll();
         assertEquals(event0, event);
 
         // Poll returns null because key is inflight.
-        assertNull(accumulator.poll(0, TimeUnit.MICROSECONDS));
+        assertNull(accumulator.poll());
         accumulator.done(event);
 
         // Poll event1.
-        event = accumulator.poll(0, TimeUnit.MICROSECONDS);
+        event = accumulator.poll();
         assertEquals(event1, event);
 
         // Poll returns null because key is inflight.
-        assertNull(accumulator.poll(0, TimeUnit.MICROSECONDS));
+        assertNull(accumulator.poll());
         accumulator.done(event);
 
         // Poll event2.
-        event = accumulator.poll(0, TimeUnit.MICROSECONDS);
+        event = accumulator.poll();
         assertEquals(event2, event);
 
         // Poll returns null because key is inflight.
-        assertNull(accumulator.poll(0, TimeUnit.MICROSECONDS));
+        assertNull(accumulator.poll());
         accumulator.done(event);
 
         accumulator.close();
@@ -160,9 +160,9 @@ public void testDoneUnblockWaitingThreads() throws ExecutionException, Interrupt
         MockEvent event1 = new MockEvent(1, 1);
         MockEvent event2 = new MockEvent(1, 2);
 
-        CompletableFuture<MockEvent> future0 = CompletableFuture.supplyAsync(accumulator::poll);
-        CompletableFuture<MockEvent> future1 = CompletableFuture.supplyAsync(accumulator::poll);
-        CompletableFuture<MockEvent> future2 = CompletableFuture.supplyAsync(accumulator::poll);
+        CompletableFuture<MockEvent> future0 = CompletableFuture.supplyAsync(accumulator::take);
+        CompletableFuture<MockEvent> future1 = CompletableFuture.supplyAsync(accumulator::take);
+        CompletableFuture<MockEvent> future2 = CompletableFuture.supplyAsync(accumulator::take);
         List<CompletableFuture<MockEvent>> futures = Arrays.asList(future0, future1, future2);
 
         assertFalse(future0.isDone());
@@ -215,9 +215,9 @@ public void testDoneUnblockWaitingThreads() throws ExecutionException, Interrupt
     public void testCloseUnblockWaitingThreads() throws ExecutionException, InterruptedException, TimeoutException {
         EventAccumulator<Integer, MockEvent> accumulator = new EventAccumulator<>();
 
-        CompletableFuture<MockEvent> future0 = CompletableFuture.supplyAsync(accumulator::poll);
-        CompletableFuture<MockEvent> future1 = CompletableFuture.supplyAsync(accumulator::poll);
-        CompletableFuture<MockEvent> future2 = CompletableFuture.supplyAsync(accumulator::poll);
+        CompletableFuture<MockEvent> future0 = CompletableFuture.supplyAsync(accumulator::take);
+        CompletableFuture<MockEvent> future1 = CompletableFuture.supplyAsync(accumulator::take);
+        CompletableFuture<MockEvent> future2 = CompletableFuture.supplyAsync(accumulator::take);
 
         assertFalse(future0.isDone());
         assertFalse(future1.isDone());
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
index 2714188f65e70..3708141827cc0 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
@@ -27,15 +27,12 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.LinkedList;
 import java.util.List;
-import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
@@ -53,53 +50,19 @@
 
 @Timeout(value = 60)
 public class MultiThreadedEventProcessorTest {
-    private static class MockEventAccumulator<T> extends EventAccumulator<TopicPartition, CoordinatorEvent> {
+    private static class DelayEventAccumulator extends EventAccumulator<TopicPartition, CoordinatorEvent> {
         private final Time time;
-        private final Queue<CoordinatorEvent> events;
-        private final long timeToPollMs;
-        private final AtomicBoolean isClosed;
+        private final long takeDelayMs;
 
-        public MockEventAccumulator(Time time, long timeToPollMs) {
+        public DelayEventAccumulator(Time time, long takeDelayMs) {
             this.time = time;
-            this.events = new LinkedList<>();
-            this.timeToPollMs = timeToPollMs;
-            this.isClosed = new AtomicBoolean(false);
+            this.takeDelayMs = takeDelayMs;
         }
 
         @Override
-        public CoordinatorEvent poll() {
-            synchronized (events) {
-                while (events.isEmpty() && !isClosed.get()) {
-                    try {
-                        events.wait();
-                    } catch (Exception ignored) {
-                        
-                    }
-                }
-                time.sleep(timeToPollMs);
-                return events.poll();
-            }
-        }
-
-        @Override
-        public CoordinatorEvent poll(long timeout, TimeUnit unit) {
-            return null;
-        }
-
-        @Override
-        public void add(CoordinatorEvent event) throws RejectedExecutionException {
-            synchronized (events) {
-                events.add(event);
-                events.notifyAll();
-            }
-        }
-
-        @Override
-        public void close() {
-            isClosed.set(true);
-            synchronized (events) {
-                events.notifyAll();
-            }
+        public CoordinatorEvent take() {
+            time.sleep(takeDelayMs);
+            return super.take();
         }
     }
 
@@ -353,7 +316,11 @@ public void testEventsAreDrainedWhenClosed() throws Exception {
             AtomicInteger numEventsExecuted = new AtomicInteger(0);
 
             // Special event which blocks until the latch is released.
-            FutureEvent<Integer> blockingEvent = new FutureEvent<>(new TopicPartition("foo", 0), numEventsExecuted::incrementAndGet, true);
+            FutureEvent<Integer> blockingEvent = new FutureEvent<>(
+                new TopicPartition("foo", 0),
+                numEventsExecuted::incrementAndGet,
+                true
+            );
 
             List<FutureEvent<Integer>> events = Arrays.asList(
                 new FutureEvent<>(new TopicPartition("foo", 0), numEventsExecuted::incrementAndGet),
@@ -428,7 +395,7 @@ public void testMetrics() throws Exception {
             1, // Use a single thread to block event in the processor.
             mockTime,
             mockRuntimeMetrics,
-            new MockEventAccumulator<>(mockTime, 500L)
+            new DelayEventAccumulator(mockTime, 500L)
         )) {
             // Enqueue the blocking event.
             eventProcessor.enqueue(blockingEvent);
@@ -501,7 +468,7 @@ public void testRecordThreadIdleRatioTwoThreads() throws Exception {
             2,
             Time.SYSTEM,
             mockRuntimeMetrics,
-            new MockEventAccumulator<>(Time.SYSTEM, 100L)
+            new DelayEventAccumulator(Time.SYSTEM, 100L)
         )) {
             List<Double> recordedRatios = new ArrayList<>();
             AtomicInteger numEventsExecuted = new AtomicInteger(0);

From 1bb9a851744972d6b2ce534fd75f03b18a61974d Mon Sep 17 00:00:00 2001
From: John Yu <54207775+chiacyu@users.noreply.github.com>
Date: Thu, 29 Feb 2024 08:14:35 +0800
Subject: [PATCH 089/258] MINOR: Remove the space between two words (#15439)

Remove the space between two words

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../kafka/server/log/remote/storage/RemoteLogMetadata.java    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/storage/api/src/main/java/org/apache/kafka/server/log/remote/storage/RemoteLogMetadata.java b/storage/api/src/main/java/org/apache/kafka/server/log/remote/storage/RemoteLogMetadata.java
index 74d5c3d28a796..4d8eaa95bcb97 100644
--- a/storage/api/src/main/java/org/apache/kafka/server/log/remote/storage/RemoteLogMetadata.java
+++ b/storage/api/src/main/java/org/apache/kafka/server/log/remote/storage/RemoteLogMetadata.java
@@ -32,7 +32,7 @@ public abstract class RemoteLogMetadata {
     private final int brokerId;
 
     /**
-     * Epoch time in milli seconds at which this event is generated.
+     * Epoch time in milliseconds at which this event is generated.
      */
     private final long eventTimestampMs;
 
@@ -42,7 +42,7 @@ protected RemoteLogMetadata(int brokerId, long eventTimestampMs) {
     }
 
     /**
-     * @return Epoch time in milli seconds at which this event is occurred.
+     * @return Epoch time in milliseconds at which this event is occurred.
      */
     public long eventTimestampMs() {
         return eventTimestampMs;

From 55a6d30ccbe971f4d2e99aeb3b1a773ffe5792a2 Mon Sep 17 00:00:00 2001
From: Christo Lolov <lolovc@amazon.com>
Date: Thu, 29 Feb 2024 00:19:55 +0000
Subject: [PATCH 090/258] KAFKA-16154: Broker returns offset for
 LATEST_TIERED_TIMESTAMP (#15213)

This is the first part of the implementation of KIP-1005

The purpose of this pull request is for the broker to start returning the correct offset when it receives a -5 as a timestamp in a ListOffsets API request

Reviewers: Luke Chen <showuon@gmail.com>, Kamal Chandraprakash <kamal.chandraprakash@gmail.com>, Satish Duggana <satishd@apache.org>
---
 .../common/requests/ListOffsetsRequest.java   |  2 +
 .../src/main/scala/kafka/log/UnifiedLog.scala | 56 ++++++++----
 .../scala/unit/kafka/log/UnifiedLogTest.scala | 86 +++++++++++++++++++
 3 files changed, 127 insertions(+), 17 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetsRequest.java
index efdc7da2afe1e..fc996453d6470 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListOffsetsRequest.java
@@ -47,6 +47,8 @@ public class ListOffsetsRequest extends AbstractRequest {
      */
     public static final long EARLIEST_LOCAL_TIMESTAMP = -4L;
 
+    public static final long LATEST_TIERED_TIMESTAMP = -5L;
+
     public static final int CONSUMER_REPLICA_ID = -1;
     public static final int DEBUGGING_REPLICA_ID = -2;
 
diff --git a/core/src/main/scala/kafka/log/UnifiedLog.scala b/core/src/main/scala/kafka/log/UnifiedLog.scala
index c0bb9d8cd66c6..f6198ce9d219e 100644
--- a/core/src/main/scala/kafka/log/UnifiedLog.scala
+++ b/core/src/main/scala/kafka/log/UnifiedLog.scala
@@ -41,7 +41,7 @@ import org.apache.kafka.server.record.BrokerCompressionType
 import org.apache.kafka.server.util.Scheduler
 import org.apache.kafka.storage.internals.checkpoint.{LeaderEpochCheckpointFile, PartitionMetadataFile}
 import org.apache.kafka.storage.internals.epoch.LeaderEpochFileCache
-import org.apache.kafka.storage.internals.log.{AbortedTxn, AppendOrigin, BatchMetadata, CompletedTxn, EpochEntry, FetchDataInfo, FetchIsolation, LastRecord, LeaderHwChange, LogAppendInfo, LogConfig, LogDirFailureChannel, LogFileUtils, LogOffsetMetadata, LogOffsetSnapshot, LogOffsetsListener, LogSegment, LogSegments, LogStartOffsetIncrementReason, LogValidator, ProducerAppendInfo, ProducerStateManager, ProducerStateManagerConfig, RollParams, VerificationGuard}
+import org.apache.kafka.storage.internals.log.{AbortedTxn, AppendOrigin, BatchMetadata, CompletedTxn, FetchDataInfo, FetchIsolation, LastRecord, LeaderHwChange, LogAppendInfo, LogConfig, LogDirFailureChannel, LogFileUtils, LogOffsetMetadata, LogOffsetSnapshot, LogOffsetsListener, LogSegment, LogSegments, LogStartOffsetIncrementReason, LogValidator, ProducerAppendInfo, ProducerStateManager, ProducerStateManagerConfig, RollParams, VerificationGuard}
 
 import java.io.{File, IOException}
 import java.nio.file.{Files, Path}
@@ -150,7 +150,9 @@ class UnifiedLog(@volatile var logStartOffset: Long,
   def localLogStartOffset(): Long = _localLogStartOffset
 
   // This is the offset(inclusive) until which segments are copied to the remote storage.
-  @volatile private var highestOffsetInRemoteStorage: Long = -1L
+  @volatile private[kafka] var _highestOffsetInRemoteStorage: Long = -1L
+
+  def highestOffsetInRemoteStorage(): Long = _highestOffsetInRemoteStorage
 
   locally {
     def updateLocalLogStartOffset(offset: Long): Unit = {
@@ -544,8 +546,8 @@ class UnifiedLog(@volatile var logStartOffset: Long,
 
   def updateHighestOffsetInRemoteStorage(offset: Long): Unit = {
     if (!remoteLogEnabled())
-      warn(s"Unable to update the highest offset in remote storage with offset $offset since remote storage is not enabled. The existing highest offset is $highestOffsetInRemoteStorage.")
-    else if (offset > highestOffsetInRemoteStorage) highestOffsetInRemoteStorage = offset
+      warn(s"Unable to update the highest offset in remote storage with offset $offset since remote storage is not enabled. The existing highest offset is ${highestOffsetInRemoteStorage()}.")
+    else if (offset > highestOffsetInRemoteStorage()) _highestOffsetInRemoteStorage = offset
   }
 
   // Rebuild producer state until lastOffset. This method may be called from the recovery code path, and thus must be
@@ -1279,7 +1281,6 @@ class UnifiedLog(@volatile var logStartOffset: Long,
 
       if (config.messageFormatVersion.isLessThan(IBP_0_10_0_IV0) &&
         targetTimestamp != ListOffsetsRequest.EARLIEST_TIMESTAMP &&
-        targetTimestamp != ListOffsetsRequest.EARLIEST_LOCAL_TIMESTAMP &&
         targetTimestamp != ListOffsetsRequest.LATEST_TIMESTAMP)
         throw new UnsupportedForMessageFormatException(s"Cannot search offsets based on timestamp because message format version " +
           s"for partition $topicPartition is ${config.messageFormatVersion} which is earlier than the minimum " +
@@ -1300,18 +1301,39 @@ class UnifiedLog(@volatile var logStartOffset: Long,
       } else if (targetTimestamp == ListOffsetsRequest.EARLIEST_LOCAL_TIMESTAMP) {
         val curLocalLogStartOffset = localLogStartOffset()
 
-        val earliestLocalLogEpochEntry = leaderEpochCache.asJava.flatMap(cache => {
-          val epoch = cache.epochForOffset(curLocalLogStartOffset)
-          if (epoch.isPresent) cache.epochEntry(epoch.getAsInt) else Optional.empty[EpochEntry]()
-        })
-
-        val epochOpt = if (earliestLocalLogEpochEntry.isPresent && earliestLocalLogEpochEntry.get().startOffset <= curLocalLogStartOffset)
-          Optional.of[Integer](earliestLocalLogEpochEntry.get().epoch)
-        else Optional.empty[Integer]()
+        val epochResult: Optional[Integer] =
+          if (leaderEpochCache.isDefined) {
+            val epochOpt = leaderEpochCache.get.epochForOffset(curLocalLogStartOffset)
+            if (epochOpt.isPresent) Optional.of(epochOpt.getAsInt) else Optional.empty()
+          } else {
+            Optional.empty()
+          }
 
-        Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, curLocalLogStartOffset, epochOpt))
+        Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, curLocalLogStartOffset, epochResult))
       } else if (targetTimestamp == ListOffsetsRequest.LATEST_TIMESTAMP) {
         Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, logEndOffset, latestEpochAsOptional(leaderEpochCache)))
+      } else if (targetTimestamp == ListOffsetsRequest.LATEST_TIERED_TIMESTAMP) {
+        if (remoteLogEnabled()) {
+          val curHighestRemoteOffset = highestOffsetInRemoteStorage()
+
+          val epochResult: Optional[Integer] =
+            if (leaderEpochCache.isDefined) {
+              val epochOpt = leaderEpochCache.get.epochForOffset(curHighestRemoteOffset)
+              if (epochOpt.isPresent) {
+                Optional.of(epochOpt.getAsInt)
+              } else if (curHighestRemoteOffset == -1) {
+                Optional.of(RecordBatch.NO_PARTITION_LEADER_EPOCH)
+              } else {
+                Optional.empty()
+              }
+            } else {
+              Optional.empty()
+            }
+
+          Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, curHighestRemoteOffset, epochResult))
+        } else {
+          Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, -1L, Optional.of(-1)))
+        }
       } else if (targetTimestamp == ListOffsetsRequest.MAX_TIMESTAMP) {
         // Cache to avoid race conditions. `toBuffer` is faster than most alternatives and provides
         // constant time access while being safe to use with concurrent collections unlike `toArray`.
@@ -1448,7 +1470,7 @@ class UnifiedLog(@volatile var logStartOffset: Long,
       //    1. they are uploaded to the remote storage
       //    2. log-start-offset was incremented higher than the largest offset in the candidate segment
       if (remoteLogEnabled()) {
-        (upperBoundOffset > 0 && upperBoundOffset - 1 <= highestOffsetInRemoteStorage) ||
+        (upperBoundOffset > 0 && upperBoundOffset - 1 <= highestOffsetInRemoteStorage()) ||
           allowDeletionDueToLogStartOffsetIncremented
       } else {
         true
@@ -1582,13 +1604,13 @@ class UnifiedLog(@volatile var logStartOffset: Long,
    * The log size in bytes for all segments that are only in local log but not yet in remote log.
    */
   def onlyLocalLogSegmentsSize: Long =
-    UnifiedLog.sizeInBytes(logSegments.stream.filter(_.baseOffset >= highestOffsetInRemoteStorage).collect(Collectors.toList[LogSegment]))
+    UnifiedLog.sizeInBytes(logSegments.stream.filter(_.baseOffset >= highestOffsetInRemoteStorage()).collect(Collectors.toList[LogSegment]))
 
   /**
    * The number of segments that are only in local log but not yet in remote log.
    */
   def onlyLocalLogSegmentsCount: Long =
-    logSegments.stream().filter(_.baseOffset >= highestOffsetInRemoteStorage).count()
+    logSegments.stream().filter(_.baseOffset >= highestOffsetInRemoteStorage()).count()
 
   /**
    * The offset of the next message that will be appended to the log
diff --git a/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala b/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
index 7b92a9e2df794..ffa585a460916 100755
--- a/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
+++ b/core/src/test/scala/unit/kafka/log/UnifiedLogTest.scala
@@ -2127,6 +2127,92 @@ class UnifiedLogTest {
       log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIMESTAMP, Some(remoteLogManager)))
   }
 
+  @Test
+  def testFetchLatestTieredTimestampNoRemoteStorage(): Unit = {
+    val logConfig = LogTestUtils.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
+    val log = createLog(logDir, logConfig)
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, -1, Optional.of(-1))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIERED_TIMESTAMP))
+
+    val firstTimestamp = mockTime.milliseconds
+    val leaderEpoch = 0
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = firstTimestamp),
+      leaderEpoch = leaderEpoch)
+
+    val secondTimestamp = firstTimestamp + 1
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = secondTimestamp),
+      leaderEpoch = leaderEpoch)
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, -1, Optional.of(-1))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIERED_TIMESTAMP))
+  }
+
+  @Test
+  def testFetchLatestTieredTimestampWithRemoteStorage(): Unit = {
+    val remoteLogManager = mock(classOf[RemoteLogManager])
+    val logConfig = LogTestUtils.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1,
+      remoteLogStorageEnable = true)
+    val log = createLog(logDir, logConfig, remoteStorageSystemEnable = true, remoteLogManager = Some(remoteLogManager))
+    when(remoteLogManager.findOffsetByTimestamp(log.topicPartition, 0, 0, log.leaderEpochCache.get))
+      .thenReturn(Optional.empty[TimestampAndOffset]())
+
+    assertEquals(None, log.fetchOffsetByTimestamp(0L, Some(remoteLogManager)))
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 0, Optional.empty())),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.EARLIEST_LOCAL_TIMESTAMP, Some(remoteLogManager)))
+
+    val firstTimestamp = mockTime.milliseconds
+    val firstLeaderEpoch = 0
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = firstTimestamp),
+      leaderEpoch = firstLeaderEpoch)
+
+    val secondTimestamp = firstTimestamp + 1
+    val secondLeaderEpoch = 1
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = secondTimestamp),
+      leaderEpoch = secondLeaderEpoch)
+
+    when(remoteLogManager.findOffsetByTimestamp(ArgumentMatchers.eq(log.topicPartition),
+      anyLong(), anyLong(), ArgumentMatchers.eq(log.leaderEpochCache.get)))
+      .thenAnswer(ans => {
+        val timestamp = ans.getArgument(1).asInstanceOf[Long]
+        Optional.of(timestamp)
+          .filter(_ == firstTimestamp)
+          .map[TimestampAndOffset](x => new TimestampAndOffset(x, 0L, Optional.of(firstLeaderEpoch)))
+      })
+    log._localLogStartOffset = 1
+    log._highestOffsetInRemoteStorage = 0
+
+    // In the assertions below we test that offset 0 (first timestamp) is in remote and offset 1 (second timestamp) is in local storage.
+    assertEquals(Some(new TimestampAndOffset(firstTimestamp, 0L, Optional.of(firstLeaderEpoch))),
+      log.fetchOffsetByTimestamp(firstTimestamp, Some(remoteLogManager)))
+    assertEquals(Some(new TimestampAndOffset(secondTimestamp, 1L, Optional.of(secondLeaderEpoch))),
+      log.fetchOffsetByTimestamp(secondTimestamp, Some(remoteLogManager)))
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 0L, Optional.of(firstLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.EARLIEST_TIMESTAMP, Some(remoteLogManager)))
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 0L, Optional.of(firstLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIERED_TIMESTAMP, Some(remoteLogManager)))
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 1L, Optional.of(secondLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.EARLIEST_LOCAL_TIMESTAMP, Some(remoteLogManager)))
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 2L, Optional.of(secondLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIMESTAMP, Some(remoteLogManager)))
+
+    // The cache can be updated directly after a leader change.
+    // The new latest offset should reflect the updated epoch.
+    log.maybeAssignEpochStartOffset(2, 2L)
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetsResponse.UNKNOWN_TIMESTAMP, 2L, Optional.of(2))),
+      log.fetchOffsetByTimestamp(ListOffsetsRequest.LATEST_TIMESTAMP, Some(remoteLogManager)))
+  }
+
   /**
    * Test the Log truncate operations
    */

From 96c68096a26ea5e7c2333308dfbaef47cb1eac72 Mon Sep 17 00:00:00 2001
From: Ritika Reddy <98577846+rreddy-22@users.noreply.github.com>
Date: Thu, 29 Feb 2024 00:38:42 -0800
Subject: [PATCH 091/258] KAFKA-15462: Add Group Type Filter for List Group to
 the Admin Client  (#15150)

In KIP-848, we introduce the notion of Group Types based on the protocol type that the members in the consumer group use. As of now we support two types of groups:
* Classic : Members use the classic consumer group protocol ( existing one )
* Consumer : Members use the consumer group protocol introduced in KIP-848.
Currently List Groups allows users to list all the consumer groups available. KIP-518 introduced filtering the consumer groups by the state that they are in. We now want to allow users to filter consumer groups by type.

This patch includes the changes to the admin client and related files. It also includes changes to parameterize the tests to include permutations of the old GC and the new GC with the different protocol types.

Reviewers: David Jacot <djacot@confluent.io>
---
 checkstyle/suppressions.xml                   |   1 +
 .../clients/admin/ConsumerGroupListing.java   |  64 +--
 .../kafka/clients/admin/KafkaAdminClient.java |  20 +-
 .../admin/ListConsumerGroupsOptions.java      |  25 +-
 .../clients/admin/KafkaAdminClientTest.java   | 124 +++++-
 .../kafka/admin/ConsumerGroupCommand.scala    | 107 +++--
 .../kafka/api/BaseConsumerTest.scala          |  19 +-
 .../admin/ConsumerGroupCommandTest.scala      |   7 +-
 .../apache/kafka/tools/ToolsTestUtils.java    |   2 +
 .../group/ConsumerGroupCommandTest.java       |  13 +-
 .../consumer/group/ListConsumerGroupTest.java | 386 ++++++++++++++++--
 11 files changed, 669 insertions(+), 99 deletions(-)

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 7486ef9a80d3e..c65cd675a9ef3 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -46,6 +46,7 @@
 
     <!-- server tests -->
     <suppress checks="MethodLength|JavaNCSS|NPath" files="DescribeTopicPartitionsRequestHandlerTest.java"/>
+    <suppress checks="CyclomaticComplexity" files="ListConsumerGroupTest.java"/>
 
     <!-- Clients -->
     <suppress id="dontUseSystemExit"
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java
index 0abc3e01ca9de..01c23796d41ae 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java
@@ -21,6 +21,7 @@
 import java.util.Optional;
 
 import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.GroupType;
 
 /**
  * A listing of a consumer group in the cluster.
@@ -29,6 +30,7 @@ public class ConsumerGroupListing {
     private final String groupId;
     private final boolean isSimpleConsumerGroup;
     private final Optional<ConsumerGroupState> state;
+    private final Optional<GroupType> type;
 
     /**
      * Create an instance with the specified parameters.
@@ -37,7 +39,7 @@ public class ConsumerGroupListing {
      * @param isSimpleConsumerGroup If consumer group is simple or not.
      */
     public ConsumerGroupListing(String groupId, boolean isSimpleConsumerGroup) {
-        this(groupId, isSimpleConsumerGroup, Optional.empty());
+        this(groupId, isSimpleConsumerGroup, Optional.empty(), Optional.empty());
     }
 
     /**
@@ -48,9 +50,27 @@ public ConsumerGroupListing(String groupId, boolean isSimpleConsumerGroup) {
      * @param state The state of the consumer group
      */
     public ConsumerGroupListing(String groupId, boolean isSimpleConsumerGroup, Optional<ConsumerGroupState> state) {
+        this(groupId, isSimpleConsumerGroup, state, Optional.empty());
+    }
+
+    /**
+     * Create an instance with the specified parameters.
+     *
+     * @param groupId                   Group Id.
+     * @param isSimpleConsumerGroup     If consumer group is simple or not.
+     * @param state                     The state of the consumer group.
+     * @param type                      The type of the consumer group.
+     */
+    public ConsumerGroupListing(
+        String groupId,
+        boolean isSimpleConsumerGroup,
+        Optional<ConsumerGroupState> state,
+        Optional<GroupType> type
+    ) {
         this.groupId = groupId;
         this.isSimpleConsumerGroup = isSimpleConsumerGroup;
         this.state = Objects.requireNonNull(state);
+        this.type = Objects.requireNonNull(type);
     }
 
     /**
@@ -74,42 +94,38 @@ public Optional<ConsumerGroupState> state() {
         return state;
     }
 
+    /**
+     * The type of the consumer group.
+     *
+     * @return An Optional containing the type, if available.
+     */
+    public Optional<GroupType> type() {
+        return type;
+    }
+
     @Override
     public String toString() {
         return "(" +
             "groupId='" + groupId + '\'' +
             ", isSimpleConsumerGroup=" + isSimpleConsumerGroup +
             ", state=" + state +
+            ", type=" + type +
             ')';
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(groupId, isSimpleConsumerGroup, state);
+        return Objects.hash(groupId, isSimpleConsumerGroup(), state, type);
     }
 
     @Override
-    public boolean equals(Object obj) {
-        if (this == obj)
-            return true;
-        if (obj == null)
-            return false;
-        if (getClass() != obj.getClass())
-            return false;
-        ConsumerGroupListing other = (ConsumerGroupListing) obj;
-        if (groupId == null) {
-            if (other.groupId != null)
-                return false;
-        } else if (!groupId.equals(other.groupId))
-            return false;
-        if (isSimpleConsumerGroup != other.isSimpleConsumerGroup)
-            return false;
-        if (state == null) {
-            if (other.state != null)
-                return false;
-        } else if (!state.equals(other.state))
-            return false;
-        return true;
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof ConsumerGroupListing)) return false;
+        ConsumerGroupListing that = (ConsumerGroupListing) o;
+        return isSimpleConsumerGroup() == that.isSimpleConsumerGroup() &&
+            Objects.equals(groupId, that.groupId) &&
+            Objects.equals(state, that.state) &&
+            Objects.equals(type, that.type);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
index 85c82e2514437..d98ad8ac04e9d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
@@ -58,6 +58,7 @@
 import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.GroupType;
 import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.KafkaFuture;
@@ -3382,7 +3383,14 @@ ListGroupsRequest.Builder createRequest(int timeoutMs) {
                                     .stream()
                                     .map(ConsumerGroupState::toString)
                                     .collect(Collectors.toList());
-                            return new ListGroupsRequest.Builder(new ListGroupsRequestData().setStatesFilter(states));
+                            List<String> groupTypes = options.types()
+                                    .stream()
+                                    .map(GroupType::toString)
+                                    .collect(Collectors.toList());
+                            return new ListGroupsRequest.Builder(new ListGroupsRequestData()
+                                .setStatesFilter(states)
+                                .setTypesFilter(groupTypes)
+                            );
                         }
 
                         private void maybeAddConsumerGroup(ListGroupsResponseData.ListedGroup group) {
@@ -3392,7 +3400,15 @@ private void maybeAddConsumerGroup(ListGroupsResponseData.ListedGroup group) {
                                 final Optional<ConsumerGroupState> state = group.groupState().equals("")
                                         ? Optional.empty()
                                         : Optional.of(ConsumerGroupState.parse(group.groupState()));
-                                final ConsumerGroupListing groupListing = new ConsumerGroupListing(groupId, protocolType.isEmpty(), state);
+                                final Optional<GroupType> type = group.groupType().equals("")
+                                        ? Optional.empty()
+                                        : Optional.of(GroupType.parse(group.groupType()));
+                                final ConsumerGroupListing groupListing = new ConsumerGroupListing(
+                                        groupId,
+                                        protocolType.isEmpty(),
+                                        state,
+                                        type
+                                    );
                                 results.addListing(groupListing);
                             }
                         }
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsOptions.java
index 9f1f38dd4a8e6..c240da159ff77 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsOptions.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsOptions.java
@@ -22,6 +22,7 @@
 import java.util.Set;
 
 import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.GroupType;
 import org.apache.kafka.common.annotation.InterfaceStability;
 
 /**
@@ -34,20 +35,38 @@ public class ListConsumerGroupsOptions extends AbstractOptions<ListConsumerGroup
 
     private Set<ConsumerGroupState> states = Collections.emptySet();
 
+    private Set<GroupType> types = Collections.emptySet();
+
     /**
-     * If states is set, only groups in these states will be returned by listConsumerGroups()
+     * If states is set, only groups in these states will be returned by listConsumerGroups().
      * Otherwise, all groups are returned.
      * This operation is supported by brokers with version 2.6.0 or later.
      */
     public ListConsumerGroupsOptions inStates(Set<ConsumerGroupState> states) {
-        this.states = (states == null) ? Collections.emptySet() : new HashSet<>(states);
+        this.states = (states == null || states.isEmpty()) ? Collections.emptySet() : new HashSet<>(states);
         return this;
     }
 
     /**
-     * Returns the list of States that are requested or empty if no states have been specified
+     * If types is set, only groups of these types will be returned by listConsumerGroups().
+     * Otherwise, all groups are returned.
+     */
+    public ListConsumerGroupsOptions withTypes(Set<GroupType> types) {
+        this.types = (types == null || types.isEmpty()) ? Collections.emptySet() : new HashSet<>(types);
+        return this;
+    }
+
+    /**
+     * Returns the list of States that are requested or empty if no states have been specified.
      */
     public Set<ConsumerGroupState> states() {
         return states;
     }
+
+    /**
+     * Returns the list of group types that are requested or empty if no types have been specified.
+     */
+    public Set<GroupType> types() {
+        return types;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
index b8b3d54ef439a..43d391a220efc 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
@@ -243,6 +243,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
@@ -2811,6 +2812,68 @@ public void testListConsumerGroupsWithStates() throws Exception {
         }
     }
 
+    @Test
+    public void testListConsumerGroupsWithTypes() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Test with a specific state filter but no type filter in list consumer group options.
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+
+            env.kafkaClient().prepareResponseFrom(
+                expectListGroupsRequestWithFilters(singleton(ConsumerGroupState.STABLE.toString()), Collections.emptySet()),
+                new ListGroupsResponse(new ListGroupsResponseData()
+                    .setErrorCode(Errors.NONE.code())
+                    .setGroups(Arrays.asList(
+                        new ListGroupsResponseData.ListedGroup()
+                            .setGroupId("group-1")
+                            .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                            .setGroupState("Stable")
+                            .setGroupType(GroupType.CLASSIC.toString())))),
+                env.cluster().nodeById(0));
+
+            final ListConsumerGroupsOptions options = new ListConsumerGroupsOptions().inStates(singleton(ConsumerGroupState.STABLE));
+            final ListConsumerGroupsResult result = env.adminClient().listConsumerGroups(options);
+            Collection<ConsumerGroupListing> listings = result.valid().get();
+
+            assertEquals(1, listings.size());
+            List<ConsumerGroupListing> expected = new ArrayList<>();
+            expected.add(new ConsumerGroupListing("group-1", false, Optional.of(ConsumerGroupState.STABLE), Optional.of(GroupType.CLASSIC)));
+            assertEquals(expected, listings);
+            assertEquals(0, result.errors().get().size());
+
+            // Test with list consumer group options.
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+
+            env.kafkaClient().prepareResponseFrom(
+                expectListGroupsRequestWithFilters(Collections.emptySet(), singleton(GroupType.CONSUMER.toString())),
+                new ListGroupsResponse(new ListGroupsResponseData()
+                    .setErrorCode(Errors.NONE.code())
+                    .setGroups(Arrays.asList(
+                        new ListGroupsResponseData.ListedGroup()
+                            .setGroupId("group-1")
+                            .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                            .setGroupState("Stable")
+                            .setGroupType(GroupType.CONSUMER.toString()),
+                        new ListGroupsResponseData.ListedGroup()
+                            .setGroupId("group-2")
+                            .setGroupState("Empty")
+                            .setGroupType(GroupType.CONSUMER.toString())))),
+                env.cluster().nodeById(0));
+
+            final ListConsumerGroupsOptions options2 = new ListConsumerGroupsOptions().withTypes(singleton(GroupType.CONSUMER));
+            final ListConsumerGroupsResult result2 = env.adminClient().listConsumerGroups(options2);
+            Collection<ConsumerGroupListing> listings2 = result2.valid().get();
+
+            assertEquals(2, listings2.size());
+            List<ConsumerGroupListing> expected2 = new ArrayList<>();
+            expected2.add(new ConsumerGroupListing("group-2", true, Optional.of(ConsumerGroupState.EMPTY), Optional.of(GroupType.CONSUMER)));
+            expected2.add(new ConsumerGroupListing("group-1", false, Optional.of(ConsumerGroupState.STABLE), Optional.of(GroupType.CONSUMER)));
+            assertEquals(expected2, listings2);
+            assertEquals(0, result.errors().get().size());
+        }
+    }
+
     @Test
     public void testListConsumerGroupsWithStatesOlderBrokerVersion() throws Exception {
         ApiVersion listGroupV3 = new ApiVersion()
@@ -2835,7 +2898,7 @@ public void testListConsumerGroupsWithStatesOlderBrokerVersion() throws Exceptio
             ListConsumerGroupsResult result = env.adminClient().listConsumerGroups(options);
             Collection<ConsumerGroupListing> listing = result.all().get();
             assertEquals(1, listing.size());
-            List<ConsumerGroupListing> expected = Collections.singletonList(new ConsumerGroupListing("group-1", false, Optional.empty()));
+            List<ConsumerGroupListing> expected = Collections.singletonList(new ConsumerGroupListing("group-1", false));
             assertEquals(expected, listing);
 
             // But we cannot set a state filter with older broker
@@ -2849,6 +2912,65 @@ public void testListConsumerGroupsWithStatesOlderBrokerVersion() throws Exceptio
         }
     }
 
+    @Test
+    public void testListConsumerGroupsWithTypesOlderBrokerVersion() throws Exception {
+        ApiVersion listGroupV4 = new ApiVersion()
+            .setApiKey(ApiKeys.LIST_GROUPS.id)
+            .setMinVersion((short) 0)
+            .setMaxVersion((short) 4);
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create(Collections.singletonList(listGroupV4)));
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+
+            // Check if we can list groups with older broker if we specify states and don't specify types.
+            env.kafkaClient().prepareResponseFrom(
+                expectListGroupsRequestWithFilters(singleton(ConsumerGroupState.STABLE.toString()), Collections.emptySet()),
+                new ListGroupsResponse(new ListGroupsResponseData()
+                    .setErrorCode(Errors.NONE.code())
+                    .setGroups(Collections.singletonList(
+                        new ListGroupsResponseData.ListedGroup()
+                            .setGroupId("group-1")
+                            .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                            .setGroupState(ConsumerGroupState.STABLE.toString())))),
+                env.cluster().nodeById(0));
+
+            ListConsumerGroupsOptions options = new ListConsumerGroupsOptions().inStates(singleton(ConsumerGroupState.STABLE));
+            ListConsumerGroupsResult result = env.adminClient().listConsumerGroups(options);
+
+            Collection<ConsumerGroupListing> listing = result.all().get();
+            assertEquals(1, listing.size());
+            List<ConsumerGroupListing> expected = Collections.singletonList(
+                new ConsumerGroupListing("group-1", false, Optional.of(ConsumerGroupState.STABLE))
+            );
+            assertEquals(expected, listing);
+
+            // Check that we cannot set a type filter with an older broker.
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+            env.kafkaClient().prepareUnsupportedVersionResponse(request ->
+                request instanceof ListGroupsRequest && !((ListGroupsRequest) request).data().typesFilter().isEmpty()
+            );
+
+            options = new ListConsumerGroupsOptions().withTypes(singleton(GroupType.CLASSIC));
+            result = env.adminClient().listConsumerGroups(options);
+            TestUtils.assertFutureThrows(result.all(), UnsupportedVersionException.class);
+        }
+    }
+
+    private MockClient.RequestMatcher expectListGroupsRequestWithFilters(
+        Set<String> expectedStates,
+        Set<String> expectedTypes
+    ) {
+        return body -> {
+            if (body instanceof ListGroupsRequest) {
+                ListGroupsRequest request = (ListGroupsRequest) body;
+                return Objects.equals(new HashSet<>(request.data().statesFilter()), expectedStates)
+                    && Objects.equals(new HashSet<>(request.data().typesFilter()), expectedTypes);
+            }
+            return false;
+        };
+    }
+
     @Test
     public void testOffsetCommitNumRetries() throws Exception {
         final Cluster cluster = mockCluster(3, 0);
diff --git a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
index 4187274a22d6e..160b9a70aaec9 100755
--- a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
@@ -29,7 +29,7 @@ import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.consumer.OffsetAndMetadata
 import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.common.{KafkaException, Node, TopicPartition}
+import org.apache.kafka.common.{ConsumerGroupState, GroupType, KafkaException, Node, TopicPartition}
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 
 import scala.jdk.CollectionConverters._
@@ -41,7 +41,6 @@ import org.apache.kafka.common.protocol.Errors
 
 import scala.collection.immutable.TreeMap
 import scala.reflect.ClassTag
-import org.apache.kafka.common.ConsumerGroupState
 import org.apache.kafka.common.requests.ListOffsetsResponse
 
 object ConsumerGroupCommand extends Logging {
@@ -104,6 +103,15 @@ object ConsumerGroupCommand extends Logging {
     parsedStates
   }
 
+  def consumerGroupTypesFromString(input: String): Set[GroupType] = {
+    val parsedTypes = input.toLowerCase.split(',').map(s => GroupType.parse(s.trim)).toSet
+    if (parsedTypes.contains(GroupType.UNKNOWN)) {
+      val validTypes = GroupType.values().filter(_ != GroupType.UNKNOWN)
+      throw new IllegalArgumentException(s"Invalid types list '$input'. Valid types are: ${validTypes.mkString(", ")}")
+    }
+    parsedTypes
+  }
+
   val MISSING_COLUMN_VALUE = "-"
 
   private def printError(msg: String, e: Option[Throwable] = None): Unit = {
@@ -135,7 +143,7 @@ object ConsumerGroupCommand extends Logging {
   private[admin] case class MemberAssignmentState(group: String, consumerId: String, host: String, clientId: String, groupInstanceId: String,
                                              numPartitions: Int, assignment: List[TopicPartition])
 
-  case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
+  private[admin] case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
 
   private[admin] sealed trait CsvRecord
   private[admin] case class CsvRecordWithGroup(group: String, topic: String, partition: Int, offset: Long) extends CsvRecord
@@ -189,16 +197,65 @@ object ConsumerGroupCommand extends Logging {
     }
 
     def listGroups(): Unit = {
-      if (opts.options.has(opts.stateOpt)) {
-        val stateValue = opts.options.valueOf(opts.stateOpt)
-        val states = if (stateValue == null || stateValue.isEmpty)
-          Set[ConsumerGroupState]()
-        else
-          consumerGroupStatesFromString(stateValue)
-        val listings = listConsumerGroupsWithState(states)
-        printGroupStates(listings.map(e => (e.groupId, e.state.get.toString)))
-      } else
+      val includeType = opts.options.has(opts.typeOpt)
+      val includeState = opts.options.has(opts.stateOpt)
+
+      if (includeType || includeState) {
+        val types = typeValues()
+        val states = stateValues()
+        val listings = listConsumerGroupsWithFilters(types, states)
+
+        printGroupInfo(listings, includeType, includeState)
+
+      } else {
         listConsumerGroups().foreach(println(_))
+      }
+    }
+
+    private def stateValues(): Set[ConsumerGroupState] = {
+      val stateValue = opts.options.valueOf(opts.stateOpt)
+      if (stateValue == null || stateValue.isEmpty)
+        Set[ConsumerGroupState]()
+      else
+        consumerGroupStatesFromString(stateValue)
+    }
+
+    private def typeValues(): Set[GroupType] = {
+      val typeValue = opts.options.valueOf(opts.typeOpt)
+      if (typeValue == null || typeValue.isEmpty)
+        Set[GroupType]()
+      else
+        consumerGroupTypesFromString(typeValue)
+    }
+
+    private def printGroupInfo(groups: List[ConsumerGroupListing], includeType: Boolean, includeState: Boolean): Unit = {
+      def groupId(groupListing: ConsumerGroupListing): String = groupListing.groupId
+      def groupType(groupListing: ConsumerGroupListing): String = groupListing.`type`().orElse(GroupType.UNKNOWN).toString
+      def groupState(groupListing: ConsumerGroupListing): String = groupListing.state.orElse(ConsumerGroupState.UNKNOWN).toString
+
+      val maxGroupLen = groups.foldLeft(15)((maxLen, groupListing) => Math.max(maxLen, groupId(groupListing).length)) + 10
+      var format = s"%-${maxGroupLen}s"
+      var header = List("GROUP")
+      var extractors: List[ConsumerGroupListing => String] = List(groupId)
+
+      if (includeType) {
+        header = header :+ "TYPE"
+        extractors = extractors :+ groupType _
+        format += " %-20s"
+      }
+
+      if (includeState) {
+        header = header :+ "STATE"
+        extractors = extractors :+ groupState _
+        format += " %-20s"
+      }
+
+      println(format.format(header: _*))
+
+      groups.foreach { groupListing =>
+        val info = extractors.map(extractor => extractor(groupListing))
+        println(format.format(info: _*))
+      }
     }
 
     def listConsumerGroups(): List[String] = {
@@ -207,26 +264,15 @@ object ConsumerGroupCommand extends Logging {
       listings.map(_.groupId).toList
     }
 
-    def listConsumerGroupsWithState(states: Set[ConsumerGroupState]): List[ConsumerGroupListing] = {
+    def listConsumerGroupsWithFilters(types: Set[GroupType], states: Set[ConsumerGroupState]): List[ConsumerGroupListing] = {
       val listConsumerGroupsOptions = withTimeoutMs(new ListConsumerGroupsOptions())
-      listConsumerGroupsOptions.inStates(states.asJava)
+      listConsumerGroupsOptions
+        .inStates(states.asJava)
+        .withTypes(types.asJava)
       val result = adminClient.listConsumerGroups(listConsumerGroupsOptions)
       result.all.get.asScala.toList
     }
 
-    private def printGroupStates(groupsAndStates: List[(String, String)]): Unit = {
-      // find proper columns width
-      var maxGroupLen = 15
-      for ((groupId, _) <- groupsAndStates) {
-        maxGroupLen = Math.max(maxGroupLen, groupId.length)
-      }
-      val format = s"%${-maxGroupLen}s %s"
-      println(format.format("GROUP", "STATE"))
-      for ((groupId, state) <- groupsAndStates) {
-        println(format.format(groupId, state))
-      }
-    }
-
     private def shouldPrintMemberState(group: String, state: Option[String], numRows: Option[Int]): Boolean = {
       // numRows contains the number of data rows, if any, compiled from the API call in the caller method.
       // if it's undefined or 0, there is no relevant group information to display.
@@ -1024,6 +1070,9 @@ object ConsumerGroupCommand extends Logging {
       "When specified with '--list', it displays the state of all groups. It can also be used to list groups with specific states." + nl +
       "Example: --bootstrap-server localhost:9092 --list --state stable,empty" + nl +
       "This option may be used with '--describe', '--list' and '--bootstrap-server' options only."
+    private val TypeDoc = "When specified with '--list', it displays the types of all the groups. It can also be used to list groups with specific types." + nl +
+      "Example: --bootstrap-server localhost:9092 --list --type classic,consumer" + nl +
+      "This option may be used with the '--list' option only."
     private val DeleteOffsetsDoc = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics."
 
     val bootstrapServerOpt: OptionSpec[String] = parser.accepts("bootstrap-server", BootstrapServerDoc)
@@ -1090,6 +1139,10 @@ object ConsumerGroupCommand extends Logging {
                          .availableIf(describeOpt, listOpt)
                          .withOptionalArg()
                          .ofType(classOf[String])
+    val typeOpt: OptionSpec[String] = parser.accepts("type", TypeDoc)
+                        .availableIf(listOpt)
+                        .withOptionalArg()
+                        .ofType(classOf[String])
 
     options = parser.parse(args : _*)
 
diff --git a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
index bb3259baf98bd..20159830943eb 100644
--- a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
@@ -26,6 +26,7 @@ import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.{Arguments, MethodSource}
 
+import java.util
 import java.util.Properties
 import java.util.concurrent.atomic.AtomicInteger
 import scala.jdk.CollectionConverters._
@@ -117,11 +118,12 @@ object BaseConsumerTest {
   // * KRaft with the new group coordinator enabled and the classic group protocol
   // * KRaft with the new group coordinator enabled and the consumer group protocol
   def getTestQuorumAndGroupProtocolParametersAll() : java.util.stream.Stream[Arguments] = {
-    java.util.stream.Stream.of(
+    util.Arrays.stream(Array(
         Arguments.of("zk", "classic"),
         Arguments.of("kraft", "classic"),
         Arguments.of("kraft+kip848", "classic"),
-        Arguments.of("kraft+kip848", "consumer"))
+        Arguments.of("kraft+kip848", "consumer")
+    ))
   }
 
   // In Scala 2.12, it is necessary to disambiguate the java.util.stream.Stream.of() method call
@@ -138,10 +140,19 @@ object BaseConsumerTest {
   // * KRaft and the classic group protocol
   // * KRaft with the new group coordinator enabled and the classic group protocol
   def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly() : java.util.stream.Stream[Arguments] = {
-    java.util.stream.Stream.of(
+    util.Arrays.stream(Array(
         Arguments.of("zk", "classic"),
         Arguments.of("kraft", "classic"),
-        Arguments.of("kraft+kip848", "classic"))
+        Arguments.of("kraft+kip848", "classic")
+    ))
+  }
+
+  // For tests that only work with the consumer group protocol, we want to test the following combination:
+  // * KRaft with the new group coordinator enabled and the consumer group protocol
+  def getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly(): java.util.stream.Stream[Arguments] = {
+    util.Arrays.stream(Array(
+        Arguments.of("kraft+kip848", "consumer")
+    ))
   }
 
   val updateProducerCount = new AtomicInteger()
diff --git a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
index 18c7a0a8f8142..f682df1f1dcdb 100644
--- a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
@@ -19,7 +19,7 @@ package kafka.admin
 
 import java.time.Duration
 import java.util.concurrent.{ExecutorService, Executors, TimeUnit}
-import java.util.{Collections, Properties}
+import java.util.{Collections, Properties, stream}
 import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
 import kafka.api.BaseConsumerTest
 import kafka.integration.KafkaServerTestHarness
@@ -31,6 +31,7 @@ import org.apache.kafka.common.{PartitionInfo, TopicPartition}
 import org.apache.kafka.common.errors.WakeupException
 import org.apache.kafka.common.serialization.StringDeserializer
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo}
+import org.junit.jupiter.params.provider.Arguments
 
 import scala.jdk.CollectionConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -122,7 +123,9 @@ class ConsumerGroupCommandTest extends KafkaServerTestHarness {
 }
 
 object ConsumerGroupCommandTest {
-   def getTestQuorumAndGroupProtocolParametersAll() = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+   def getTestQuorumAndGroupProtocolParametersAll(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+   def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly()
+   def getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly()
 
   abstract class AbstractConsumerRunnable(broker: String, groupId: String, customPropsOpt: Option[Properties] = None,
                                           syncCommit: Boolean = false) extends Runnable {
diff --git a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
index fdc732ea29ab1..83fa31bf5e795 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
@@ -45,6 +45,8 @@
 public class ToolsTestUtils {
     /** @see TestInfoUtils#TestWithParameterizedQuorumName()  */
     public static final String TEST_WITH_PARAMETERIZED_QUORUM_NAME = "{displayName}.{argumentsWithNames}";
+    /** @see TestInfoUtils#TestWithParameterizedQuorumAndGroupProtocolNames()  */
+    public static final String TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES = "{displayName}.quorum={0}.groupProtocol={1}";
 
     private static int randomPort = 0;
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
index b78054cb4adaa..bde3af37a1d70 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
@@ -58,6 +58,7 @@
 public class ConsumerGroupCommandTest extends kafka.integration.KafkaServerTestHarness {
     public static final String TOPIC = "foo";
     public static final String GROUP = "test.group";
+    public static final String PROTOCOL_GROUP = "protocol-group";
 
     List<ConsumerGroupCommand.ConsumerGroupService> consumerGroupService = new ArrayList<>();
     List<AbstractConsumerGroupExecutor> consumerGroupExecutors = new ArrayList<>();
@@ -154,8 +155,8 @@ ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String groupPro
         return addConsumerGroupExecutor(numConsumers, TOPIC, GROUP, RangeAssignor.class.getName(), remoteAssignor, Optional.empty(), false, groupProtocol);
     }
 
-    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String topic, String group) {
-        return addConsumerGroupExecutor(numConsumers, topic, group, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, GroupProtocol.CLASSIC.name);
+    ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String group, String groupProtocol) {
+        return addConsumerGroupExecutor(numConsumers, TOPIC, group, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), false, groupProtocol);
     }
 
     ConsumerGroupExecutor addConsumerGroupExecutor(int numConsumers, String topic, String group, String groupProtocol) {
@@ -342,6 +343,14 @@ public static Stream<Arguments> getTestQuorumAndGroupProtocolParametersAll() {
         return BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll();
     }
 
+    public static Stream<Arguments> getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly() {
+        return BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly();
+    }
+
+    public static Stream<Arguments> getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly() {
+        return BaseConsumerTest.getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly();
+    }
+
     @SuppressWarnings({"deprecation"})
     static <T> Seq<T> seq(Collection<T> seq) {
         return JavaConverters.asScalaIteratorConverter(seq.iterator()).asScala().toSeq();
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
index 894f00df5e7f6..ba5ebd254fc93 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
@@ -20,83 +20,258 @@
 import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.admin.ConsumerGroupListing;
 import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.GroupType;
 import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.ValueSource;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Properties;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class ListConsumerGroupTest extends ConsumerGroupCommandTest {
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
-    public void testListConsumerGroups(String quorum) throws Exception {
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersAll")
+    public void testListConsumerGroupsWithoutFilters(String quorum, String groupProtocol) throws Exception {
         String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
         addSimpleGroupExecutor(simpleGroup);
         addConsumerGroupExecutor(1);
+        addConsumerGroupExecutor(1, PROTOCOL_GROUP, groupProtocol);
 
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
-        scala.collection.Set<String> expectedGroups = set(Arrays.asList(GROUP, simpleGroup));
+
+        scala.collection.Set<String> expectedGroups = set(Arrays.asList(GROUP, simpleGroup, PROTOCOL_GROUP));
         final AtomicReference<scala.collection.Set> foundGroups = new AtomicReference<>();
+
         TestUtils.waitForCondition(() -> {
             foundGroups.set(service.listConsumerGroups().toSet());
             return Objects.equals(expectedGroups, foundGroups.get());
         }, "Expected --list to show groups " + expectedGroups + ", but found " + foundGroups.get() + ".");
     }
 
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
+    @Test
     public void testListWithUnrecognizedNewConsumerOption() {
         String[] cgcArgs = new String[]{"--new-consumer", "--bootstrap-server", bootstrapServers(listenerName()), "--list"};
         assertThrows(OptionException.class, () -> getConsumerGroupService(cgcArgs));
     }
 
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
-    public void testListConsumerGroupsWithStates() throws Exception {
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersAll")
+    public void testListConsumerGroupsWithStates(String quorum, String groupProtocol) throws Exception {
         String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
         addSimpleGroupExecutor(simpleGroup);
-        addConsumerGroupExecutor(1);
+        addConsumerGroupExecutor(1, groupProtocol);
 
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        scala.collection.Set<ConsumerGroupListing> expectedListing = set(Arrays.asList(
-            new ConsumerGroupListing(simpleGroup, true, Optional.of(ConsumerGroupState.EMPTY)),
-            new ConsumerGroupListing(GROUP, false, Optional.of(ConsumerGroupState.STABLE))));
+        Set<ConsumerGroupListing> expectedListing = mkSet(
+            new ConsumerGroupListing(
+                simpleGroup,
+                true,
+                Optional.of(ConsumerGroupState.EMPTY),
+                Optional.of(GroupType.CLASSIC)
+            ),
+            new ConsumerGroupListing(
+                GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.parse(groupProtocol))
+            )
+        );
 
-        final AtomicReference<scala.collection.Set> foundListing = new AtomicReference<>();
-        TestUtils.waitForCondition(() -> {
-            foundListing.set(service.listConsumerGroupsWithState(set(Arrays.asList(ConsumerGroupState.values()))).toSet());
-            return Objects.equals(expectedListing, foundListing.get());
-        }, "Expected to show groups " + expectedListing + ", but found " + foundListing.get());
+        assertGroupListing(
+            service,
+            Collections.emptySet(),
+            EnumSet.allOf(ConsumerGroupState.class),
+            expectedListing
+        );
 
-        scala.collection.Set<ConsumerGroupListing> expectedListingStable = set(Collections.singleton(
-            new ConsumerGroupListing(GROUP, false, Optional.of(ConsumerGroupState.STABLE))));
+        expectedListing = mkSet(
+            new ConsumerGroupListing(
+                GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.parse(groupProtocol))
+            )
+        );
 
-        foundListing.set(null);
+        assertGroupListing(
+            service,
+            Collections.emptySet(),
+            mkSet(ConsumerGroupState.STABLE),
+            expectedListing
+        );
 
-        TestUtils.waitForCondition(() -> {
-            foundListing.set(service.listConsumerGroupsWithState(set(Collections.singleton(ConsumerGroupState.STABLE))).toSet());
-            return Objects.equals(expectedListingStable, foundListing.get());
-        }, "Expected to show groups " + expectedListingStable + ", but found " + foundListing.get());
+        assertGroupListing(
+            service,
+            Collections.emptySet(),
+            mkSet(ConsumerGroupState.PREPARING_REBALANCE),
+            Collections.emptySet()
+        );
     }
 
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
-    public void testConsumerGroupStatesFromString(String quorum) {
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly")
+    public void testListConsumerGroupsWithTypesClassicProtocol(String quorum, String groupProtocol) throws Exception {
+        String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        Set<ConsumerGroupListing> expectedListing = mkSet(
+            new ConsumerGroupListing(
+                simpleGroup,
+                true,
+                Optional.of(ConsumerGroupState.EMPTY),
+                Optional.of(GroupType.CLASSIC)
+            ),
+            new ConsumerGroupListing(
+                GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.CLASSIC)
+            )
+        );
+
+        // No filters explicitly mentioned. Expectation is that all groups are returned.
+        assertGroupListing(
+            service,
+            Collections.emptySet(),
+            Collections.emptySet(),
+            expectedListing
+        );
+
+        // When group type is mentioned:
+        // Old Group Coordinator returns empty listings if the type is not Classic.
+        // New Group Coordinator returns groups according to the filter.
+        assertGroupListing(
+            service,
+            mkSet(GroupType.CONSUMER),
+            Collections.emptySet(),
+            Collections.emptySet()
+        );
+
+        assertGroupListing(
+            service,
+            mkSet(GroupType.CLASSIC),
+            Collections.emptySet(),
+            expectedListing
+        );
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly")
+    public void testListConsumerGroupsWithTypesConsumerProtocol(String quorum, String groupProtocol) throws Exception {
+        String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1);
+        addConsumerGroupExecutor(1, PROTOCOL_GROUP, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        // No filters explicitly mentioned. Expectation is that all groups are returned.
+        Set<ConsumerGroupListing> expectedListing = mkSet(
+            new ConsumerGroupListing(
+                simpleGroup,
+                true,
+                Optional.of(ConsumerGroupState.EMPTY),
+                Optional.of(GroupType.CLASSIC)
+            ),
+            new ConsumerGroupListing(
+                GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.CLASSIC)
+            ),
+            new ConsumerGroupListing(
+                PROTOCOL_GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.CONSUMER)
+            )
+        );
+
+        assertGroupListing(
+            service,
+            Collections.emptySet(),
+            Collections.emptySet(),
+            expectedListing
+        );
+
+        // When group type is mentioned:
+        // New Group Coordinator returns groups according to the filter.
+        expectedListing = mkSet(
+            new ConsumerGroupListing(
+                PROTOCOL_GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.CONSUMER)
+            )
+        );
+
+        assertGroupListing(
+            service,
+            mkSet(GroupType.CONSUMER),
+            Collections.emptySet(),
+            expectedListing
+        );
+
+        expectedListing = mkSet(
+            new ConsumerGroupListing(
+                simpleGroup,
+                true,
+                Optional.of(ConsumerGroupState.EMPTY),
+                Optional.of(GroupType.CLASSIC)
+            ),
+            new ConsumerGroupListing(
+                GROUP,
+                false,
+                Optional.of(ConsumerGroupState.STABLE),
+                Optional.of(GroupType.CLASSIC)
+            )
+        );
+
+        assertGroupListing(
+            service,
+            mkSet(GroupType.CLASSIC),
+            Collections.emptySet(),
+            expectedListing
+        );
+    }
+
+    @Test
+    public void testConsumerGroupStatesFromString() {
         scala.collection.Set<ConsumerGroupState> result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable");
         assertEquals(set(Collections.singleton(ConsumerGroupState.STABLE)), result);
 
@@ -107,7 +282,7 @@ public void testConsumerGroupStatesFromString(String quorum) {
         assertEquals(set(Arrays.asList(ConsumerGroupState.DEAD, ConsumerGroupState.COMPLETING_REBALANCE)), result);
 
         result = ConsumerGroupCommand.consumerGroupStatesFromString("stable");
-        assertEquals(set(Arrays.asList(ConsumerGroupState.STABLE)), result);
+        assertEquals(set(Collections.singletonList(ConsumerGroupState.STABLE)), result);
 
         result = ConsumerGroupCommand.consumerGroupStatesFromString("stable, assigning");
         assertEquals(set(Arrays.asList(ConsumerGroupState.STABLE, ConsumerGroupState.ASSIGNING)), result);
@@ -122,10 +297,31 @@ public void testConsumerGroupStatesFromString(String quorum) {
         assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupStatesFromString("   ,   ,"));
     }
 
-    @ParameterizedTest
-    @ValueSource(strings = {"zk", "kraft"})
-    public void testListGroupCommand(String quorum) throws Exception {
+    @Test
+    public void testConsumerGroupTypesFromString() {
+        scala.collection.Set<GroupType> result = ConsumerGroupCommand.consumerGroupTypesFromString("consumer");
+        assertEquals(set(Collections.singleton(GroupType.CONSUMER)), result);
+
+        result = ConsumerGroupCommand.consumerGroupTypesFromString("consumer, classic");
+        assertEquals(set(Arrays.asList(GroupType.CONSUMER, GroupType.CLASSIC)), result);
+
+        result = ConsumerGroupCommand.consumerGroupTypesFromString("Consumer, Classic");
+        assertEquals(set(Arrays.asList(GroupType.CONSUMER, GroupType.CLASSIC)), result);
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupTypesFromString("bad, wrong"));
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupTypesFromString("  bad, generic"));
+
+        assertThrows(IllegalArgumentException.class, () -> ConsumerGroupCommand.consumerGroupTypesFromString("   ,   ,"));
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly")
+    public void testListGroupCommandClassicProtocol(String quorum, String groupProtocol) throws Exception {
         String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
         addSimpleGroupExecutor(simpleGroup);
         addConsumerGroupExecutor(1);
 
@@ -147,6 +343,24 @@ public void testListGroupCommand(String quorum) throws Exception {
             )
         );
 
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type"),
+            Arrays.asList("GROUP", "TYPE"),
+            mkSet(
+                Arrays.asList(GROUP, "Classic"),
+                Arrays.asList(simpleGroup, "Classic")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "--state"),
+            Arrays.asList("GROUP", "TYPE", "STATE"),
+            mkSet(
+                Arrays.asList(GROUP, "Classic", "Stable"),
+                Arrays.asList(simpleGroup, "Classic", "Empty")
+            )
+        );
+
         validateListOutput(
             Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "Stable"),
             Arrays.asList("GROUP", "STATE"),
@@ -155,6 +369,7 @@ public void testListGroupCommand(String quorum) throws Exception {
             )
         );
 
+        // Check case-insensitivity in state filter.
         validateListOutput(
             Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state", "stable"),
             Arrays.asList("GROUP", "STATE"),
@@ -162,6 +377,109 @@ public void testListGroupCommand(String quorum) throws Exception {
                 Arrays.asList(GROUP, "Stable")
             )
         );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "Classic"),
+            Arrays.asList("GROUP", "TYPE"),
+            mkSet(
+                Arrays.asList(GROUP, "Classic"),
+                Arrays.asList(simpleGroup, "Classic")
+            )
+        );
+
+        // Check case-insensitivity in type filter.
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "classic"),
+            Arrays.asList("GROUP", "TYPE"),
+            mkSet(
+                Arrays.asList(GROUP, "Classic"),
+                Arrays.asList(simpleGroup, "Classic")
+            )
+        );
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource("getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly")
+    public void testListGroupCommandConsumerProtocol(String quorum, String groupProtocol) throws Exception {
+        String simpleGroup = "simple-group";
+
+        createOffsetsTopic(listenerName(), new Properties());
+
+        addSimpleGroupExecutor(simpleGroup);
+        addConsumerGroupExecutor(1, PROTOCOL_GROUP, groupProtocol);
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list"),
+            Collections.emptyList(),
+            mkSet(
+                Collections.singletonList(PROTOCOL_GROUP),
+                Collections.singletonList(simpleGroup)
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--state"),
+            Arrays.asList("GROUP", "STATE"),
+            mkSet(
+                Arrays.asList(PROTOCOL_GROUP, "Stable"),
+                Arrays.asList(simpleGroup, "Empty")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type"),
+            Arrays.asList("GROUP", "TYPE"),
+            mkSet(
+                Arrays.asList(PROTOCOL_GROUP, "Consumer"),
+                Arrays.asList(simpleGroup, "Classic")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "--state"),
+            Arrays.asList("GROUP", "TYPE", "STATE"),
+            mkSet(
+                Arrays.asList(PROTOCOL_GROUP, "Consumer", "Stable"),
+                Arrays.asList(simpleGroup, "Classic", "Empty")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "consumer"),
+            Arrays.asList("GROUP", "TYPE"),
+            mkSet(
+                Arrays.asList(PROTOCOL_GROUP, "Consumer")
+            )
+        );
+
+        validateListOutput(
+            Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--list", "--type", "consumer", "--state", "Stable"),
+            Arrays.asList("GROUP", "TYPE", "STATE"),
+            mkSet(
+                Arrays.asList(PROTOCOL_GROUP, "Consumer", "Stable")
+            )
+        );
+    }
+
+    /**
+     * Validates the consumer group listings returned against expected values using specified filters.
+     *
+     * @param service           The service to list consumer groups.
+     * @param typeFilterSet     Filters for group types, empty for no filter.
+     * @param stateFilterSet    Filters for group states, empty for no filter.
+     * @param expectedListing   Expected consumer group listings.
+     */
+    private static void assertGroupListing(
+        ConsumerGroupCommand.ConsumerGroupService service,
+        Set<GroupType> typeFilterSet,
+        Set<ConsumerGroupState> stateFilterSet,
+        Set<ConsumerGroupListing> expectedListing
+    ) throws Exception {
+        final AtomicReference<scala.collection.Set> foundListing = new AtomicReference<>();
+        TestUtils.waitForCondition(() -> {
+            foundListing.set(service.listConsumerGroupsWithFilters(set(typeFilterSet), set(stateFilterSet)).toSet());
+            return Objects.equals(set(expectedListing), foundListing.get());
+        }, () -> "Expected to show groups " + expectedListing + ", but found " + foundListing.get() + ".");
     }
 
     /**

From f8eb4294d67b37e854aa14fb989d5e074df82ac2 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Thu, 29 Feb 2024 02:22:23 -0800
Subject: [PATCH 092/258] KAFKA-16191: Clean up of consumer client internal
 events (#15438)

There are a few minor issues with the event sub-classes in the
org.apache.kafka.clients.consumer.internals.events package that should be cleaned up:

- Update the names of subclasses to remove "Application" or "Background"
- Make toString() final in the base classes and clean up the implementations of toStringBase()
- Fix minor whitespace inconsistencies
- Make variable/method names consistent

Reviewer: Bruno Cadonna <cadonna@apache.org>
---
 .../internals/AsyncKafkaConsumer.java         |  89 ++++++------
 .../internals/CoordinatorRequestManager.java  |   6 +-
 .../internals/HeartbeatRequestManager.java    |   4 +-
 .../internals/MembershipManagerImpl.java      |   4 +-
 .../internals/OffsetsRequestManager.java      |   4 +-
 .../events/AbstractTopicMetadataEvent.java    |  41 ++++++
 .../events/AllTopicsMetadataEvent.java        |  24 ++++
 .../internals/events/ApplicationEvent.java    |  35 +++--
 .../events/ApplicationEventHandler.java       |   6 +-
 .../events/ApplicationEventProcessor.java     | 123 +++++++++-------
 ...nEvent.java => AssignmentChangeEvent.java} |  33 +----
 ...cationEvent.java => AsyncCommitEvent.java} |  15 +-
 .../internals/events/BackgroundEvent.java     |  30 ++--
 .../events/BackgroundEventHandler.java        |   2 +-
 ...ApplicationEvent.java => CommitEvent.java} |  21 +--
 ...tionEvent.java => CommitOnCloseEvent.java} |  11 +-
 .../events/CompletableApplicationEvent.java   |  45 +-----
 .../events/CompletableBackgroundEvent.java    |  28 +---
 .../internals/events/CompletableEvent.java    |   1 -
 ...balanceListenerCallbackCompletedEvent.java |  31 +---
 ...rRebalanceListenerCallbackNeededEvent.java |  30 +---
 .../events/ErrorBackgroundEvent.java          |  59 --------
 .../consumer/internals/events/ErrorEvent.java |  38 +++++
 ...t.java => FetchCommittedOffsetsEvent.java} |  31 +---
 .../events/GroupMetadataUpdateEvent.java      |  32 +----
 ...ationEvent.java => LeaveOnCloseEvent.java} |  12 +-
 ...cationEvent.java => ListOffsetsEvent.java} |  33 +----
 .../NewTopicsMetadataUpdateRequestEvent.java  |   7 -
 .../events/PollApplicationEvent.java          |  57 --------
 .../consumer/internals/events/PollEvent.java  |  36 +++++
 ...ionEvent.java => ResetPositionsEvent.java} |   4 +-
 ...vent.java => SubscriptionChangeEvent.java} |   4 +-
 ...icationEvent.java => SyncCommitEvent.java} |  17 +--
 .../events/TopicMetadataApplicationEvent.java |  78 ----------
 .../internals/events/TopicMetadataEvent.java  |  38 +++++
 ...cationEvent.java => UnsubscribeEvent.java} |   5 +-
 ...Event.java => ValidatePositionsEvent.java} |   4 +-
 .../internals/AsyncKafkaConsumerTest.java     | 134 +++++++++---------
 .../internals/ConsumerNetworkThreadTest.java  |  54 +++----
 .../CoordinatorRequestManagerTest.java        |   6 +-
 .../HeartbeatRequestManagerTest.java          |  22 +--
 .../internals/MembershipManagerImplTest.java  |   8 +-
 .../internals/OffsetsRequestManagerTest.java  |   6 +-
 .../events/ApplicationEventProcessorTest.java |   4 +-
 44 files changed, 519 insertions(+), 753 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{AssignmentChangeApplicationEvent.java => AssignmentChangeEvent.java} (56%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{AsyncCommitApplicationEvent.java => AsyncCommitEvent.java} (73%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{CommitApplicationEvent.java => CommitEvent.java} (71%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{CommitOnCloseApplicationEvent.java => CommitOnCloseEvent.java} (76%)
 delete mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorBackgroundEvent.java
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorEvent.java
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{FetchCommittedOffsetsApplicationEvent.java => FetchCommittedOffsetsEvent.java} (60%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{LeaveOnCloseApplicationEvent.java => LeaveOnCloseEvent.java} (76%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{ListOffsetsApplicationEvent.java => ListOffsetsEvent.java} (69%)
 delete mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollApplicationEvent.java
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollEvent.java
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{ResetPositionsApplicationEvent.java => ResetPositionsEvent.java} (89%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{SubscriptionChangeApplicationEvent.java => SubscriptionChangeEvent.java} (90%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{SyncCommitApplicationEvent.java => SyncCommitEvent.java} (73%)
 delete mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataApplicationEvent.java
 create mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{UnsubscribeApplicationEvent.java => UnsubscribeEvent.java} (91%)
 rename clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/{ValidatePositionsApplicationEvent.java => ValidatePositionsEvent.java} (89%)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index d4b461b0140d8..d810c5f053bb2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -36,32 +36,33 @@
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.internals.events.AllTopicsMetadataEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
-import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.CommitEvent;
+import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
-import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
-import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ListOffsetsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseEvent;
+import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
-import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.TopicMetadataApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.PollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ResetPositionsEvent;
+import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitEvent;
+import org.apache.kafka.clients.consumer.internals.events.TopicMetadataEvent;
+import org.apache.kafka.clients.consumer.internals.events.UnsubscribeEvent;
+import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.KafkaConsumerMetrics;
 import org.apache.kafka.clients.consumer.internals.metrics.RebalanceCallbackMetricsManager;
 import org.apache.kafka.common.Cluster;
@@ -179,7 +180,7 @@ public BackgroundEventProcessor(final LogContext logContext,
 
         /**
          * Process the events—if any—that were produced by the {@link ConsumerNetworkThread network thread}.
-         * It is possible that {@link org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent an error}
+         * It is possible that {@link ErrorEvent an error}
          * could occur when processing the events. In such cases, the processor will take a reference to the first
          * error, continue to process the remaining events, and then throw the first error that occurred.
          */
@@ -209,7 +210,7 @@ public boolean process() {
         public void process(final BackgroundEvent event) {
             switch (event.type()) {
                 case ERROR:
-                    process((ErrorBackgroundEvent) event);
+                    process((ErrorEvent) event);
                     break;
 
                 case GROUP_METADATA_UPDATE:
@@ -226,7 +227,7 @@ public void process(final BackgroundEvent event) {
             }
         }
 
-        private void process(final ErrorBackgroundEvent event) {
+        private void process(final ErrorEvent event) {
             throw event.error();
         }
 
@@ -703,7 +704,7 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
             do {
 
                 // Make sure to let the background thread know that we are still polling.
-                applicationEventHandler.add(new PollApplicationEvent(timer.currentTimeMs()));
+                applicationEventHandler.add(new PollEvent(timer.currentTimeMs()));
 
                 // We must not allow wake-ups between polling for fetches and returning the records.
                 // If the polled fetches are not empty the consumed position has already been updated in the polling
@@ -768,7 +769,7 @@ public void commitAsync(OffsetCommitCallback callback) {
     public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCommitCallback callback) {
         acquireAndEnsureOpen();
         try {
-            AsyncCommitApplicationEvent asyncCommitEvent = new AsyncCommitApplicationEvent(offsets);
+            AsyncCommitEvent asyncCommitEvent = new AsyncCommitEvent(offsets);
             CompletableFuture<Void> future = commit(asyncCommitEvent);
             future.whenComplete((r, t) -> {
 
@@ -790,7 +791,7 @@ public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCo
         }
     }
 
-    private CompletableFuture<Void> commit(final CommitApplicationEvent commitEvent) {
+    private CompletableFuture<Void> commit(final CommitEvent commitEvent) {
         maybeInvokeCommitCallbacks();
         maybeThrowFencedInstanceException();
         maybeThrowInvalidGroupIdException();
@@ -936,7 +937,7 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
                 return Collections.emptyMap();
             }
 
-            final FetchCommittedOffsetsApplicationEvent event = new FetchCommittedOffsetsApplicationEvent(
+            final FetchCommittedOffsetsEvent event = new FetchCommittedOffsetsEvent(
                 partitions,
                 timeout.toMillis());
             wakeupTrigger.setActiveTask(event.future());
@@ -988,12 +989,11 @@ public List<PartitionInfo> partitionsFor(String topic, Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final TopicMetadataApplicationEvent topicMetadataApplicationEvent =
-                    new TopicMetadataApplicationEvent(topic, timeout.toMillis());
-            wakeupTrigger.setActiveTask(topicMetadataApplicationEvent.future());
+            final TopicMetadataEvent topicMetadataEvent = new TopicMetadataEvent(topic, timeout.toMillis());
+            wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
                 Map<String, List<PartitionInfo>> topicMetadata =
-                        applicationEventHandler.addAndGet(topicMetadataApplicationEvent, time.timer(timeout));
+                        applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
 
                 return topicMetadata.getOrDefault(topic, Collections.emptyList());
             } finally {
@@ -1017,11 +1017,10 @@ public Map<String, List<PartitionInfo>> listTopics(Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final TopicMetadataApplicationEvent topicMetadataApplicationEvent =
-                    new TopicMetadataApplicationEvent(timeout.toMillis());
-            wakeupTrigger.setActiveTask(topicMetadataApplicationEvent.future());
+            final AllTopicsMetadataEvent topicMetadataEvent = new AllTopicsMetadataEvent(timeout.toMillis());
+            wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
-                return applicationEventHandler.addAndGet(topicMetadataApplicationEvent, time.timer(timeout));
+                return applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
             } finally {
                 wakeupTrigger.clearTask();
             }
@@ -1089,7 +1088,7 @@ public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(Map<TopicPartitio
             if (timestampsToSearch.isEmpty()) {
                 return Collections.emptyMap();
             }
-            final ListOffsetsApplicationEvent listOffsetsEvent = new ListOffsetsApplicationEvent(
+            final ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampsToSearch,
                 true);
 
@@ -1139,7 +1138,7 @@ private Map<TopicPartition, Long> beginningOrEndOffset(Collection<TopicPartition
             Map<TopicPartition, Long> timestampToSearch = partitions
                 .stream()
                 .collect(Collectors.toMap(Function.identity(), tp -> timestamp));
-            ListOffsetsApplicationEvent listOffsetsEvent = new ListOffsetsApplicationEvent(
+            ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampToSearch,
                 false);
             Map<TopicPartition, OffsetAndTimestamp> offsetAndTimestampMap = applicationEventHandler.addAndGet(
@@ -1267,11 +1266,11 @@ void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstEx
         if (!groupMetadata.isPresent())
             return;
         maybeAutoCommitSync(autoCommitEnabled, timer, firstException);
-        applicationEventHandler.add(new CommitOnCloseApplicationEvent());
+        applicationEventHandler.add(new CommitOnCloseEvent());
         completeQuietly(
             () -> {
                 maybeRevokePartitions();
-                applicationEventHandler.addAndGet(new LeaveOnCloseApplicationEvent(), timer);
+                applicationEventHandler.addAndGet(new LeaveOnCloseEvent(), timer);
             },
             "Failed to send leaveGroup heartbeat with a timeout(ms)=" + timer.timeoutMs(), firstException);
     }
@@ -1349,7 +1348,7 @@ public void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration
         long commitStart = time.nanoseconds();
         try {
             Timer requestTimer = time.timer(timeout.toMillis());
-            SyncCommitApplicationEvent syncCommitEvent = new SyncCommitApplicationEvent(offsets, timeout.toMillis());
+            SyncCommitEvent syncCommitEvent = new SyncCommitEvent(offsets, timeout.toMillis());
             CompletableFuture<Void> commitFuture = commit(syncCommitEvent);
             wakeupTrigger.setActiveTask(commitFuture);
             ConsumerUtils.getResult(commitFuture, requestTimer);
@@ -1429,7 +1428,7 @@ public void assign(Collection<TopicPartition> partitions) {
             // be no following rebalance.
             //
             // See the ApplicationEventProcessor.process() method that handles this event for more detail.
-            applicationEventHandler.add(new AssignmentChangeApplicationEvent(subscriptions.allConsumed(), time.milliseconds()));
+            applicationEventHandler.add(new AssignmentChangeEvent(subscriptions.allConsumed(), time.milliseconds()));
 
             log.info("Assigned to partition(s): {}", join(partitions, ", "));
             if (subscriptions.assignFromUser(new HashSet<>(partitions)))
@@ -1463,13 +1462,13 @@ public void unsubscribe() {
         try {
             fetchBuffer.retainAll(Collections.emptySet());
             if (groupMetadata.isPresent()) {
-                UnsubscribeApplicationEvent unsubscribeApplicationEvent = new UnsubscribeApplicationEvent();
-                applicationEventHandler.add(unsubscribeApplicationEvent);
+                UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent();
+                applicationEventHandler.add(unsubscribeEvent);
                 log.info("Unsubscribing all topics or patterns and assigned partitions");
                 Timer timer = time.timer(Long.MAX_VALUE);
 
                 try {
-                    processBackgroundEvents(backgroundEventProcessor, unsubscribeApplicationEvent.future(), timer);
+                    processBackgroundEvents(backgroundEventProcessor, unsubscribeEvent.future(), timer);
                     log.info("Unsubscribed all topics or patterns and assigned partitions");
                 } catch (TimeoutException e) {
                     log.error("Failed while waiting for the unsubscribe event to complete");
@@ -1567,7 +1566,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // Validate positions using the partition leader end offsets, to detect if any partition
             // has been truncated due to a leader change. This will trigger an OffsetForLeaderEpoch
             // request, retrieve the partition end offsets, and validate the current position against it.
-            applicationEventHandler.addAndGet(new ValidatePositionsApplicationEvent(), timer);
+            applicationEventHandler.addAndGet(new ValidatePositionsEvent(), timer);
 
             cachedSubscriptionHasAllFetchPositions = subscriptions.hasAllFetchPositions();
             if (cachedSubscriptionHasAllFetchPositions) return true;
@@ -1590,7 +1589,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // which are awaiting reset. This will trigger a ListOffset request, retrieve the
             // partition offsets according to the strategy (ex. earliest, latest), and update the
             // positions.
-            applicationEventHandler.addAndGet(new ResetPositionsApplicationEvent(), timer);
+            applicationEventHandler.addAndGet(new ResetPositionsEvent(), timer);
             return true;
         } catch (TimeoutException e) {
             return false;
@@ -1620,8 +1619,8 @@ private boolean initWithCommittedOffsetsIfNeeded(Timer timer) {
 
         log.debug("Refreshing committed offsets for partitions {}", initializingPartitions);
         try {
-            final FetchCommittedOffsetsApplicationEvent event =
-                new FetchCommittedOffsetsApplicationEvent(
+            final FetchCommittedOffsetsEvent event =
+                new FetchCommittedOffsetsEvent(
                     initializingPartitions,
                     timer.remainingMs());
             final Map<TopicPartition, OffsetAndMetadata> offsets = applicationEventHandler.addAndGet(event, timer);
@@ -1770,7 +1769,7 @@ private void subscribeInternal(Collection<String> topics, Optional<ConsumerRebal
 
                 // Trigger subscribe event to effectively join the group if not already part of it,
                 // or just send the new subscription to the broker.
-                applicationEventHandler.add(new SubscriptionChangeApplicationEvent());
+                applicationEventHandler.add(new SubscriptionChangeEvent());
             }
         } finally {
             release();
@@ -1797,7 +1796,7 @@ private void subscribeInternal(Collection<String> topics, Optional<ConsumerRebal
      * <p/>
      *
      * As an example, take {@link #unsubscribe()}. To start unsubscribing, the application thread enqueues an
-     * {@link UnsubscribeApplicationEvent} on the application event queue. That event will eventually trigger the
+     * {@link UnsubscribeEvent} on the application event queue. That event will eventually trigger the
      * rebalancing logic in the background thread. Critically, as part of this rebalancing work, the
      * {@link ConsumerRebalanceListener#onPartitionsRevoked(Collection)} callback needs to be invoked. However,
      * this callback must be executed on the application thread. To achieve this, the background thread enqueues a
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManager.java
index d6a72812a52f9..a6cc28fb0f4c6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManager.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
@@ -175,12 +175,12 @@ private void onFailedResponse(final long currentTimeMs, final Throwable exceptio
         if (exception == Errors.GROUP_AUTHORIZATION_FAILED.exception()) {
             log.debug("FindCoordinator request failed due to authorization error {}", exception.getMessage());
             KafkaException groupAuthorizationException = GroupAuthorizationException.forGroupId(this.groupId);
-            backgroundEventHandler.add(new ErrorBackgroundEvent(groupAuthorizationException));
+            backgroundEventHandler.add(new ErrorEvent(groupAuthorizationException));
             return;
         }
 
         log.warn("FindCoordinator request failed due to fatal exception", exception);
-        backgroundEventHandler.add(new ErrorBackgroundEvent(exception));
+        backgroundEventHandler.add(new ErrorEvent(exception));
     }
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index d551dbe2508c0..826774a6a64e0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -21,7 +21,7 @@
 import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.HeartbeatMetricsManager;
 import org.apache.kafka.common.Uuid;
@@ -446,7 +446,7 @@ private void logInfo(final String message,
     }
 
     private void handleFatalFailure(Throwable error) {
-        backgroundEventHandler.add(new ErrorBackgroundEvent(error));
+        backgroundEventHandler.add(new ErrorEvent(error));
         membershipManager.transitionToFatal();
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 6f3947eea46e5..e74c7e30a2cad 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -26,7 +26,7 @@
 import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicIdPartition;
@@ -255,7 +255,7 @@ public class MembershipManagerImpl implements MembershipManager {
     /**
      * Serves as the conduit by which we can report events to the application thread. This is needed as we send
      * {@link ConsumerRebalanceListenerCallbackNeededEvent callbacks} and, if needed,
-     * {@link ErrorBackgroundEvent errors} to the application thread.
+     * {@link ErrorEvent errors} to the application thread.
      */
     private final BackgroundEventHandler backgroundEventHandler;
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManager.java
index 34f4b30c44dd3..c5156e9e0b939 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManager.java
@@ -25,7 +25,7 @@
 import org.apache.kafka.clients.consumer.internals.OffsetFetcherUtils.ListOffsetData;
 import org.apache.kafka.clients.consumer.internals.OffsetFetcherUtils.ListOffsetResult;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.common.ClusterResource;
 import org.apache.kafka.common.ClusterResourceListener;
 import org.apache.kafka.common.IsolationLevel;
@@ -199,7 +199,7 @@ public CompletableFuture<Void> resetPositionsIfNeeded() {
         try {
             offsetResetTimestamps = offsetFetcherUtils.getOffsetResetTimestamp();
         } catch (Exception e) {
-            backgroundEventHandler.add(new ErrorBackgroundEvent(e));
+            backgroundEventHandler.add(new ErrorEvent(e));
             return CompletableFuture.completedFuture(null);
         }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
new file mode 100644
index 0000000000000..31c21817d85a3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+import org.apache.kafka.common.PartitionInfo;
+
+import java.util.List;
+import java.util.Map;
+
+public abstract class AbstractTopicMetadataEvent extends CompletableApplicationEvent<Map<String, List<PartitionInfo>>> {
+
+    private final long timeoutMs;
+
+    protected AbstractTopicMetadataEvent(final Type type, final long timeoutMs) {
+        super(type);
+        this.timeoutMs = timeoutMs;
+    }
+
+    public long timeoutMs() {
+        return timeoutMs;
+    }
+
+    @Override
+    public String toStringBase() {
+        return super.toStringBase() + ", timeoutMs=" + timeoutMs;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
new file mode 100644
index 0000000000000..154703aaee15b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+public class AllTopicsMetadataEvent extends AbstractTopicMetadataEvent {
+
+    public AllTopicsMetadataEvent(final long timeoutMs) {
+        super(Type.ALL_TOPICS_METADATA, timeoutMs);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
index ac7ccc56c55f2..2897117da8bab 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
@@ -16,52 +16,65 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer;
+import org.apache.kafka.common.Uuid;
+
 import java.util.Objects;
 
 /**
- * This is the abstract definition of the events created by the KafkaConsumer API
+ * This is the abstract definition of the events created by the {@link AsyncKafkaConsumer} on the user's
+ * application thread.
  */
 public abstract class ApplicationEvent {
 
     public enum Type {
         COMMIT_ASYNC, COMMIT_SYNC, POLL, FETCH_COMMITTED_OFFSETS, NEW_TOPICS_METADATA_UPDATE, ASSIGNMENT_CHANGE,
-        LIST_OFFSETS, RESET_POSITIONS, VALIDATE_POSITIONS, TOPIC_METADATA, SUBSCRIPTION_CHANGE,
+        LIST_OFFSETS, RESET_POSITIONS, VALIDATE_POSITIONS, TOPIC_METADATA, ALL_TOPICS_METADATA, SUBSCRIPTION_CHANGE,
         UNSUBSCRIBE, CONSUMER_REBALANCE_LISTENER_CALLBACK_COMPLETED,
         COMMIT_ON_CLOSE, LEAVE_ON_CLOSE
     }
 
     private final Type type;
 
+    /**
+     * This identifies a particular event. It is used to disambiguate events via {@link #hashCode()} and
+     * {@link #equals(Object)} and can be used in log messages when debugging.
+     */
+    private final Uuid id;
+
     protected ApplicationEvent(Type type) {
         this.type = Objects.requireNonNull(type);
+        this.id = Uuid.randomUuid();
     }
 
     public Type type() {
         return type;
     }
 
+    public Uuid id() {
+        return id;
+    }
+
     @Override
-    public boolean equals(Object o) {
+    public final boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         ApplicationEvent that = (ApplicationEvent) o;
-
-        return type == that.type;
+        return type == that.type && id.equals(that.id);
     }
 
     @Override
-    public int hashCode() {
-        return type.hashCode();
+    public final int hashCode() {
+        return Objects.hash(type, id);
     }
 
     protected String toStringBase() {
-        return "type=" + type;
+        return "type=" + type + ", id=" + id;
     }
 
     @Override
-    public String toString() {
-        return "ApplicationEvent{" +
+    public final String toString() {
+        return getClass().getSimpleName() + "{" +
                 toStringBase() +
                 '}';
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventHandler.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventHandler.java
index 7535edf5970b3..eac1cc3d62f92 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventHandler.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer.internals.events;
 
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkThread;
+import org.apache.kafka.clients.consumer.internals.ConsumerUtils;
 import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate;
 import org.apache.kafka.clients.consumer.internals.RequestManagers;
 import org.apache.kafka.common.internals.IdempotentCloser;
@@ -99,10 +100,9 @@ public long maximumTimeToWait() {
      *
      * <p/>
      *
-     * See {@link CompletableApplicationEvent#get(Timer)} and {@link Future#get(long, TimeUnit)} for more details.
+     * See {@link ConsumerUtils#getResult(Future, Timer)} and {@link Future#get(long, TimeUnit)} for more details.
      *
      * @param event A {@link CompletableApplicationEvent} created by the polling thread
-     * @param timer Timer for which to wait for the event to complete
      * @return      Value that is the result of the event
      * @param <T>   Type of return value of the event
      */
@@ -110,7 +110,7 @@ public <T> T addAndGet(final CompletableApplicationEvent<T> event, final Timer t
         Objects.requireNonNull(event, "CompletableApplicationEvent provided to addAndGet must be non-null");
         Objects.requireNonNull(timer, "Timer provided to addAndGet must be non-null");
         add(event);
-        return event.get(timer);
+        return ConsumerUtils.getResult(event.future(), timer);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
index 9e48b4de6daad..c86aa8815f250 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.internals.CachedSupplier;
 import org.apache.kafka.clients.consumer.internals.CommitRequestManager;
@@ -34,6 +35,7 @@
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 
 /**
@@ -65,23 +67,24 @@ public boolean process() {
         return process((event, error) -> error.ifPresent(e -> log.warn("Error processing event {}", e.getMessage(), e)));
     }
 
+    @SuppressWarnings({"CyclomaticComplexity"})
     @Override
     public void process(ApplicationEvent event) {
         switch (event.type()) {
             case COMMIT_ASYNC:
-                process((AsyncCommitApplicationEvent) event);
+                process((AsyncCommitEvent) event);
                 return;
 
             case COMMIT_SYNC:
-                process((SyncCommitApplicationEvent) event);
+                process((SyncCommitEvent) event);
                 return;
 
             case POLL:
-                process((PollApplicationEvent) event);
+                process((PollEvent) event);
                 return;
 
             case FETCH_COMMITTED_OFFSETS:
-                process((FetchCommittedOffsetsApplicationEvent) event);
+                process((FetchCommittedOffsetsEvent) event);
                 return;
 
             case NEW_TOPICS_METADATA_UPDATE:
@@ -89,31 +92,35 @@ public void process(ApplicationEvent event) {
                 return;
 
             case ASSIGNMENT_CHANGE:
-                process((AssignmentChangeApplicationEvent) event);
+                process((AssignmentChangeEvent) event);
                 return;
 
             case TOPIC_METADATA:
-                process((TopicMetadataApplicationEvent) event);
+                process((TopicMetadataEvent) event);
+                return;
+
+            case ALL_TOPICS_METADATA:
+                process((AllTopicsMetadataEvent) event);
                 return;
 
             case LIST_OFFSETS:
-                process((ListOffsetsApplicationEvent) event);
+                process((ListOffsetsEvent) event);
                 return;
 
             case RESET_POSITIONS:
-                process((ResetPositionsApplicationEvent) event);
+                process((ResetPositionsEvent) event);
                 return;
 
             case VALIDATE_POSITIONS:
-                process((ValidatePositionsApplicationEvent) event);
+                process((ValidatePositionsEvent) event);
                 return;
 
             case SUBSCRIPTION_CHANGE:
-                process((SubscriptionChangeApplicationEvent) event);
+                process((SubscriptionChangeEvent) event);
                 return;
 
             case UNSUBSCRIBE:
-                process((UnsubscribeApplicationEvent) event);
+                process((UnsubscribeEvent) event);
                 return;
 
             case CONSUMER_REBALANCE_LISTENER_CALLBACK_COMPLETED:
@@ -121,11 +128,11 @@ public void process(ApplicationEvent event) {
                 return;
 
             case COMMIT_ON_CLOSE:
-                process((CommitOnCloseApplicationEvent) event);
+                process((CommitOnCloseEvent) event);
                 return;
 
             case LEAVE_ON_CLOSE:
-                process((LeaveOnCloseApplicationEvent) event);
+                process((LeaveOnCloseEvent) event);
                 return;
 
             default:
@@ -133,7 +140,7 @@ public void process(ApplicationEvent event) {
         }
     }
 
-    private void process(final PollApplicationEvent event) {
+    private void process(final PollEvent event) {
         if (!requestManagers.commitRequestManager.isPresent()) {
             return;
         }
@@ -142,20 +149,28 @@ private void process(final PollApplicationEvent event) {
         requestManagers.heartbeatRequestManager.ifPresent(hrm -> hrm.resetPollTimer(event.pollTimeMs()));
     }
 
-    private void process(final AsyncCommitApplicationEvent event) {
+    private void process(final AsyncCommitEvent event) {
+        if (!requestManagers.commitRequestManager.isPresent()) {
+            return;
+        }
+
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        CompletableFuture<Void> commitResult = manager.commitAsync(event.offsets());
-        event.chain(commitResult);
+        CompletableFuture<Void> future = manager.commitAsync(event.offsets());
+        future.whenComplete(complete(event.future()));
     }
 
-    private void process(final SyncCommitApplicationEvent event) {
+    private void process(final SyncCommitEvent event) {
+        if (!requestManagers.commitRequestManager.isPresent()) {
+            return;
+        }
+
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
         long expirationTimeoutMs = getExpirationTimeForTimeout(event.retryTimeoutMs());
-        CompletableFuture<Void> commitResult = manager.commitSync(event.offsets(), expirationTimeoutMs);
-        event.chain(commitResult);
+        CompletableFuture<Void> future = manager.commitSync(event.offsets(), expirationTimeoutMs);
+        future.whenComplete(complete(event.future()));
     }
 
-    private void process(final FetchCommittedOffsetsApplicationEvent event) {
+    private void process(final FetchCommittedOffsetsEvent event) {
         if (!requestManagers.commitRequestManager.isPresent()) {
             event.future().completeExceptionally(new KafkaException("Unable to fetch committed " +
                     "offset because the CommittedRequestManager is not available. Check if group.id was set correctly"));
@@ -163,19 +178,19 @@ private void process(final FetchCommittedOffsetsApplicationEvent event) {
         }
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
         long expirationTimeMs = getExpirationTimeForTimeout(event.timeout());
-        event.chain(manager.fetchOffsets(event.partitions(), expirationTimeMs));
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future = manager.fetchOffsets(event.partitions(), expirationTimeMs);
+        future.whenComplete(complete(event.future()));
     }
 
     private void process(final NewTopicsMetadataUpdateRequestEvent ignored) {
         metadata.requestUpdateForNewTopics();
     }
 
-
     /**
      * Commit all consumed if auto-commit is enabled. Note this will trigger an async commit,
      * that will not be retried if the commit request fails.
      */
-    private void process(final AssignmentChangeApplicationEvent event) {
+    private void process(final AssignmentChangeEvent event) {
         if (!requestManagers.commitRequestManager.isPresent()) {
             return;
         }
@@ -184,11 +199,11 @@ private void process(final AssignmentChangeApplicationEvent event) {
         manager.maybeAutoCommitAsync();
     }
 
-    private void process(final ListOffsetsApplicationEvent event) {
+    private void process(final ListOffsetsEvent event) {
         final CompletableFuture<Map<TopicPartition, OffsetAndTimestamp>> future =
                 requestManagers.offsetsRequestManager.fetchOffsets(event.timestampsToSearch(),
                         event.requireTimestamps());
-        event.chain(future);
+        future.whenComplete(complete(event.future()));
     }
 
     /**
@@ -196,7 +211,7 @@ private void process(final ListOffsetsApplicationEvent event) {
      * consumer join the group if it is not part of it yet, or send the updated subscription if
      * it is already a member.
      */
-    private void process(final SubscriptionChangeApplicationEvent ignored) {
+    private void process(final SubscriptionChangeEvent ignored) {
         if (!requestManagers.heartbeatRequestManager.isPresent()) {
             log.warn("Group membership manager not present when processing a subscribe event");
             return;
@@ -213,38 +228,39 @@ private void process(final SubscriptionChangeApplicationEvent ignored) {
      *              execution for releasing the assignment completes, and the request to leave
      *              the group is sent out.
      */
-    private void process(final UnsubscribeApplicationEvent event) {
+    private void process(final UnsubscribeEvent event) {
         if (!requestManagers.heartbeatRequestManager.isPresent()) {
             KafkaException error = new KafkaException("Group membership manager not present when processing an unsubscribe event");
             event.future().completeExceptionally(error);
             return;
         }
         MembershipManager membershipManager = requestManagers.heartbeatRequestManager.get().membershipManager();
-        CompletableFuture<Void> result = membershipManager.leaveGroup();
-        event.chain(result);
+        CompletableFuture<Void> future = membershipManager.leaveGroup();
+        future.whenComplete(complete(event.future()));
     }
 
-    private void process(final ResetPositionsApplicationEvent event) {
-        CompletableFuture<Void> result = requestManagers.offsetsRequestManager.resetPositionsIfNeeded();
-        event.chain(result);
+    private void process(final ResetPositionsEvent event) {
+        CompletableFuture<Void> future = requestManagers.offsetsRequestManager.resetPositionsIfNeeded();
+        future.whenComplete(complete(event.future()));
     }
 
-    private void process(final ValidatePositionsApplicationEvent event) {
-        CompletableFuture<Void> result = requestManagers.offsetsRequestManager.validatePositionsIfNeeded();
-        event.chain(result);
+    private void process(final ValidatePositionsEvent event) {
+        CompletableFuture<Void> future = requestManagers.offsetsRequestManager.validatePositionsIfNeeded();
+        future.whenComplete(complete(event.future()));
     }
 
-    private void process(final TopicMetadataApplicationEvent event) {
-        final CompletableFuture<Map<String, List<PartitionInfo>>> future;
-
-        long expirationTimeMs = getExpirationTimeForTimeout(event.getTimeoutMs());
-        if (event.isAllTopics()) {
-            future = requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(expirationTimeMs);
-        } else {
-            future = requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), expirationTimeMs);
-        }
+    private void process(final TopicMetadataEvent event) {
+        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
+        final CompletableFuture<Map<String, List<PartitionInfo>>> future =
+                requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), expirationTimeMs);
+        future.whenComplete(complete(event.future()));
+    }
 
-        event.chain(future);
+    private void process(final AllTopicsMetadataEvent event) {
+        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
+        final CompletableFuture<Map<String, List<PartitionInfo>>> future =
+                requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(expirationTimeMs);
+        future.whenComplete(complete(event.future()));
     }
 
     private void process(final ConsumerRebalanceListenerCallbackCompletedEvent event) {
@@ -259,14 +275,14 @@ private void process(final ConsumerRebalanceListenerCallbackCompletedEvent event
         manager.consumerRebalanceListenerCallbackCompleted(event);
     }
 
-    private void process(final CommitOnCloseApplicationEvent event) {
+    private void process(final CommitOnCloseEvent event) {
         if (!requestManagers.commitRequestManager.isPresent())
             return;
         log.debug("Signal CommitRequestManager closing");
         requestManagers.commitRequestManager.get().signalClose();
     }
 
-    private void process(final LeaveOnCloseApplicationEvent event) {
+    private void process(final LeaveOnCloseEvent event) {
         if (!requestManagers.heartbeatRequestManager.isPresent()) {
             event.future().complete(null);
             return;
@@ -277,7 +293,7 @@ private void process(final LeaveOnCloseApplicationEvent event) {
         log.debug("Leaving group before closing");
         CompletableFuture<Void> future = membershipManager.leaveGroup();
         // The future will be completed on heartbeat sent
-        event.chain(future);
+        future.whenComplete(complete(event.future()));
     }
 
     /**
@@ -293,6 +309,15 @@ long getExpirationTimeForTimeout(final long timeoutMs) {
         return expiration;
     }
 
+    private <T> BiConsumer<? super T, ? super Throwable> complete(final CompletableFuture<T> b) {
+        return (value, exception) -> {
+            if (exception != null)
+                b.completeExceptionally(exception);
+            else
+                b.complete(value);
+        };
+    }
+
     /**
      * Creates a {@link Supplier} for deferred creation during invocation by
      * {@link ConsumerNetworkThread}.
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeEvent.java
similarity index 56%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeEvent.java
index ccf7199f2606b..c9efa2e9dff89 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AssignmentChangeEvent.java
@@ -22,13 +22,12 @@
 import java.util.Collections;
 import java.util.Map;
 
-public class AssignmentChangeApplicationEvent extends ApplicationEvent {
+public class AssignmentChangeEvent extends ApplicationEvent {
 
     private final Map<TopicPartition, OffsetAndMetadata> offsets;
     private final long currentTimeMs;
 
-    public AssignmentChangeApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                            final long currentTimeMs) {
+    public AssignmentChangeEvent(final Map<TopicPartition, OffsetAndMetadata> offsets, final long currentTimeMs) {
         super(Type.ASSIGNMENT_CHANGE);
         this.offsets = Collections.unmodifiableMap(offsets);
         this.currentTimeMs = currentTimeMs;
@@ -43,31 +42,7 @@ public long currentTimeMs() {
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        AssignmentChangeApplicationEvent that = (AssignmentChangeApplicationEvent) o;
-
-        if (currentTimeMs != that.currentTimeMs) return false;
-        return offsets.equals(that.offsets);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + offsets.hashCode();
-        result = 31 * result + (int) (currentTimeMs ^ (currentTimeMs >>> 32));
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "AssignmentChangeApplicationEvent{" +
-                toStringBase() +
-                ", offsets=" + offsets +
-                ", currentTimeMs=" + currentTimeMs +
-                '}';
+    protected String toStringBase() {
+        return super.toStringBase() + ", offsets=" + offsets + ", currentTimeMs=" + currentTimeMs;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
similarity index 73%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
index 7a939ce3cfd16..2f03fdfb1e543 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
@@ -18,22 +18,15 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+
 import java.util.Map;
 
 /**
  * Event to commit offsets without waiting for a response, so the request won't be retried.
  */
-public class AsyncCommitApplicationEvent extends CommitApplicationEvent {
-
-    public AsyncCommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets) {
-        super(offsets, Type.COMMIT_ASYNC);
-    }
+public class AsyncCommitEvent extends CommitEvent {
 
-    @Override
-    public String toString() {
-        return "AsyncCommitApplicationEvent{" +
-            toStringBase() +
-            ", offsets=" + offsets() +
-            '}';
+    public AsyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets) {
+        super(Type.COMMIT_ASYNC, offsets);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
index e5d522201ef0a..9bc3fbebc30a6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer.internals.events;
 
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkThread;
+import org.apache.kafka.common.Uuid;
 
 import java.util.Objects;
 
@@ -31,36 +32,45 @@ public enum Type {
 
     private final Type type;
 
-    public BackgroundEvent(Type type) {
+    /**
+     * This identifies a particular event. It is used to disambiguate events via {@link #hashCode()} and
+     * {@link #equals(Object)} and can be used in log messages when debugging.
+     */
+    private final Uuid id;
+
+    protected BackgroundEvent(Type type) {
         this.type = Objects.requireNonNull(type);
+        this.id = Uuid.randomUuid();
     }
 
     public Type type() {
         return type;
     }
 
+    public Uuid id() {
+        return id;
+    }
+
     @Override
-    public boolean equals(Object o) {
+    public final boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         BackgroundEvent that = (BackgroundEvent) o;
-
-        return type == that.type;
+        return type == that.type && id.equals(that.id);
     }
 
     @Override
-    public int hashCode() {
-        return type.hashCode();
+    public final int hashCode() {
+        return Objects.hash(type, id);
     }
 
     protected String toStringBase() {
-        return "type=" + type;
+        return "type=" + type + ", id=" + id;
     }
 
     @Override
-    public String toString() {
-        return "BackgroundEvent{" +
+    public final String toString() {
+        return getClass().getSimpleName() + "{" +
                 toStringBase() +
                 '}';
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEventHandler.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEventHandler.java
index 103493d25314f..48421484f1d3b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEventHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEventHandler.java
@@ -26,7 +26,7 @@
 /**
  * An event handler that receives {@link BackgroundEvent background events} from the
  * {@link ConsumerNetworkThread network thread} which are then made available to the application thread
- * via the {@link BackgroundEventProcessor}.
+ * via an {@link EventProcessor}.
  */
 
 public class BackgroundEventHandler {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
similarity index 71%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
index 69d969d7b0f46..253d27e2573b8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
@@ -22,14 +22,14 @@
 import java.util.Collections;
 import java.util.Map;
 
-public abstract class CommitApplicationEvent extends CompletableApplicationEvent<Void> {
+public abstract class CommitEvent extends CompletableApplicationEvent<Void> {
 
     /**
      * Offsets to commit per partition.
      */
     private final Map<TopicPartition, OffsetAndMetadata> offsets;
 
-    public CommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets, Type type) {
+    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets) {
         super(type);
         this.offsets = Collections.unmodifiableMap(offsets);
 
@@ -45,20 +45,7 @@ public Map<TopicPartition, OffsetAndMetadata> offsets() {
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        CommitApplicationEvent that = (CommitApplicationEvent) o;
-
-        return offsets.equals(that.offsets);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + offsets.hashCode();
-        return result;
+    protected String toStringBase() {
+        return super.toStringBase() + ", offsets=" + offsets;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseEvent.java
similarity index 76%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseEvent.java
index 4cc07e945f9d2..7d2e29fcedb9a 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitOnCloseEvent.java
@@ -16,16 +16,9 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
-public class CommitOnCloseApplicationEvent extends ApplicationEvent {
+public class CommitOnCloseEvent extends ApplicationEvent {
 
-    public CommitOnCloseApplicationEvent() {
+    public CommitOnCloseEvent() {
         super(Type.COMMIT_ON_CLOSE);
     }
-
-    @Override
-    public String toString() {
-        return "CommitOnCloseApplicationEvent{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
index 365c620e0c0c0..a62c3aaa4c43b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
@@ -16,9 +16,6 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
-import org.apache.kafka.clients.consumer.internals.ConsumerUtils;
-import org.apache.kafka.common.utils.Timer;
-
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -31,56 +28,18 @@ public abstract class CompletableApplicationEvent<T> extends ApplicationEvent im
 
     private final CompletableFuture<T> future;
 
-    protected CompletableApplicationEvent(Type type) {
+    protected CompletableApplicationEvent(final Type type) {
         super(type);
         this.future = new CompletableFuture<>();
     }
 
+    @Override
     public CompletableFuture<T> future() {
         return future;
     }
 
-    public T get(Timer timer) {
-        return ConsumerUtils.getResult(future, timer);
-    }
-
-    public void chain(final CompletableFuture<T> providedFuture) {
-        providedFuture.whenComplete((value, exception) -> {
-            if (exception != null) {
-                this.future.completeExceptionally(exception);
-            } else {
-                this.future.complete(value);
-            }
-        });
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        CompletableApplicationEvent<?> that = (CompletableApplicationEvent<?>) o;
-
-        return future.equals(that.future);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + future.hashCode();
-        return result;
-    }
-
     @Override
     protected String toStringBase() {
         return super.toStringBase() + ", future=" + future;
     }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + "{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableBackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableBackgroundEvent.java
index 640ee6103af9b..1a58515a5cbce 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableBackgroundEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableBackgroundEvent.java
@@ -28,42 +28,18 @@ public abstract class CompletableBackgroundEvent<T> extends BackgroundEvent impl
 
     private final CompletableFuture<T> future;
 
-    protected CompletableBackgroundEvent(Type type) {
+    protected CompletableBackgroundEvent(final Type type) {
         super(type);
         this.future = new CompletableFuture<>();
     }
 
+    @Override
     public CompletableFuture<T> future() {
         return future;
     }
 
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        CompletableBackgroundEvent<?> that = (CompletableBackgroundEvent<?>) o;
-
-        return future.equals(that.future);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + future.hashCode();
-        return result;
-    }
-
     @Override
     protected String toStringBase() {
         return super.toStringBase() + ", future=" + future;
     }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + "{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableEvent.java
index 8fdcc20fa8363..97559d8cb9be2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableEvent.java
@@ -21,5 +21,4 @@
 public interface CompletableEvent<T> {
 
     CompletableFuture<T> future();
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
index b260c6154ea5f..a10e98df1d061 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
@@ -34,9 +34,9 @@ public class ConsumerRebalanceListenerCallbackCompletedEvent extends Application
     private final CompletableFuture<Void> future;
     private final Optional<KafkaException> error;
 
-    public ConsumerRebalanceListenerCallbackCompletedEvent(ConsumerRebalanceListenerMethodName methodName,
-                                                           CompletableFuture<Void> future,
-                                                           Optional<KafkaException> error) {
+    public ConsumerRebalanceListenerCallbackCompletedEvent(final ConsumerRebalanceListenerMethodName methodName,
+                                                           final CompletableFuture<Void> future,
+                                                           final Optional<KafkaException> error) {
         super(Type.CONSUMER_REBALANCE_LISTENER_CALLBACK_COMPLETED);
         this.methodName = Objects.requireNonNull(methodName);
         this.future = Objects.requireNonNull(future);
@@ -55,24 +55,6 @@ public Optional<KafkaException> error() {
         return error;
     }
 
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        ConsumerRebalanceListenerCallbackCompletedEvent that = (ConsumerRebalanceListenerCallbackCompletedEvent) o;
-
-        return methodName == that.methodName &&
-                future.equals(that.future) &&
-                error.equals(that.error);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(methodName, future, error);
-    }
-
     @Override
     protected String toStringBase() {
         return super.toStringBase() +
@@ -80,11 +62,4 @@ protected String toStringBase() {
                 ", future=" + future +
                 ", error=" + error;
     }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + "{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackNeededEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackNeededEvent.java
index 7b17c034abdbd..6ce833580c88d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackNeededEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackNeededEvent.java
@@ -37,8 +37,8 @@ public class ConsumerRebalanceListenerCallbackNeededEvent extends CompletableBac
     private final ConsumerRebalanceListenerMethodName methodName;
     private final SortedSet<TopicPartition> partitions;
 
-    public ConsumerRebalanceListenerCallbackNeededEvent(ConsumerRebalanceListenerMethodName methodName,
-                                                        SortedSet<TopicPartition> partitions) {
+    public ConsumerRebalanceListenerCallbackNeededEvent(final ConsumerRebalanceListenerMethodName methodName,
+                                                        final SortedSet<TopicPartition> partitions) {
         super(Type.CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED);
         this.methodName = Objects.requireNonNull(methodName);
         this.partitions = Collections.unmodifiableSortedSet(partitions);
@@ -52,36 +52,10 @@ public SortedSet<TopicPartition> partitions() {
         return partitions;
     }
 
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        ConsumerRebalanceListenerCallbackNeededEvent that = (ConsumerRebalanceListenerCallbackNeededEvent) o;
-
-        return methodName == that.methodName && partitions.equals(that.partitions);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + methodName.hashCode();
-        result = 31 * result + partitions.hashCode();
-        return result;
-    }
-
     @Override
     protected String toStringBase() {
         return super.toStringBase() +
                 ", methodName=" + methodName +
                 ", partitions=" + partitions;
     }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + "{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorBackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorBackgroundEvent.java
deleted file mode 100644
index 2945f22986b18..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorBackgroundEvent.java
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer.internals.events;
-
-import org.apache.kafka.common.KafkaException;
-
-public class ErrorBackgroundEvent extends BackgroundEvent {
-
-    private final RuntimeException error;
-
-    public ErrorBackgroundEvent(Throwable t) {
-        super(Type.ERROR);
-        this.error = t instanceof RuntimeException ? (RuntimeException) t : new KafkaException(t);
-    }
-
-    public RuntimeException error() {
-        return error;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        ErrorBackgroundEvent that = (ErrorBackgroundEvent) o;
-
-        return error.equals(that.error);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + error.hashCode();
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "ErrorBackgroundEvent{" +
-            toStringBase() +
-            ", error=" + error +
-            '}';
-    }
-}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorEvent.java
new file mode 100644
index 0000000000000..5e6d822382348
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ErrorEvent.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+import org.apache.kafka.common.KafkaException;
+
+public class ErrorEvent extends BackgroundEvent {
+
+    private final RuntimeException error;
+
+    public ErrorEvent(Throwable t) {
+        super(Type.ERROR);
+        this.error = t instanceof RuntimeException ? (RuntimeException) t : new KafkaException(t);
+    }
+
+    public RuntimeException error() {
+        return error;
+    }
+
+    @Override
+    public String toStringBase() {
+        return super.toStringBase() + ", error=" + error;
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
similarity index 60%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
index 34b2d97705cd9..7cf56b990b0d9 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
@@ -23,7 +23,7 @@
 import java.util.Map;
 import java.util.Set;
 
-public class FetchCommittedOffsetsApplicationEvent extends CompletableApplicationEvent<Map<TopicPartition, OffsetAndMetadata>> {
+public class FetchCommittedOffsetsEvent extends CompletableApplicationEvent<Map<TopicPartition, OffsetAndMetadata>> {
 
     /**
      * Partitions to retrieve committed offsets for.
@@ -35,8 +35,7 @@ public class FetchCommittedOffsetsApplicationEvent extends CompletableApplicatio
      */
     private final long timeoutMs;
 
-    public FetchCommittedOffsetsApplicationEvent(final Set<TopicPartition> partitions,
-                                                 final long timeoutMs) {
+    public FetchCommittedOffsetsEvent(final Set<TopicPartition> partitions, final long timeoutMs) {
         super(Type.FETCH_COMMITTED_OFFSETS);
         this.partitions = Collections.unmodifiableSet(partitions);
         this.timeoutMs = timeoutMs;
@@ -51,29 +50,7 @@ public long timeout() {
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        FetchCommittedOffsetsApplicationEvent that = (FetchCommittedOffsetsApplicationEvent) o;
-
-        return partitions.equals(that.partitions);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + partitions.hashCode();
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + "{" +
-                toStringBase() +
-                ", partitions=" + partitions +
-                ", timeout=" + timeoutMs + "ms" +
-                '}';
+    public String toStringBase() {
+        return super.toStringBase() + ", partitions=" + partitions + ", partitions=" + partitions;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
index 120e671724209..001f549818383 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
@@ -19,8 +19,6 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkThread;
 
-import java.util.Objects;
-
 /**
  * This event is sent by the {@link ConsumerNetworkThread consumer's network thread} to the application thread
  * so that when the user calls the {@link Consumer#groupMetadata()} API, the information is up-to-date. The
@@ -29,11 +27,10 @@
  */
 public class GroupMetadataUpdateEvent extends BackgroundEvent {
 
-    final private int memberEpoch;
-    final private String memberId;
+    private final int memberEpoch;
+    private final String memberId;
 
-    public GroupMetadataUpdateEvent(final int memberEpoch,
-                                    final String memberId) {
+    public GroupMetadataUpdateEvent(final int memberEpoch, final String memberId) {
         super(Type.GROUP_METADATA_UPDATE);
         this.memberEpoch = memberEpoch;
         this.memberId = memberId;
@@ -47,33 +44,10 @@ public String memberId() {
         return memberId;
     }
 
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-        GroupMetadataUpdateEvent that = (GroupMetadataUpdateEvent) o;
-        return memberEpoch == that.memberEpoch &&
-            Objects.equals(memberId, that.memberId);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(super.hashCode(), memberEpoch, memberId);
-    }
-
     @Override
     public String toStringBase() {
         return super.toStringBase() +
             ", memberEpoch=" + memberEpoch +
             ", memberId='" + memberId + '\'';
     }
-
-    @Override
-    public String toString() {
-        return "GroupMetadataUpdateEvent{" +
-            toStringBase() +
-            '}';
-    }
-
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
similarity index 76%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
index ee0b6ffa61c7d..5ee19a7cc02da 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
@@ -16,15 +16,9 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
-public class LeaveOnCloseApplicationEvent extends CompletableApplicationEvent<Void> {
-    public LeaveOnCloseApplicationEvent() {
-        super(Type.LEAVE_ON_CLOSE);
-    }
+public class LeaveOnCloseEvent extends CompletableApplicationEvent<Void> {
 
-    @Override
-    public String toString() {
-        return "LeaveOnCloseApplicationEvent{" +
-            toStringBase() +
-            '}';
+    public LeaveOnCloseEvent() {
+        super(Type.LEAVE_ON_CLOSE);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
similarity index 69%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
index 2466d062726f8..fd3b321173f11 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
@@ -31,12 +31,12 @@
  * {@link OffsetAndTimestamp} found (offset of the first message whose timestamp is greater than
  * or equals to the target timestamp)
  */
-public class ListOffsetsApplicationEvent extends CompletableApplicationEvent<Map<TopicPartition, OffsetAndTimestamp>> {
+public class ListOffsetsEvent extends CompletableApplicationEvent<Map<TopicPartition, OffsetAndTimestamp>> {
 
     private final Map<TopicPartition, Long> timestampsToSearch;
     private final boolean requireTimestamps;
 
-    public ListOffsetsApplicationEvent(Map<TopicPartition, Long> timestampToSearch, boolean requireTimestamps) {
+    public ListOffsetsEvent(final Map<TopicPartition, Long> timestampToSearch, final boolean requireTimestamps) {
         super(Type.LIST_OFFSETS);
         this.timestampsToSearch = Collections.unmodifiableMap(timestampToSearch);
         this.requireTimestamps = requireTimestamps;
@@ -64,31 +64,10 @@ public boolean requireTimestamps() {
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        ListOffsetsApplicationEvent that = (ListOffsetsApplicationEvent) o;
-
-        if (requireTimestamps != that.requireTimestamps) return false;
-        return timestampsToSearch.equals(that.timestampsToSearch);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + timestampsToSearch.hashCode();
-        result = 31 * result + (requireTimestamps ? 1 : 0);
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + " {" +
-                toStringBase() +
-                ", timestampsToSearch=" + timestampsToSearch + ", " +
-                "requireTimestamps=" + requireTimestamps + '}';
+    public String toStringBase() {
+        return super.toStringBase() +
+                ", timestampsToSearch=" + timestampsToSearch +
+                ", requireTimestamps=" + requireTimestamps;
     }
 
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/NewTopicsMetadataUpdateRequestEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/NewTopicsMetadataUpdateRequestEvent.java
index c06a3a717dceb..b06bd456f5f14 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/NewTopicsMetadataUpdateRequestEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/NewTopicsMetadataUpdateRequestEvent.java
@@ -21,11 +21,4 @@ public class NewTopicsMetadataUpdateRequestEvent extends ApplicationEvent {
     public NewTopicsMetadataUpdateRequestEvent() {
         super(Type.NEW_TOPICS_METADATA_UPDATE);
     }
-
-    @Override
-    public String toString() {
-        return "NewTopicsMetadataUpdateRequestEvent{" +
-                toStringBase() +
-                '}';
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollApplicationEvent.java
deleted file mode 100644
index b958f0ec41703..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollApplicationEvent.java
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer.internals.events;
-
-public class PollApplicationEvent extends ApplicationEvent {
-
-    private final long pollTimeMs;
-
-    public PollApplicationEvent(final long pollTimeMs) {
-        super(Type.POLL);
-        this.pollTimeMs = pollTimeMs;
-    }
-
-    public long pollTimeMs() {
-        return pollTimeMs;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        PollApplicationEvent that = (PollApplicationEvent) o;
-
-        return pollTimeMs == that.pollTimeMs;
-    }
-
-    @Override
-    public int hashCode() {
-        int result = super.hashCode();
-        result = 31 * result + (int) (pollTimeMs ^ (pollTimeMs >>> 32));
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "PollApplicationEvent{" +
-                toStringBase() +
-                ", pollTimeMs=" + pollTimeMs +
-                '}';
-    }
-}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollEvent.java
new file mode 100644
index 0000000000000..96614c06e9b6f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/PollEvent.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+public class PollEvent extends ApplicationEvent {
+
+    private final long pollTimeMs;
+
+    public PollEvent(final long pollTimeMs) {
+        super(Type.POLL);
+        this.pollTimeMs = pollTimeMs;
+    }
+
+    public long pollTimeMs() {
+        return pollTimeMs;
+    }
+
+    @Override
+    public String toStringBase() {
+        return super.toStringBase() + ", pollTimeMs=" + pollTimeMs;
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
similarity index 89%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
index 5d9b07f9de05f..06f6ebbb68a32 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
@@ -22,9 +22,9 @@
  * asynchronous event that generates ListOffsets requests, and completes by updating in-memory
  * positions when responses are received.
  */
-public class ResetPositionsApplicationEvent extends CompletableApplicationEvent<Void> {
+public class ResetPositionsEvent extends CompletableApplicationEvent<Void> {
 
-    public ResetPositionsApplicationEvent() {
+    public ResetPositionsEvent() {
         super(Type.RESET_POSITIONS);
     }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeEvent.java
similarity index 90%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeEvent.java
index 73fd15fb14408..ad5fd34c06fcf 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SubscriptionChangeEvent.java
@@ -22,9 +22,9 @@
  * calls the subscribe API. This will make the consumer join a consumer group if not part of it
  * yet, or just send the updated subscription to the broker if it's already a member of the group.
  */
-public class SubscriptionChangeApplicationEvent extends ApplicationEvent {
+public class SubscriptionChangeEvent extends ApplicationEvent {
 
-    public SubscriptionChangeApplicationEvent() {
+    public SubscriptionChangeEvent() {
         super(Type.SUBSCRIPTION_CHANGE);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
similarity index 73%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
index 43dfee6ab18b5..7e00e0da59683 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
@@ -18,22 +18,23 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+
 import java.util.Map;
 
 /**
  * Event to commit offsets waiting for a response and retrying on expected retriable errors until
  * the timer expires.
  */
-public class SyncCommitApplicationEvent extends CommitApplicationEvent {
+public class SyncCommitEvent extends CommitEvent {
 
     /**
      * Time to wait for a response, retrying on retriable errors.
      */
     private final long retryTimeoutMs;
 
-    public SyncCommitApplicationEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                                      final long retryTimeoutMs) {
-        super(offsets, Type.COMMIT_SYNC);
+    public SyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
+                           final long retryTimeoutMs) {
+        super(Type.COMMIT_SYNC, offsets);
         this.retryTimeoutMs = retryTimeoutMs;
     }
 
@@ -42,11 +43,7 @@ public Long retryTimeoutMs() {
     }
 
     @Override
-    public String toString() {
-        return "SyncCommitApplicationEvent{" +
-            toStringBase() +
-            ", offsets=" + offsets() +
-            ", retryTimeout=" + retryTimeoutMs + "ms" +
-            '}';
+    public String toStringBase() {
+        return super.toStringBase() + ", offsets=" + offsets() + ", retryTimeoutMs=" + retryTimeoutMs;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataApplicationEvent.java
deleted file mode 100644
index dd6f842cc2674..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataApplicationEvent.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer.internals.events;
-
-import org.apache.kafka.common.PartitionInfo;
-
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-
-public class TopicMetadataApplicationEvent extends CompletableApplicationEvent<Map<String, List<PartitionInfo>>> {
-    private final String topic;
-    private final boolean allTopics;
-    private final long timeoutMs;
-
-    public TopicMetadataApplicationEvent(final long timeoutMs) {
-        super(Type.TOPIC_METADATA);
-        this.topic = null;
-        this.allTopics = true;
-        this.timeoutMs = timeoutMs;
-    }
-
-    public TopicMetadataApplicationEvent(final String topic, final long timeoutMs) {
-        super(Type.TOPIC_METADATA);
-        this.topic = topic;
-        this.allTopics = false;
-        this.timeoutMs = timeoutMs;
-    }
-
-    public String topic() {
-        return topic;
-    }
-
-    public boolean isAllTopics() {
-        return allTopics;
-    }
-
-    public long getTimeoutMs() {
-        return timeoutMs;
-    }
-    @Override
-    public String toString() {
-        return getClass().getSimpleName() + " {" + toStringBase() +
-                ", topic=" + topic +
-                ", allTopics=" + allTopics +
-                ", timeoutMs=" + timeoutMs + "}";
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (!(o instanceof TopicMetadataApplicationEvent)) return false;
-        if (!super.equals(o)) return false;
-
-        TopicMetadataApplicationEvent that = (TopicMetadataApplicationEvent) o;
-
-        return topic.equals(that.topic) && (allTopics == that.allTopics) && (timeoutMs == that.timeoutMs);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(super.hashCode(), topic, allTopics, timeoutMs);
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
new file mode 100644
index 0000000000000..ebbb2a6c46892
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals.events;
+
+import java.util.Objects;
+
+public class TopicMetadataEvent extends AbstractTopicMetadataEvent {
+
+    private final String topic;
+
+    public TopicMetadataEvent(final String topic, final long timeoutMs) {
+        super(Type.TOPIC_METADATA, timeoutMs);
+        this.topic = Objects.requireNonNull(topic);
+    }
+
+    public String topic() {
+        return topic;
+    }
+
+    @Override
+    public String toStringBase() {
+        return super.toStringBase() + ", topic=" + topic;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
similarity index 91%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
index a1ccb896fdf57..07af36e5feb85 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
@@ -24,8 +24,9 @@
  * complete and the heartbeat to leave the group is sent out (minimal effort to send the
  * leave group heartbeat, without waiting for any response or considering timeouts).
  */
-public class UnsubscribeApplicationEvent extends CompletableApplicationEvent<Void> {
-    public UnsubscribeApplicationEvent() {
+public class UnsubscribeEvent extends CompletableApplicationEvent<Void> {
+
+    public UnsubscribeEvent() {
         super(Type.UNSUBSCRIBE);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
similarity index 89%
rename from clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsApplicationEvent.java
rename to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
index 3b093e0b68353..efa358b4c7882 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
@@ -22,9 +22,9 @@
  * detected. This is an asynchronous event that generates OffsetForLeaderEpoch requests, and
  * completes by validating in-memory positions against the offsets received in the responses.
  */
-public class ValidatePositionsApplicationEvent extends CompletableApplicationEvent<Void> {
+public class ValidatePositionsEvent extends CompletableApplicationEvent<Void> {
 
-    public ValidatePositionsApplicationEvent() {
+    public ValidatePositionsEvent() {
         super(Type.VALIDATE_POSITIONS);
     }
 }
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 976677dec8628..2d6899612f056 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -31,26 +31,26 @@
 import org.apache.kafka.clients.consumer.RetriableCommitFailedException;
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
-import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
-import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
-import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ListOffsetsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseEvent;
+import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
-import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.PollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ResetPositionsEvent;
+import org.apache.kafka.clients.consumer.internals.events.SubscriptionChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitEvent;
+import org.apache.kafka.clients.consumer.internals.events.UnsubscribeEvent;
+import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsEvent;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
@@ -256,9 +256,9 @@ public void testCommitAsyncWithNullCallback() {
 
         consumer.commitAsync(offsets, null);
 
-        final ArgumentCaptor<AsyncCommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitApplicationEvent.class);
+        final ArgumentCaptor<AsyncCommitEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitEvent.class);
         verify(applicationEventHandler).add(commitEventCaptor.capture());
-        final AsyncCommitApplicationEvent commitEvent = commitEventCaptor.getValue();
+        final AsyncCommitEvent commitEvent = commitEventCaptor.getValue();
         assertEquals(offsets, commitEvent.offsets());
         assertDoesNotThrow(() -> commitEvent.future().complete(null));
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, null));
@@ -310,9 +310,9 @@ public void testCommitAsyncWithFencedException() {
 
         assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
 
-        final ArgumentCaptor<AsyncCommitApplicationEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitApplicationEvent.class);
+        final ArgumentCaptor<AsyncCommitEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitEvent.class);
         verify(applicationEventHandler).add(commitEventCaptor.capture());
-        final AsyncCommitApplicationEvent commitEvent = commitEventCaptor.getValue();
+        final AsyncCommitEvent commitEvent = commitEventCaptor.getValue();
         commitEvent.future().completeExceptionally(Errors.FENCED_INSTANCE_ID.exception());
 
         assertThrows(Errors.FENCED_INSTANCE_ID.exception().getClass(), () -> consumer.commitAsync());
@@ -325,7 +325,7 @@ public void testCommitted() {
         completeFetchedCommittedOffsetApplicationEventSuccessfully(topicPartitionOffsets);
 
         assertEquals(topicPartitionOffsets, consumer.committed(topicPartitionOffsets.keySet(), Duration.ofMillis(1000)));
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsApplicationEvent.class), any());
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class), any());
         final Metric metric = consumer.metrics()
             .get(consumer.metricsRegistry().metricName("committed-time-ns-total", "consumer-metrics"));
         assertTrue((double) metric.metricValue() > 0);
@@ -347,7 +347,7 @@ public void testCommittedLeaderEpochUpdate() {
 
         verify(metadata).updateLastSeenEpochIfNewer(t0, 2);
         verify(metadata).updateLastSeenEpochIfNewer(t2, 3);
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsApplicationEvent.class), any());
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class), any());
     }
 
     @Test
@@ -355,9 +355,9 @@ public void testCommittedExceptionThrown() {
         consumer = newConsumer();
         Map<TopicPartition, OffsetAndMetadata> offsets = mockTopicPartitionOffset();
         when(applicationEventHandler.addAndGet(
-            any(FetchCommittedOffsetsApplicationEvent.class), any())).thenAnswer(invocation -> {
+            any(FetchCommittedOffsetsEvent.class), any())).thenAnswer(invocation -> {
                 CompletableApplicationEvent<?> event = invocation.getArgument(0);
-                assertInstanceOf(FetchCommittedOffsetsApplicationEvent.class, event);
+                assertInstanceOf(FetchCommittedOffsetsEvent.class, event);
                 throw new KafkaException("Test exception");
             });
 
@@ -530,7 +530,7 @@ public void testCommitSyncLeaderEpochUpdate() {
 
         verify(metadata).updateLastSeenEpochIfNewer(t0, 2);
         verify(metadata).updateLastSeenEpochIfNewer(t1, 1);
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitEvent.class));
     }
 
     @Test
@@ -564,7 +564,7 @@ public void testCommitAsyncLeaderEpochUpdate() {
 
         verify(metadata).updateLastSeenEpochIfNewer(t0, 2);
         verify(metadata).updateLastSeenEpochIfNewer(t1, 1);
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
     }
 
     @Test
@@ -598,8 +598,8 @@ public void testVerifyApplicationEventOnShutdown() {
         consumer = newConsumer();
         doReturn(null).when(applicationEventHandler).addAndGet(any(), any());
         consumer.close();
-        verify(applicationEventHandler).addAndGet(any(LeaveOnCloseApplicationEvent.class), any());
-        verify(applicationEventHandler).add(any(CommitOnCloseApplicationEvent.class));
+        verify(applicationEventHandler).addAndGet(any(LeaveOnCloseEvent.class), any());
+        verify(applicationEventHandler).add(any(CommitOnCloseEvent.class));
     }
 
     @Test
@@ -641,7 +641,7 @@ public void testFailedPartitionRevocationOnClose() {
         subscriptions.assignFromSubscribed(singleton(tp));
         doThrow(new KafkaException()).when(listener).onPartitionsRevoked(eq(singleton(tp)));
         assertThrows(KafkaException.class, () -> consumer.close(Duration.ZERO));
-        verify(applicationEventHandler, never()).addAndGet(any(LeaveOnCloseApplicationEvent.class), any());
+        verify(applicationEventHandler, never()).addAndGet(any(LeaveOnCloseEvent.class), any());
         verify(listener).onPartitionsRevoked(eq(singleton(tp)));
         assertEquals(emptySet(), subscriptions.assignedPartitions());
     }
@@ -677,7 +677,7 @@ public void testAutoCommitSyncEnabled() {
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
         consumer.maybeAutoCommitSync(true, time.timer(100), null);
-        verify(applicationEventHandler).add(any(SyncCommitApplicationEvent.class));
+        verify(applicationEventHandler).add(any(SyncCommitEvent.class));
     }
 
     @Test
@@ -695,7 +695,7 @@ public void testAutoCommitSyncDisabled() {
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
         consumer.maybeAutoCommitSync(false, time.timer(100), null);
-        verify(applicationEventHandler, never()).add(any(SyncCommitApplicationEvent.class));
+        verify(applicationEventHandler, never()).add(any(SyncCommitEvent.class));
     }
 
     private void assertMockCommitCallbackInvoked(final Executable task,
@@ -729,7 +729,7 @@ public void testAssign() {
         consumer.assign(singleton(tp));
         assertTrue(consumer.subscription().isEmpty());
         assertTrue(consumer.assignment().contains(tp));
-        verify(applicationEventHandler).add(any(AssignmentChangeApplicationEvent.class));
+        verify(applicationEventHandler).add(any(AssignmentChangeEvent.class));
         verify(applicationEventHandler).add(any(NewTopicsMetadataUpdateRequestEvent.class));
     }
 
@@ -781,7 +781,7 @@ public void testBeginningOffsets() {
         Map<TopicPartition, Long> expectedOffsets = expectedOffsetsAndTimestamp.entrySet().stream()
             .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().offset()));
         assertEquals(expectedOffsets, result);
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsApplicationEvent.class),
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class),
             ArgumentMatchers.isA(Timer.class));
     }
 
@@ -792,13 +792,13 @@ public void testBeginningOffsetsThrowsKafkaExceptionForUnderlyingExecutionFailur
         Throwable eventProcessingFailure = new KafkaException("Unexpected failure " +
             "processing List Offsets event");
         doThrow(eventProcessingFailure).when(applicationEventHandler).addAndGet(
-            any(ListOffsetsApplicationEvent.class),
+            any(ListOffsetsEvent.class),
             any());
         Throwable consumerError = assertThrows(KafkaException.class,
             () -> consumer.beginningOffsets(partitions,
                 Duration.ofMillis(1)));
         assertEquals(eventProcessingFailure, consumerError);
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsApplicationEvent.class),
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class),
             ArgumentMatchers.isA(Timer.class));
     }
 
@@ -810,7 +810,7 @@ public void testBeginningOffsetsTimeoutOnEventProcessingTimeout() {
             () -> consumer.beginningOffsets(
                 Collections.singletonList(new TopicPartition("t1", 0)),
                 Duration.ofMillis(1)));
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsApplicationEvent.class),
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class),
             ArgumentMatchers.isA(Timer.class));
     }
 
@@ -850,7 +850,7 @@ public void testOffsetsForTimes() {
         Map<TopicPartition, OffsetAndTimestamp> result =
                 assertDoesNotThrow(() -> consumer.offsetsForTimes(timestampToSearch, Duration.ofMillis(1)));
         assertEquals(expectedResult, result);
-        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsApplicationEvent.class),
+        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class),
                 ArgumentMatchers.isA(Timer.class));
     }
 
@@ -869,7 +869,7 @@ public void testOffsetsForTimesWithZeroTimeout() {
                 assertDoesNotThrow(() -> consumer.offsetsForTimes(timestampToSearch,
                         Duration.ZERO));
         assertEquals(expectedResult, result);
-        verify(applicationEventHandler, never()).addAndGet(ArgumentMatchers.isA(ListOffsetsApplicationEvent.class),
+        verify(applicationEventHandler, never()).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class),
             ArgumentMatchers.isA(Timer.class));
     }
 
@@ -880,12 +880,12 @@ public void testWakeupCommitted() {
         doAnswer(invocation -> {
             CompletableApplicationEvent<?> event = invocation.getArgument(0);
             Timer timer = invocation.getArgument(1);
-            assertInstanceOf(FetchCommittedOffsetsApplicationEvent.class, event);
+            assertInstanceOf(FetchCommittedOffsetsEvent.class, event);
             assertTrue(event.future().isCompletedExceptionally());
             return ConsumerUtils.getResult(event.future(), timer);
         })
             .when(applicationEventHandler)
-            .addAndGet(any(FetchCommittedOffsetsApplicationEvent.class), any(Timer.class));
+            .addAndGet(any(FetchCommittedOffsetsEvent.class), any(Timer.class));
 
         consumer.wakeup();
         assertThrows(WakeupException.class, () -> consumer.committed(offsets.keySet()));
@@ -1011,7 +1011,7 @@ public void testSubscribeGeneratesEvent() {
         consumer.subscribe(singletonList(topic));
         assertEquals(singleton(topic), consumer.subscription());
         assertTrue(consumer.assignment().isEmpty());
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(SubscriptionChangeApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(SubscriptionChangeEvent.class));
     }
 
     @Test
@@ -1023,7 +1023,7 @@ public void testUnsubscribeGeneratesUnsubscribeEvent() {
 
         assertTrue(consumer.subscription().isEmpty());
         assertTrue(consumer.assignment().isEmpty());
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
     }
 
     @Test
@@ -1034,7 +1034,7 @@ public void testSubscribeToEmptyListActsAsUnsubscribe() {
         consumer.subscribe(Collections.emptyList());
         assertTrue(consumer.subscription().isEmpty());
         assertTrue(consumer.assignment().isEmpty());
-        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeApplicationEvent.class));
+        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
     }
 
     @Test
@@ -1356,8 +1356,8 @@ public void testBackgroundError() {
         consumer = newConsumer(config);
 
         final KafkaException expectedException = new KafkaException("Nobody expects the Spanish Inquisition");
-        final ErrorBackgroundEvent errorBackgroundEvent = new ErrorBackgroundEvent(expectedException);
-        backgroundEventQueue.add(errorBackgroundEvent);
+        final ErrorEvent errorEvent = new ErrorEvent(expectedException);
+        backgroundEventQueue.add(errorEvent);
         consumer.assign(singletonList(new TopicPartition("topic", 0)));
         final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
 
@@ -1371,11 +1371,11 @@ public void testMultipleBackgroundErrors() {
         consumer = newConsumer(config);
 
         final KafkaException expectedException1 = new KafkaException("Nobody expects the Spanish Inquisition");
-        final ErrorBackgroundEvent errorBackgroundEvent1 = new ErrorBackgroundEvent(expectedException1);
-        backgroundEventQueue.add(errorBackgroundEvent1);
+        final ErrorEvent errorEvent1 = new ErrorEvent(expectedException1);
+        backgroundEventQueue.add(errorEvent1);
         final KafkaException expectedException2 = new KafkaException("Spam, Spam, Spam");
-        final ErrorBackgroundEvent errorBackgroundEvent2 = new ErrorBackgroundEvent(expectedException2);
-        backgroundEventQueue.add(errorBackgroundEvent2);
+        final ErrorEvent errorEvent2 = new ErrorEvent(expectedException2);
+        backgroundEventQueue.add(errorEvent2);
         consumer.assign(singletonList(new TopicPartition("topic", 0)));
         final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
 
@@ -1472,7 +1472,7 @@ public void testEnsurePollEventSentOnConsumerPoll() {
 
         consumer.subscribe(singletonList("topic1"));
         consumer.poll(Duration.ofMillis(100));
-        verify(applicationEventHandler).add(any(PollApplicationEvent.class));
+        verify(applicationEventHandler).add(any(PollEvent.class));
     }
 
     private void testInvalidGroupId(final String groupId) {
@@ -1510,20 +1510,20 @@ private void testUpdateFetchPositionsWithFetchCommittedOffsetsTimeout(boolean co
         consumer.poll(Duration.ZERO);
 
         verify(applicationEventHandler, atLeast(1))
-            .addAndGet(ArgumentMatchers.isA(ValidatePositionsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+            .addAndGet(ArgumentMatchers.isA(ValidatePositionsEvent.class), ArgumentMatchers.isA(Timer.class));
 
         if (committedOffsetsEnabled) {
             // Verify there was an FetchCommittedOffsets event and no ResetPositions event
             verify(applicationEventHandler, atLeast(1))
-                .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+                .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class), ArgumentMatchers.isA(Timer.class));
             verify(applicationEventHandler, never())
-                .addAndGet(ArgumentMatchers.isA(ResetPositionsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+                .addAndGet(ArgumentMatchers.isA(ResetPositionsEvent.class), ArgumentMatchers.isA(Timer.class));
         } else {
             // Verify there was not any FetchCommittedOffsets event but there should be a ResetPositions
             verify(applicationEventHandler, never())
-                .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+                .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class), ArgumentMatchers.isA(Timer.class));
             verify(applicationEventHandler, atLeast(1))
-                .addAndGet(ArgumentMatchers.isA(ResetPositionsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+                .addAndGet(ArgumentMatchers.isA(ResetPositionsEvent.class), ArgumentMatchers.isA(Timer.class));
         }
     }
 
@@ -1538,11 +1538,11 @@ private void testRefreshCommittedOffsetsSuccess(Set<TopicPartition> partitions,
         consumer.poll(Duration.ZERO);
 
         verify(applicationEventHandler, atLeast(1))
-            .addAndGet(ArgumentMatchers.isA(ValidatePositionsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+            .addAndGet(ArgumentMatchers.isA(ValidatePositionsEvent.class), ArgumentMatchers.isA(Timer.class));
         verify(applicationEventHandler, atLeast(1))
-            .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+            .addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class), ArgumentMatchers.isA(Timer.class));
         verify(applicationEventHandler, atLeast(1))
-            .addAndGet(ArgumentMatchers.isA(ResetPositionsApplicationEvent.class), ArgumentMatchers.isA(Timer.class));
+            .addAndGet(ArgumentMatchers.isA(ResetPositionsEvent.class), ArgumentMatchers.isA(Timer.class));
     }
 
     @Test
@@ -1697,54 +1697,54 @@ private HashMap<TopicPartition, Long> mockTimestampToSearch() {
 
     private void completeCommitAsyncApplicationEventExceptionally(Exception ex) {
         doAnswer(invocation -> {
-            AsyncCommitApplicationEvent event = invocation.getArgument(0);
+            AsyncCommitEvent event = invocation.getArgument(0);
             event.future().completeExceptionally(ex);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
     }
 
     private void completeCommitSyncApplicationEventExceptionally(Exception ex) {
         doAnswer(invocation -> {
-            SyncCommitApplicationEvent event = invocation.getArgument(0);
+            SyncCommitEvent event = invocation.getArgument(0);
             event.future().completeExceptionally(ex);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitEvent.class));
     }
 
     private void completeCommitAsyncApplicationEventSuccessfully() {
         doAnswer(invocation -> {
-            AsyncCommitApplicationEvent event = invocation.getArgument(0);
+            AsyncCommitEvent event = invocation.getArgument(0);
             event.future().complete(null);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
     }
 
     private void completeCommitSyncApplicationEventSuccessfully() {
         doAnswer(invocation -> {
-            SyncCommitApplicationEvent event = invocation.getArgument(0);
+            SyncCommitEvent event = invocation.getArgument(0);
             event.future().complete(null);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitEvent.class));
     }
 
     private void completeFetchedCommittedOffsetApplicationEventSuccessfully(final Map<TopicPartition, OffsetAndMetadata> committedOffsets) {
         doReturn(committedOffsets)
             .when(applicationEventHandler)
-            .addAndGet(any(FetchCommittedOffsetsApplicationEvent.class), any(Timer.class));
+            .addAndGet(any(FetchCommittedOffsetsEvent.class), any(Timer.class));
     }
 
     private void completeFetchedCommittedOffsetApplicationEventExceptionally(Exception ex) {
         doThrow(ex)
             .when(applicationEventHandler)
-            .addAndGet(any(FetchCommittedOffsetsApplicationEvent.class), any(Timer.class));
+            .addAndGet(any(FetchCommittedOffsetsEvent.class), any(Timer.class));
     }
 
     private void completeUnsubscribeApplicationEventSuccessfully() {
         doAnswer(invocation -> {
-            UnsubscribeApplicationEvent event = invocation.getArgument(0);
+            UnsubscribeEvent event = invocation.getArgument(0);
             event.future().complete(null);
             return null;
-        }).when(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeApplicationEvent.class));
+        }).when(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
     }
 
     private void forceCommitCallbackInvocation() {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
index a491df417de45..cbd56d8b5e2bd 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
@@ -20,16 +20,16 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
-import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.AsyncCommitApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.AsyncCommitEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ListOffsetsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
-import org.apache.kafka.clients.consumer.internals.events.PollApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ResetPositionsApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.SyncCommitApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.TopicMetadataApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.PollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ResetPositionsEvent;
+import org.apache.kafka.clients.consumer.internals.events.SyncCommitEvent;
+import org.apache.kafka.clients.consumer.internals.events.TopicMetadataEvent;
+import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsEvent;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.message.FindCoordinatorRequestData;
@@ -137,7 +137,7 @@ public void testStartupAndTearDown() throws InterruptedException {
 
     @Test
     public void testApplicationEvent() {
-        ApplicationEvent e = new PollApplicationEvent(100);
+        ApplicationEvent e = new PollEvent(100);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor, times(1)).process(e);
@@ -153,36 +153,36 @@ public void testMetadataUpdateEvent() {
 
     @Test
     public void testAsyncCommitEvent() {
-        ApplicationEvent e = new AsyncCommitApplicationEvent(new HashMap<>());
+        ApplicationEvent e = new AsyncCommitEvent(new HashMap<>());
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(AsyncCommitApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(AsyncCommitEvent.class));
     }
 
     @Test
     public void testSyncCommitEvent() {
-        ApplicationEvent e = new SyncCommitApplicationEvent(new HashMap<>(), 100L);
+        ApplicationEvent e = new SyncCommitEvent(new HashMap<>(), 100L);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(SyncCommitApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(SyncCommitEvent.class));
     }
 
     @Test
     public void testListOffsetsEventIsProcessed() {
         Map<TopicPartition, Long> timestamps = Collections.singletonMap(new TopicPartition("topic1", 1), 5L);
-        ApplicationEvent e = new ListOffsetsApplicationEvent(timestamps, true);
+        ApplicationEvent e = new ListOffsetsEvent(timestamps, true);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(ListOffsetsApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(ListOffsetsEvent.class));
         assertTrue(applicationEventsQueue.isEmpty());
     }
 
     @Test
     public void testResetPositionsEventIsProcessed() {
-        ResetPositionsApplicationEvent e = new ResetPositionsApplicationEvent();
+        ResetPositionsEvent e = new ResetPositionsEvent();
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(ResetPositionsApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(ResetPositionsEvent.class));
         assertTrue(applicationEventsQueue.isEmpty());
     }
 
@@ -190,19 +190,19 @@ public void testResetPositionsEventIsProcessed() {
     public void testResetPositionsProcessFailureIsIgnored() {
         doThrow(new NullPointerException()).when(offsetsRequestManager).resetPositionsIfNeeded();
 
-        ResetPositionsApplicationEvent event = new ResetPositionsApplicationEvent();
+        ResetPositionsEvent event = new ResetPositionsEvent();
         applicationEventsQueue.add(event);
         assertDoesNotThrow(() -> consumerNetworkThread.runOnce());
 
-        verify(applicationEventProcessor).process(any(ResetPositionsApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(ResetPositionsEvent.class));
     }
 
     @Test
     public void testValidatePositionsEventIsProcessed() {
-        ValidatePositionsApplicationEvent e = new ValidatePositionsApplicationEvent();
+        ValidatePositionsEvent e = new ValidatePositionsEvent();
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(ValidatePositionsApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(ValidatePositionsEvent.class));
         assertTrue(applicationEventsQueue.isEmpty());
     }
 
@@ -211,11 +211,11 @@ public void testAssignmentChangeEvent() {
         HashMap<TopicPartition, OffsetAndMetadata> offset = mockTopicPartitionOffset();
 
         final long currentTimeMs = time.milliseconds();
-        ApplicationEvent e = new AssignmentChangeApplicationEvent(offset, currentTimeMs);
+        ApplicationEvent e = new AssignmentChangeEvent(offset, currentTimeMs);
         applicationEventsQueue.add(e);
 
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(AssignmentChangeApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(AssignmentChangeEvent.class));
         verify(networkClient, times(1)).poll(anyLong(), anyLong());
         verify(commitRequestManager, times(1)).updateAutoCommitTimer(currentTimeMs);
         // Assignment change should generate an async commit (not retried).
@@ -224,9 +224,9 @@ public void testAssignmentChangeEvent() {
 
     @Test
     void testFetchTopicMetadata() {
-        applicationEventsQueue.add(new TopicMetadataApplicationEvent("topic", Long.MAX_VALUE));
+        applicationEventsQueue.add(new TopicMetadataEvent("topic", Long.MAX_VALUE));
         consumerNetworkThread.runOnce();
-        verify(applicationEventProcessor).process(any(TopicMetadataApplicationEvent.class));
+        verify(applicationEventProcessor).process(any(TopicMetadataEvent.class));
     }
 
     @Test
@@ -283,8 +283,8 @@ void testEnsureEventsAreCompleted() {
         coordinatorRequestManager.markCoordinatorUnknown("test", time.milliseconds());
         client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, "group-id", node));
         prepareOffsetCommitRequest(new HashMap<>(), Errors.NONE, false);
-        CompletableApplicationEvent<Void> event1 = spy(new AsyncCommitApplicationEvent(Collections.emptyMap()));
-        ApplicationEvent event2 = new AsyncCommitApplicationEvent(Collections.emptyMap());
+        CompletableApplicationEvent<Void> event1 = spy(new AsyncCommitEvent(Collections.emptyMap()));
+        ApplicationEvent event2 = new AsyncCommitEvent(Collections.emptyMap());
         CompletableFuture<Void> future = new CompletableFuture<>();
         when(event1.future()).thenReturn(future);
         applicationEventsQueue.add(event1);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
index d7ad1b55738c4..afee3fff39675 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.TimeoutException;
@@ -114,10 +114,10 @@ public void testPropagateAndBackoffAfterFatalError() {
         expectFindCoordinatorRequest(coordinatorManager, Errors.GROUP_AUTHORIZATION_FAILED);
 
         verify(backgroundEventHandler).add(argThat(backgroundEvent -> {
-            if (!(backgroundEvent instanceof ErrorBackgroundEvent))
+            if (!(backgroundEvent instanceof ErrorEvent))
                 return false;
 
-            RuntimeException exception = ((ErrorBackgroundEvent) backgroundEvent).error();
+            RuntimeException exception = ((ErrorEvent) backgroundEvent).error();
 
             if (!(exception instanceof GroupAuthorizationException))
                 return false;
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 4d4492bcb47e4..90cfb90bb1bda 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -328,12 +328,8 @@ public void testValidateConsumerGroupHeartbeatRequest(final short version) {
     @Test
     public void testConsumerGroupMetadataFirstUpdate() {
         final GroupMetadataUpdateEvent groupMetadataUpdateEvent = makeFirstGroupMetadataUpdate(memberId, memberEpoch);
-
-        final GroupMetadataUpdateEvent expectedGroupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            memberEpoch,
-            memberId
-        );
-        assertEquals(expectedGroupMetadataUpdateEvent, groupMetadataUpdateEvent);
+        assertEquals(memberEpoch, groupMetadataUpdateEvent.memberEpoch());
+        assertEquals(memberId, groupMetadataUpdateEvent.memberId());
     }
 
     @Test
@@ -370,11 +366,8 @@ public void testConsumerGroupMetadataUpdateWithMemberIdNullButMemberEpochUpdated
         final BackgroundEvent eventWithUpdatedMemberEpoch = backgroundEventQueue.poll();
         assertEquals(BackgroundEvent.Type.GROUP_METADATA_UPDATE, eventWithUpdatedMemberEpoch.type());
         final GroupMetadataUpdateEvent groupMetadataUpdateEvent = (GroupMetadataUpdateEvent) eventWithUpdatedMemberEpoch;
-        final GroupMetadataUpdateEvent expectedGroupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            updatedMemberEpoch,
-            memberId
-        );
-        assertEquals(expectedGroupMetadataUpdateEvent, groupMetadataUpdateEvent);
+        assertEquals(updatedMemberEpoch, groupMetadataUpdateEvent.memberEpoch());
+        assertEquals(memberId, groupMetadataUpdateEvent.memberId());
     }
 
     @Test
@@ -398,11 +391,8 @@ public void testConsumerGroupMetadataUpdateWithMemberIdUpdatedAndMemberEpochSame
         final BackgroundEvent eventWithUpdatedMemberEpoch = backgroundEventQueue.poll();
         assertEquals(BackgroundEvent.Type.GROUP_METADATA_UPDATE, eventWithUpdatedMemberEpoch.type());
         final GroupMetadataUpdateEvent groupMetadataUpdateEvent = (GroupMetadataUpdateEvent) eventWithUpdatedMemberEpoch;
-        final GroupMetadataUpdateEvent expectedGroupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            memberEpoch,
-            updatedMemberId
-        );
-        assertEquals(expectedGroupMetadataUpdateEvent, groupMetadataUpdateEvent);
+        assertEquals(memberEpoch, groupMetadataUpdateEvent.memberEpoch());
+        assertEquals(updatedMemberId, groupMetadataUpdateEvent.memberId());
     }
 
     private GroupMetadataUpdateEvent makeFirstGroupMetadataUpdate(final String memberId, final int memberEpoch) {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 8a0fcf85758e1..49c5819be7b79 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -1890,10 +1890,10 @@ private Map<Uuid, SortedSet<Integer>> topicIdPartitionsMap(Uuid topicId, int...
     }
 
     private ConsumerRebalanceListenerCallbackCompletedEvent performCallback(MembershipManagerImpl membershipManager,
-                                 ConsumerRebalanceListenerInvoker invoker,
-                                 ConsumerRebalanceListenerMethodName expectedMethodName,
-                                 SortedSet<TopicPartition> expectedPartitions,
-                                 boolean complete) {
+                                                                            ConsumerRebalanceListenerInvoker invoker,
+                                                                            ConsumerRebalanceListenerMethodName expectedMethodName,
+                                                                            SortedSet<TopicPartition> expectedPartitions,
+                                                                            boolean complete) {
         // We expect only our enqueued event in the background queue.
         assertEquals(1, backgroundEventQueue.size());
         assertNotNull(backgroundEventQueue.peek());
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
index 5ca034d636029..59ed7d440e264 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
@@ -24,7 +24,7 @@
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.ClusterResource;
 import org.apache.kafka.common.IsolationLevel;
@@ -553,8 +553,8 @@ public void testResetOffsetsAuthorizationFailure() {
         assertNotNull(event);
 
         // Check that the event itself is of the expected type
-        assertInstanceOf(ErrorBackgroundEvent.class, event);
-        ErrorBackgroundEvent errorEvent = (ErrorBackgroundEvent) event;
+        assertInstanceOf(ErrorEvent.class, event);
+        ErrorEvent errorEvent = (ErrorEvent) event;
         assertNotNull(errorEvent.error());
 
         // Check that the error held in the event is of the expected type
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
index 8ea8cb7a729f5..f3e2557ae9417 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
@@ -92,7 +92,7 @@ public void setup() {
     public void testPrepClosingCommitEvents() {
         List<NetworkClientDelegate.UnsentRequest> results = mockCommitResults();
         doReturn(new NetworkClientDelegate.PollResult(100, results)).when(commitRequestManager).pollOnClose();
-        processor.process(new CommitOnCloseApplicationEvent());
+        processor.process(new CommitOnCloseEvent());
         verify(commitRequestManager).signalClose();
     }
 
@@ -107,7 +107,7 @@ public void testExpirationCalculation() {
 
     @Test
     public void testPrepClosingLeaveGroupEvent() {
-        LeaveOnCloseApplicationEvent event = new LeaveOnCloseApplicationEvent();
+        LeaveOnCloseEvent event = new LeaveOnCloseEvent();
         when(heartbeatRequestManager.membershipManager()).thenReturn(membershipManager);
         when(membershipManager.leaveGroup()).thenReturn(CompletableFuture.completedFuture(null));
         processor.process(event);

From d066b94c8103cca166d7ea01a4b5bf5f65a3b838 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Thu, 29 Feb 2024 07:01:21 -0800
Subject: [PATCH 093/258] MINOR: Fix UpdatedImage and HighWatermarkUpdated
 events' logs (#15432)

I have noticed the following log when a __consumer_offsets partition immigrate from a broker. It appends because the event is queued up after the event that unloads the state machine. This patch fixes it and fixes another similar one.

```
[2024-02-06 17:14:51,359] ERROR [GroupCoordinator id=1] Execution of UpdateImage(tp=__consumer_offsets-28, offset=13251) failed due to This is not the correct coordinator.. (org.apache.kafka.coordinator.group.runtime.CoordinatorRuntime)
org.apache.kafka.common.errors.NotCoordinatorException: This is not the correct coordinator.
```

Reviewers: Justine Olshan <jolshan@confluent.io>
---
 .../group/runtime/CoordinatorRuntime.java     | 144 +++++++++---------
 1 file changed, 76 insertions(+), 68 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
index b0be84c7a900b..ccb4caf04f98e 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
@@ -1192,11 +1192,28 @@ public void onHighWatermarkUpdated(
         ) {
             log.debug("High watermark of {} incremented to {}.", tp, offset);
             scheduleInternalOperation("HighWatermarkUpdated(tp=" + tp + ", offset=" + offset + ")", tp, () -> {
-                withActiveContextOrThrow(tp, context -> {
-                    context.coordinator.updateLastCommittedOffset(offset);
-                    context.deferredEventQueue.completeUpTo(offset);
-                    coordinatorMetrics.onUpdateLastCommittedOffset(tp, offset);
-                });
+                CoordinatorContext context = coordinators.get(tp);
+                if (context != null) {
+                    context.lock.lock();
+                    try {
+                        if (context.state == CoordinatorState.ACTIVE) {
+                            // The updated high watermark can be applied to the coordinator only if the coordinator
+                            // exists and is in the active state.
+                            log.debug("Updating high watermark of {} to {}.", tp, offset);
+                            context.coordinator.updateLastCommittedOffset(offset);
+                            context.deferredEventQueue.completeUpTo(offset);
+                            coordinatorMetrics.onUpdateLastCommittedOffset(tp, offset);
+                        } else {
+                            log.debug("Ignored high watermark updated for {} to {} because the coordinator is not active.",
+                                tp, offset);
+                        }
+                    } finally {
+                        context.lock.unlock();
+                    }
+                } else {
+                    log.debug("Ignored high watermark updated for {} to {} because the coordinator does not exist.",
+                        tp, offset);
+                }
             });
         }
     }
@@ -1350,14 +1367,11 @@ private void enqueue(CoordinatorEvent event) {
     }
 
     /**
-     * Creates the context if it does not exist.
-     *
-     * @param tp    The topic partition.
-     *
-     * Visible for testing.
+     * @return The coordinator context or a new context if it does not exist.
+     * Package private for testing.
      */
-    void maybeCreateContext(TopicPartition tp) {
-        coordinators.computeIfAbsent(tp, CoordinatorContext::new);
+    CoordinatorContext maybeCreateContext(TopicPartition tp) {
+        return coordinators.computeIfAbsent(tp, CoordinatorContext::new);
     }
 
     /**
@@ -1376,29 +1390,6 @@ CoordinatorContext contextOrThrow(TopicPartition tp) throws NotCoordinatorExcept
         }
     }
 
-    /**
-     * Calls the provided function with the context; throws an exception otherwise.
-     * This method ensures that the context lock is acquired before calling the
-     * function and releases afterwards.
-     *
-     * @param tp    The topic partition.
-     * @param func  The function that will receive the context.
-     * @throws NotCoordinatorException
-     */
-    private void withContextOrThrow(
-        TopicPartition tp,
-        Consumer<CoordinatorContext> func
-    ) throws NotCoordinatorException {
-        CoordinatorContext context = contextOrThrow(tp);
-
-        try {
-            context.lock.lock();
-            func.accept(context);
-        } finally {
-            context.lock.unlock();
-        }
-    }
-
     /**
      * Calls the provided function with the context iff the context is active; throws
      * an exception otherwise. This method ensures that the context lock is acquired
@@ -1609,7 +1600,11 @@ public void scheduleLoadOperation(
         maybeCreateContext(tp);
 
         scheduleInternalOperation("Load(tp=" + tp + ", epoch=" + partitionEpoch + ")", tp, () -> {
-            withContextOrThrow(tp, context -> {
+            // The context is re-created if it does not exist.
+            CoordinatorContext context = maybeCreateContext(tp);
+
+            context.lock.lock();
+            try {
                 if (context.epoch < partitionEpoch) {
                     context.epoch = partitionEpoch;
 
@@ -1617,16 +1612,13 @@ public void scheduleLoadOperation(
                         case FAILED:
                         case INITIAL:
                             context.transitionTo(CoordinatorState.LOADING);
-                            loader.load(
-                                tp,
-                                context.coordinator
-                            ).whenComplete((summary, exception) -> {
+                            loader.load(tp, context.coordinator).whenComplete((summary, exception) -> {
                                 scheduleInternalOperation("CompleteLoad(tp=" + tp + ", epoch=" + partitionEpoch + ")", tp, () -> {
-                                    withContextOrThrow(tp, ctx -> {
+                                    CoordinatorContext ctx = coordinators.get(tp);
+                                    if (ctx != null)  {
                                         if (ctx.state != CoordinatorState.LOADING) {
-                                            log.info("Ignoring load completion from {} because context is in {} state.",
-                                                ctx.tp, ctx.state
-                                            );
+                                            log.info("Ignored load completion from {} because context is in {} state.",
+                                                ctx.tp, ctx.state);
                                             return;
                                         }
                                         try {
@@ -1635,18 +1627,19 @@ public void scheduleLoadOperation(
                                             if (summary != null) {
                                                 runtimeMetrics.recordPartitionLoadSensor(summary.startTimeMs(), summary.endTimeMs());
                                                 log.info("Finished loading of metadata from {} with epoch {} in {}ms where {}ms " +
-                                                    "was spent in the scheduler. Loaded {} records which total to {} bytes.",
+                                                         "was spent in the scheduler. Loaded {} records which total to {} bytes.",
                                                     tp, partitionEpoch, summary.endTimeMs() - summary.startTimeMs(),
-                                                    summary.schedulerQueueTimeMs(), summary.numRecords(), summary.numBytes()
-                                                );
+                                                    summary.schedulerQueueTimeMs(), summary.numRecords(), summary.numBytes());
                                             }
                                         } catch (Throwable ex) {
                                             log.error("Failed to load metadata from {} with epoch {} due to {}.",
-                                                tp, partitionEpoch, ex.toString()
-                                            );
+                                                tp, partitionEpoch, ex.toString());
                                             ctx.transitionTo(CoordinatorState.FAILED);
                                         }
-                                    });
+                                    } else {
+                                        log.debug("Failed to complete the loading of metadata for {} in epoch {} since the coordinator does not exist.",
+                                            tp, partitionEpoch);
+                                    }
                                 });
                             });
                             break;
@@ -1663,11 +1656,12 @@ public void scheduleLoadOperation(
                             log.error("Cannot load coordinator {} in state {}.", tp, context.state);
                     }
                 } else {
-                    log.info("Ignoring loading metadata from {} since current epoch {} is larger than or equals to {}.",
-                        context.tp, context.epoch, partitionEpoch
-                    );
+                    log.info("Ignored loading metadata from {} since current epoch {} is larger than or equals to {}.",
+                        context.tp, context.epoch, partitionEpoch);
                 }
-            });
+            } finally {
+                context.lock.unlock();
+            }
         });
     }
 
@@ -1689,8 +1683,8 @@ public void scheduleUnloadOperation(
         scheduleInternalOperation("UnloadCoordinator(tp=" + tp + ", epoch=" + partitionEpoch + ")", tp, () -> {
             CoordinatorContext context = coordinators.get(tp);
             if (context != null) {
+                context.lock.lock();
                 try {
-                    context.lock.lock();
                     if (!partitionEpoch.isPresent() || context.epoch < partitionEpoch.getAsInt()) {
                         log.info("Started unloading metadata for {} with epoch {}.", tp, partitionEpoch);
                         context.transitionTo(CoordinatorState.CLOSED);
@@ -1698,16 +1692,14 @@ public void scheduleUnloadOperation(
                         log.info("Finished unloading metadata for {} with epoch {}.", tp, partitionEpoch);
                     } else {
                         log.info("Ignored unloading metadata for {} in epoch {} since current epoch is {}.",
-                            tp, partitionEpoch, context.epoch
-                        );
+                            tp, partitionEpoch, context.epoch);
                     }
                 } finally {
                     context.lock.unlock();
                 }
             } else {
                 log.info("Ignored unloading metadata for {} in epoch {} since metadata was never loaded.",
-                    tp, partitionEpoch
-                );
+                    tp, partitionEpoch);
             }
         });
     }
@@ -1731,15 +1723,26 @@ public void onNewMetadataImage(
         // Push an event for each coordinator.
         coordinators.keySet().forEach(tp -> {
             scheduleInternalOperation("UpdateImage(tp=" + tp + ", offset=" + newImage.offset() + ")", tp, () -> {
-                withContextOrThrow(tp, context -> {
-                    if (context.state == CoordinatorState.ACTIVE) {
-                        log.debug("Applying new metadata image with offset {} to {}.", newImage.offset(), tp);
-                        context.coordinator.onNewMetadataImage(newImage, delta);
-                    } else {
-                        log.debug("Ignoring new metadata image with offset {} for {} because the coordinator is not active.",
-                            newImage.offset(), tp);
+                CoordinatorContext context = coordinators.get(tp);
+                if (context != null) {
+                    context.lock.lock();
+                    try {
+                        if (context.state == CoordinatorState.ACTIVE) {
+                            // The new image can be applied to the coordinator only if the coordinator
+                            // exists and is in the active state.
+                            log.debug("Applying new metadata image with offset {} to {}.", newImage.offset(), tp);
+                            context.coordinator.onNewMetadataImage(newImage, delta);
+                        } else {
+                            log.debug("Ignored new metadata image with offset {} for {} because the coordinator is not active.",
+                                newImage.offset(), tp);
+                        }
+                    } finally {
+                        context.lock.unlock();
                     }
-                });
+                } else {
+                    log.debug("Ignored new metadata image with offset {} for {} because the coordinator does not exist.",
+                        newImage.offset(), tp);
+                }
             });
         });
     }
@@ -1764,7 +1767,12 @@ public void close() throws Exception {
         Utils.closeQuietly(processor, "event processor");
         // Unload all the coordinators.
         coordinators.forEach((tp, context) -> {
-            context.transitionTo(CoordinatorState.CLOSED);
+            context.lock.lock();
+            try {
+                context.transitionTo(CoordinatorState.CLOSED);
+            } finally {
+                context.lock.unlock();
+            }
         });
         coordinators.clear();
         Utils.closeQuietly(runtimeMetrics, "runtime metrics");

From c8843f06841d7f3c94b640ec9dbf69ec4682ec11 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Fri, 1 Mar 2024 11:14:13 +0100
Subject: [PATCH 094/258] KAFKA-16167: Disable wakeups during autocommit on
 close (#15445)

When the consumer is closed, we perform a sychronous autocommit. We don't want to be woken up here, because we are already executing a close operation under a deadline. This is in line with the behavior of the old consumer.

This fixes PlaintextConsumerTest.testAutoCommitOnCloseAfterWakeup which is flaky on trunk - because we return immediately from the synchronous commit with a WakeupException, which causes us to not wait for the commit to finish and thereby sometimes miss the committed offset when a new consumer is created.

Reviewers: Matthias J. Sax <matthias@confluent.io>, Bruno Cadonna <cadonna@apache.org>
---
 .../internals/AsyncKafkaConsumer.java         |  8 ++--
 .../consumer/internals/WakeupTrigger.java     | 11 ++++++
 .../internals/AsyncKafkaConsumerTest.java     | 31 ++++++++++++++-
 .../consumer/internals/WakeupTriggerTest.java | 39 +++++++++++++++++++
 4 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index d810c5f053bb2..e706898b70a40 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -1228,6 +1228,9 @@ private void close(Duration timeout, boolean swallowException) {
         log.trace("Closing the Kafka consumer");
         AtomicReference<Throwable> firstException = new AtomicReference<>();
 
+        // We are already closing with a timeout, don't allow wake-ups from here on.
+        wakeupTrigger.disableWakeups();
+
         final Timer closeTimer = time.timer(timeout);
         clientTelemetryReporter.ifPresent(reporter -> reporter.initiateClose(timeout.toMillis()));
         closeTimer.update();
@@ -1265,7 +1268,7 @@ private void close(Duration timeout, boolean swallowException) {
     void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstException) {
         if (!groupMetadata.isPresent())
             return;
-        maybeAutoCommitSync(autoCommitEnabled, timer, firstException);
+        maybeAutoCommitSync(autoCommitEnabled, timer);
         applicationEventHandler.add(new CommitOnCloseEvent());
         completeQuietly(
             () -> {
@@ -1277,8 +1280,7 @@ void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstEx
 
     // Visible for testing
     void maybeAutoCommitSync(final boolean shouldAutoCommit,
-                             final Timer timer,
-                             final AtomicReference<Throwable> firstException) {
+                             final Timer timer) {
         if (!shouldAutoCommit)
             return;
         Map<TopicPartition, OffsetAndMetadata> allConsumed = subscriptions.allConsumed();
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/WakeupTrigger.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/WakeupTrigger.java
index 5a030f6307db9..209d5e41beef0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/WakeupTrigger.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/WakeupTrigger.java
@@ -73,6 +73,8 @@ public <T> CompletableFuture<T> setActiveTask(final CompletableFuture<T> current
             } else if (task instanceof WakeupFuture) {
                 currentTask.completeExceptionally(new WakeupException());
                 return null;
+            } else if (task instanceof DisabledWakeups) {
+                return task;
             }
             // last active state is still active
             throw new KafkaException("Last active task is still active");
@@ -88,6 +90,8 @@ public void setFetchAction(final FetchBuffer fetchBuffer) {
             } else if (task instanceof WakeupFuture) {
                 throwWakeupException.set(true);
                 return null;
+            } else if (task instanceof DisabledWakeups) {
+                return task;
             }
             // last active state is still active
             throw new IllegalStateException("Last active task is still active");
@@ -97,6 +101,10 @@ public void setFetchAction(final FetchBuffer fetchBuffer) {
         }
     }
 
+    public void disableWakeups() {
+        pendingTask.set(new DisabledWakeups());
+    }
+
     public void clearTask() {
         pendingTask.getAndUpdate(task -> {
             if (task == null) {
@@ -131,6 +139,9 @@ Wakeupable getPendingTask() {
 
     interface Wakeupable { }
 
+    // Set to block wakeups from happening and pending actions to be registered.
+    static class DisabledWakeups implements Wakeupable { }
+
     static class ActiveFuture implements Wakeupable {
         private final CompletableFuture<?> future;
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 2d6899612f056..35e742fbd0af8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -30,6 +30,7 @@
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.clients.consumer.RetriableCommitFailedException;
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
+import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeEvent;
 import org.apache.kafka.clients.consumer.internals.events.AsyncCommitEvent;
@@ -120,6 +121,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -676,7 +678,7 @@ public void testAutoCommitSyncEnabled() {
         consumer.subscribe(singleton("topic"), mock(ConsumerRebalanceListener.class));
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
-        consumer.maybeAutoCommitSync(true, time.timer(100), null);
+        consumer.maybeAutoCommitSync(true, time.timer(100));
         verify(applicationEventHandler).add(any(SyncCommitEvent.class));
     }
 
@@ -694,7 +696,7 @@ public void testAutoCommitSyncDisabled() {
         consumer.subscribe(singleton("topic"), mock(ConsumerRebalanceListener.class));
         subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
         subscriptions.seek(new TopicPartition("topic", 0), 100);
-        consumer.maybeAutoCommitSync(false, time.timer(100), null);
+        consumer.maybeAutoCommitSync(false, time.timer(100));
         verify(applicationEventHandler, never()).add(any(SyncCommitEvent.class));
     }
 
@@ -892,6 +894,31 @@ public void testWakeupCommitted() {
         assertNull(consumer.wakeupTrigger().getPendingTask());
     }
 
+    @Test
+    public void testNoWakeupInCloseCommit() {
+        TopicPartition tp = new TopicPartition("topic1", 0);
+        consumer = newConsumer();
+        consumer.assign(Collections.singleton(tp));
+        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
+        consumer.seek(tp, 10);
+        consumer.wakeup();
+
+        AtomicReference<SyncCommitEvent> capturedEvent = new AtomicReference<>();
+        doAnswer(invocation -> {
+            ApplicationEvent event = invocation.getArgument(0);
+            if (event instanceof SyncCommitEvent) {
+                capturedEvent.set((SyncCommitEvent) event);
+            }
+            return null;
+        }).when(applicationEventHandler).add(any());
+
+        consumer.close(Duration.ZERO);
+
+        // A commit was triggered and not completed exceptionally by the wakeup
+        assertNotNull(capturedEvent.get());
+        assertFalse(capturedEvent.get().future().isCompletedExceptionally());
+    }
+
     @Test
     public void testInterceptorAutoCommitOnClose() {
         Properties props = requiredConsumerPropertiesAndGroupId("test-id");
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
index 235ec78168d5e..3e1518814e7b8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
@@ -28,12 +28,14 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 @MockitoSettings(strictness = Strictness.STRICT_STUBS)
@@ -133,6 +135,43 @@ public void testManualTriggerWhenWakeupCalledAndFetchActionSet() {
         }
     }
 
+    @Test
+    public void testDisableWakeupWithoutPendingTask() {
+        wakeupTrigger.disableWakeups();
+        wakeupTrigger.wakeup();
+        assertDoesNotThrow(() -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
+    @Test
+    public void testDisableWakeupWithPendingTask() {
+        final CompletableFuture<Void> future = new CompletableFuture<>();
+        wakeupTrigger.disableWakeups();
+        wakeupTrigger.setActiveTask(future);
+        wakeupTrigger.wakeup();
+        assertFalse(future.isCompletedExceptionally());
+        assertDoesNotThrow(() -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
+    @Test
+    public void testDisableWakeupWithFetchAction() {
+        try (final FetchBuffer fetchBuffer = mock(FetchBuffer.class)) {
+            wakeupTrigger.disableWakeups();
+            wakeupTrigger.setFetchAction(fetchBuffer);
+            wakeupTrigger.wakeup();
+            verify(fetchBuffer, never()).wakeup();
+        }
+    }
+
+    @Test
+    public void testDisableWakeupPreservedByClearTask() {
+        final CompletableFuture<Void> future = new CompletableFuture<>();
+        wakeupTrigger.disableWakeups();
+        wakeupTrigger.setActiveTask(future);
+        wakeupTrigger.clearTask();
+        wakeupTrigger.wakeup();
+        assertDoesNotThrow(() -> wakeupTrigger.maybeTriggerWakeup());
+    }
+
     private void assertWakeupExceptionIsThrown(final CompletableFuture<?> future) {
         assertTrue(future.isCompletedExceptionally());
         try {

From 8e1516f88b88b2c815ccbca074af5a5c2d14b5c9 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Fri, 1 Mar 2024 16:42:08 +0100
Subject: [PATCH 095/258] KAFKA-16261: updateSubscription fails if already
 empty subscription (#15440)

The internal SubscriptionState object keeps track of whether the assignment is user-assigned, or auto-assigned. If there are no assigned partitions, the assignment resets to NONE. If you call SubscriptionState.assignFromSubscribed in this state it fails.

This change makes sure to check SubscriptionState.hasAutoAssignedPartitions() so that assignFromSubscribed is going to be permitted.

Also, a minor refactoring to make clearing the subscription a bit easier to follow in MembershipManagerImpl.

Testing via new unit test.

Reviewers: Bruno Cadonna <cadonna@apache.org>, Andrew Schofield <aschofield@confluent.io>
---
 .../internals/MembershipManagerImpl.java      | 37 ++++++-------------
 .../internals/MembershipManagerImplTest.java  | 20 ++++++++++
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index e74c7e30a2cad..35322fb51be33 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -550,7 +550,7 @@ public void transitionToFenced() {
                 log.error("onPartitionsLost callback invocation failed while releasing assignment" +
                         " after member got fenced. Member will rejoin the group anyways.", error);
             }
-            updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+            clearSubscription();
             if (state == MemberState.FENCED) {
                 transitionToJoining();
             } else {
@@ -583,7 +583,7 @@ public void transitionToFatal() {
                 log.error("onPartitionsLost callback invocation failed while releasing assignment" +
                         "after member failed with fatal error.", error);
             }
-            updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+            clearSubscription();
         });
     }
 
@@ -597,16 +597,14 @@ public void onSubscriptionUpdated() {
     }
 
     /**
-     * Update a new assignment by setting the assigned partitions in the member subscription.
-     *
-     * @param assignedPartitions Topic partitions to take as the new subscription assignment
-     * @param clearAssignments   True if the pending assignments and metadata cache should be cleared
+     * Clear the assigned partitions in the member subscription, pending assignments and metadata cache.
      */
-    private void updateSubscription(SortedSet<TopicIdPartition> assignedPartitions,
-                                    boolean clearAssignments) {
-        Collection<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedPartitions);
-        subscriptions.assignFromSubscribed(assignedTopicPartitions);
-        updateAssignmentLocally(assignedPartitions, clearAssignments);
+    private void clearSubscription() {
+        if (subscriptions.hasAutoAssignedPartitions()) {
+            subscriptions.assignFromSubscribed(Collections.emptySet());
+        }
+        updateCurrentAssignment(Collections.emptySet());
+        clearPendingAssignmentsAndLocalNamesCache();
     }
 
     /**
@@ -621,18 +619,7 @@ private void updateSubscriptionAwaitingCallback(SortedSet<TopicIdPartition> assi
                                                     SortedSet<TopicPartition> addedPartitions) {
         Collection<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedPartitions);
         subscriptions.assignFromSubscribedAwaitingCallback(assignedTopicPartitions, addedPartitions);
-        updateAssignmentLocally(assignedPartitions, false);
-    }
-
-    /**
-     * Make assignment effective on the group manager.
-     */
-    private void updateAssignmentLocally(SortedSet<TopicIdPartition> assignedPartitions,
-                                         boolean clearAssignments) {
         updateCurrentAssignment(assignedPartitions);
-        if (clearAssignments) {
-            clearPendingAssignmentsAndLocalNamesCache();
-        }
     }
 
     /**
@@ -660,7 +647,7 @@ public void transitionToJoining() {
     public CompletableFuture<Void> leaveGroup() {
         if (isNotInGroup()) {
             if (state == MemberState.FENCED) {
-                updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+                clearSubscription();
                 transitionTo(MemberState.UNSUBSCRIBED);
             }
             return CompletableFuture.completedFuture(null);
@@ -679,7 +666,7 @@ public CompletableFuture<Void> leaveGroup() {
         CompletableFuture<Void> callbackResult = invokeOnPartitionsRevokedOrLostToReleaseAssignment();
         callbackResult.whenComplete((result, error) -> {
             // Clear the subscription, no matter if the callback execution failed or succeeded.
-            updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+            clearSubscription();
 
             // Transition to ensure that a heartbeat request is sent out to effectively leave the
             // group (even in the case where the member had no assignment to release or when the
@@ -880,7 +867,7 @@ private void transitionToStale() {
                 log.error("onPartitionsLost callback invocation failed while releasing assignment" +
                     " after member left group due to expired poll timer.", error);
             }
-            updateSubscription(new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR), true);
+            clearSubscription();
             log.debug("Member {} sent leave group heartbeat and released its assignment. It will remain " +
                 "in {} state until the poll timer is reset, and it will then rejoin the group",
                 memberId, MemberState.STALE);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 49c5819be7b79..bc7be88c3bd8d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -771,6 +771,21 @@ public void testLeaveGroupWhenMemberOwnsAssignment() {
         testLeaveGroupReleasesAssignmentAndResetsEpochToSendLeaveGroup(membershipManager);
     }
 
+    @Test
+    public void testFencedWhenAssignmentEmpty() {
+        MembershipManager membershipManager = createMemberInStableState();
+
+        // Clear the assignment
+        when(subscriptionState.assignedPartitions()).thenReturn(Collections.emptySet());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(false);
+
+        membershipManager.transitionToFenced();
+
+        // Make sure to never call `assignFromSubscribed` again
+        verify(subscriptionState, never()).assignFromSubscribed(Collections.emptySet());
+    }
+
+
     @Test
     public void testLeaveGroupWhenMemberAlreadyLeaving() {
         MembershipManager membershipManager = createMemberInStableState();
@@ -1678,6 +1693,7 @@ private void assertLeaveGroupDueToExpiredPollAndTransitionToStale(MembershipMana
     public void testTransitionToLeavingWhileReconcilingDueToStaleMember() {
         MembershipManagerImpl membershipManager = memberJoinWithAssignment();
         clearInvocations(subscriptionState);
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
@@ -1686,6 +1702,7 @@ public void testTransitionToLeavingWhileReconcilingDueToStaleMember() {
     public void testTransitionToLeavingWhileJoiningDueToStaleMember() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         doNothing().when(subscriptionState).assignFromSubscribed(any());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         assertEquals(MemberState.JOINING, membershipManager.state());
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
@@ -1695,6 +1712,7 @@ public void testTransitionToLeavingWhileStableDueToStaleMember() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
@@ -1705,6 +1723,7 @@ public void testTransitionToLeavingWhileAcknowledgingDueToStaleMember() {
         MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true);
         doNothing().when(subscriptionState).assignFromSubscribed(any());
         clearInvocations(subscriptionState);
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
@@ -1775,6 +1794,7 @@ public void testStaleMemberWaitsForCallbackToRejoinWhenTimerReset() {
     private MembershipManagerImpl mockStaleMember() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         doNothing().when(subscriptionState).assignFromSubscribed(any());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         membershipManager.transitionToSendingLeaveGroup(true);
         membershipManager.onHeartbeatRequestSent();
         return membershipManager;

From 52a3fa07446f9c108399d47dbfb1685989a5d6eb Mon Sep 17 00:00:00 2001
From: Jamie <holmes.jc@gmail.com>
Date: Sat, 2 Mar 2024 08:13:56 +1300
Subject: [PATCH 096/258] KAFKA-15878: KIP-768 - Extend support for opaque
 (i.e. non-JWT) tokens in SASL/OAUTHBEARER (#14818)

# Overview
* This change pertains to [SASL/OAUTHBEARER ](https://kafka.apache.org/documentation/#security_sasl_oauthbearer)  mechanism of Kafka authentication.
* Kafka clients can use [SASL/OAUTHBEARER ](https://kafka.apache.org/documentation/#security_sasl_oauthbearer)   mechanism by overriding the [custom call back handlers](https://kafka.apache.org/documentation/#security_sasl_oauthbearer_prod) .
* [KIP-768](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=186877575) available from v3.1  further extends the mechanism with a production grade implementation.
* Kafka's [SASL/OAUTHBEARER ](https://kafka.apache.org/documentation/#security_sasl_oauthbearer)  mechanism currently **rejects the non-JWT (i.e. opaque) tokens**. This is because of a more restrictive set of characters than what [RFC-6750](https://datatracker.ietf.org/doc/html/rfc6750#section-2.1) recommends.
* This JIRA can be considered an extension of [KIP-768](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=186877575) to support the opaque tokens as well apart from the JWT tokens.

# Solution
* Have updated the regex in the the offending class to be compliant with the [RFC-6750](https://datatracker.ietf.org/doc/html/rfc6750#section-2.1)
* Have provided a supporting test case that includes the possible character set defined in [RFC-6750](https://datatracker.ietf.org/doc/html/rfc6750#section-2.1)

---------

Co-authored-by: Anuj Sharma <philomath.anuj@gmail.com>
Co-authored-by: Jamie Holmes <jamie.holmes@tesco.com>
Co-authored-by: Christopher Webb <31657038+cwebbtw@users.noreply.github.com>
Reviewers: Manikumar Reddy <manikumar.reddy@gmail.com>, Kirk True <ktrue@confluent.io>
---
 .../internals/OAuthBearerClientInitialResponse.java  |  2 +-
 .../OAuthBearerClientInitialResponseTest.java        | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java
index 73bfcd15c1286..3b340131cf875 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java
@@ -34,7 +34,7 @@ public class OAuthBearerClientInitialResponse {
     private static final String VALUE = "[\\x21-\\x7E \t\r\n]+";
 
     private static final String KVPAIRS = String.format("(%s=%s%s)*", KEY, VALUE, SEPARATOR);
-    private static final Pattern AUTH_PATTERN = Pattern.compile("(?<scheme>[\\w]+)[ ]+(?<token>[-_\\.a-zA-Z0-9]+)");
+    private static final Pattern AUTH_PATTERN = Pattern.compile("(?<scheme>[\\w]+)[ ]+(?<token>[-_~+/\\.a-zA-Z0-9]+([=]*))");
     private static final Pattern CLIENT_INITIAL_RESPONSE_PATTERN = Pattern.compile(
             String.format("n,(a=(?<authzid>%s))?,%s(?<kvpairs>%s)%s", SASLNAME, SEPARATOR, KVPAIRS, SEPARATOR));
     public static final String AUTH_KEY = "auth";
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java
index 3b3c90bf1d21d..fc44297a2f32d 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java
@@ -102,6 +102,18 @@ public void testRfc7688Example() throws Exception {
         assertEquals("143", response.extensions().map().get("port"));
     }
 
+    // RFC 6750 token format  1*( ALPHA / DIGIT /"-" / "." / "_" / "~" / "+" / "/" ) *"="
+    @Test
+    public void testCharSupportForRfc6750Token() throws Exception {
+        String message = "n,a=user@example.com,\u0001host=server.example.com\u0001port=143\u0001" +
+                "auth=Bearer vF-9.df_t4qm~Tc2Nvb3RlckBhbHR+hdmlzdGEuY29/tCg==\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("vF-9.df_t4qm~Tc2Nvb3RlckBhbHR+hdmlzdGEuY29/tCg==", response.tokenValue());
+        assertEquals("user@example.com", response.authorizationId());
+        assertEquals("server.example.com", response.extensions().map().get("host"));
+        assertEquals("143", response.extensions().map().get("port"));
+    }
+
     @Test
     public void testNoExtensionsFromByteArray() throws Exception {
         String message = "n,a=user@example.com,\u0001" +

From 21a5bbd84c3c80bcee57ceb05f47005d9122d46e Mon Sep 17 00:00:00 2001
From: Said Boudjelda <bmscomp@gmail.com>
Date: Sun, 3 Mar 2024 17:32:28 +0100
Subject: [PATCH 097/258] MINOR: Upgrade jqwik to version 1.8.3 (#14365)

This minor pull request consist of upgrading version of jqwik library to version 1.8.0 that brings some bug fixing and some enhancements, upgrading the version now will make future upgrades easier

For breaking changes:

We are not using ArbitraryConfiguratorBase, so there is no overriding of configure method
We are not using TypeUsage.canBeAssignedTo(TypeUsage)
No breaking is related to @Provide and @ForAll usage no Exception CannotFindArbitraryException is thrown during tests running
No usage of StringArbitrary.repeatChars(0.0)
We are not affected by the removal of method TypeArbitrary.use(Executable)
We are not affected by the removal or methods ActionChainArbitrary.addAction(action)  and ActionChainArbitrary.addAction(weight, action)
For more details check the release notes: https://jqwik.net/release-notes.html#180

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>, Yash Mayya <yash.mayya@gmail.com>
---
 gradle/dependencies.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 9cd1515a79c0c..5ac978f9b00d4 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -119,7 +119,7 @@ versions += [
   jopt: "5.0.4",
   jose4j: "0.9.4",
   junit: "5.10.2",
-  jqwik: "1.7.4",
+  jqwik: "1.8.3",
   kafka_0100: "0.10.0.1",
   kafka_0101: "0.10.1.1",
   kafka_0102: "0.10.2.2",

From 44af72fd77376cd5fb0c0a6019b0e5908928f0e1 Mon Sep 17 00:00:00 2001
From: PoAn Yang <yangpoan@gmail.com>
Date: Mon, 4 Mar 2024 00:36:57 +0800
Subject: [PATCH 098/258] MINOR: fix link for
 ListTransactionsOptions#filterOnDuration (#15459)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 clients/src/main/java/org/apache/kafka/clients/admin/Admin.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
index ff0e60e766d51..e9f0a58e41b98 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
@@ -1634,7 +1634,7 @@ default ListTransactionsResult listTransactions() {
      * should typically attempt to reduce the size of the result set using
      * {@link ListTransactionsOptions#filterProducerIds(Collection)} or
      * {@link ListTransactionsOptions#filterStates(Collection)} or
-     * {@link ListTransactionsOptions#durationFilter(Long)}
+     * {@link ListTransactionsOptions#filterOnDuration(long)}.
      *
      * @param options Options to control the method behavior (including filters)
      * @return The result

From 907e945c0b70263d312038d448196fff1f49a98b Mon Sep 17 00:00:00 2001
From: Ayoub Omari <ayoubomari1@outlook.fr>
Date: Sun, 3 Mar 2024 18:03:04 +0100
Subject: [PATCH 099/258] MINOR: fix SessionStore java doc (#15412)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../streams/state/ReadOnlyKeyValueStore.java    |  2 +-
 .../streams/state/ReadOnlySessionStore.java     | 17 ++++++++---------
 .../streams/state/ReadOnlyWindowStore.java      |  4 ++--
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyKeyValueStore.java
index f905a2962565d..4ef656c2123f0 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyKeyValueStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyKeyValueStore.java
@@ -93,7 +93,7 @@ default KeyValueIterator<K, V> reverseRange(K from, K to) {
      * and must not return null values.
      * Order is not guaranteed as bytes lexicographical ordering might not represent key order.
      *
-     * @return An reverse iterator of all key/value pairs in the store, from largest to smallest key bytes.
+     * @return A reverse iterator of all key/value pairs in the store, from largest to smallest key bytes.
      * @throws InvalidStateStoreException if the store is not initialized
      */
     default KeyValueIterator<K, V> reverseAll() {
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlySessionStore.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlySessionStore.java
index 7fe11a6bea0e7..5a52f00f6063f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlySessionStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlySessionStore.java
@@ -37,7 +37,7 @@ public interface ReadOnlySessionStore<K, AGG> {
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -48,7 +48,6 @@ public interface ReadOnlySessionStore<K, AGG> {
      * <p>
      * This iterator must be closed after use.
      *
-     *
      * @param key                    the key to return sessions for
      * @param earliestSessionEndTime the end timestamp of the earliest session to search for, where
      *                               iteration starts.
@@ -72,7 +71,7 @@ default KeyValueIterator<Windowed<K>, AGG> findSessions(final K key,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -106,7 +105,7 @@ default KeyValueIterator<Windowed<K>, AGG> findSessions(final K key,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -140,7 +139,7 @@ default KeyValueIterator<Windowed<K>, AGG> backwardFindSessions(final K key,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -175,7 +174,7 @@ default KeyValueIterator<Windowed<K>, AGG> backwardFindSessions(final K key,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -213,7 +212,7 @@ default KeyValueIterator<Windowed<K>, AGG> findSessions(final K keyFrom,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -251,7 +250,7 @@ default KeyValueIterator<Windowed<K>, AGG> findSessions(final K keyFrom,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
@@ -289,7 +288,7 @@ default KeyValueIterator<Windowed<K>, AGG> backwardFindSessions(final K keyFrom,
      * is the upper bound of the search interval, and the method returns all sessions that overlap
      * with the search interval.
      * Thus, if a session ends before earliestSessionEndTime, or starts after latestSessionStartTime
-     * if won't be contained in the result:
+     * it won't be contained in the result:
      * <pre>{@code
      * earliestSessionEndTime: ESET
      * latestSessionStartTime: LSST
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyWindowStore.java b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyWindowStore.java
index 3df170d5ab620..30ca72a64cc20 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyWindowStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/ReadOnlyWindowStore.java
@@ -174,7 +174,7 @@ default KeyValueIterator<Windowed<K>, V> backwardFetch(K keyFrom, K keyTo, Insta
      * Gets all the key-value pairs in the existing windows in backward order
      * with respect to time (from end to beginning of time).
      *
-     * @return an backward iterator over windowed key-value pairs {@code <Windowed<K>, value>}, from the end to beginning of time.
+     * @return a backward iterator over windowed key-value pairs {@code <Windowed<K>, value>}, from the end to beginning of time.
      * @throws InvalidStateStoreException if the store is not initialized
      */
     default KeyValueIterator<Windowed<K>, V> backwardAll() {
@@ -199,7 +199,7 @@ default KeyValueIterator<Windowed<K>, V> backwardAll() {
      *
      * @param timeFrom the beginning of the time slot from which to search (inclusive), where iteration ends.
      * @param timeTo   the end of the time slot from which to search (inclusive), where iteration starts.
-     * @return an backward iterator over windowed key-value pairs {@code <Windowed<K>, value>}, from end to beginning of time.
+     * @return a backward iterator over windowed key-value pairs {@code <Windowed<K>, value>}, from end to beginning of time.
      * @throws InvalidStateStoreException if the store is not initialized
      * @throws NullPointerException       if {@code null} is used for any key
      * @throws IllegalArgumentException   if duration is negative or can't be represented as {@code long milliseconds}

From 2c0cab39aedab3a8635510acfac2551aaeb62ffb Mon Sep 17 00:00:00 2001
From: Ismael Juma <ismael@juma.me.uk>
Date: Sun, 3 Mar 2024 23:31:57 -0800
Subject: [PATCH 100/258] MINOR: Remove unnecessary easymock/powermock
 dependencies (#15460)

These projects don't actually use easymock/powermock.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/build.gradle b/build.gradle
index b362c2f673a5b..f2ba22e86b39d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2183,8 +2183,6 @@ project(':streams') {
     testImplementation libs.junitJupiter
     testImplementation libs.junitVintageEngine
     testImplementation libs.easymock
-    testImplementation libs.powermockJunit4
-    testImplementation libs.powermockEasymock
     testImplementation libs.bcpkix
     testImplementation libs.hamcrest
     testImplementation libs.mockitoCore
@@ -2332,7 +2330,6 @@ project(':streams:streams-scala') {
     testImplementation project(':streams:test-utils')
 
     testImplementation libs.junitJupiter
-    testImplementation libs.easymock
     testImplementation libs.mockitoCore
     testImplementation libs.mockitoJunitJupiter // supports MockitoExtension
     testImplementation libs.hamcrest
@@ -2869,7 +2866,6 @@ project(':connect:transforms') {
 
     implementation libs.slf4jApi
 
-    testImplementation libs.easymock
     testImplementation libs.junitJupiter
 
     testRuntimeOnly libs.slf4jlog4j
@@ -2909,8 +2905,7 @@ project(':connect:json') {
     api libs.jacksonAfterburner
 
     implementation libs.slf4jApi
-
-    testImplementation libs.easymock
+    
     testImplementation libs.junitJupiter
 
     testRuntimeOnly libs.slf4jlog4j

From 7dbdc15c668dfb5a4a91c79f339c22fb7178c368 Mon Sep 17 00:00:00 2001
From: Ayoub Omari <ayoubomari1@outlook.fr>
Date: Mon, 4 Mar 2024 10:19:59 +0100
Subject: [PATCH 101/258] KAFKA-15625: Do not flush global state store at each
 commit (#15361)

Global state stores are currently flushed at each commit, which may impact performance, especially for EOS (commit each 200ms).
The goal of this improvement is to flush global state stores only when the delta between the current offset and the last checkpointed offset exceeds a threshold.
This is the same logic we apply on local state store, with a threshold of 10000 records.
The implementation only flushes if the time interval elapsed and the threshold of 10000 records is exceeded.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Bruno Cadonna <cadonna@apache.org>
---
 .../internals/GlobalStateMaintainer.java      |   2 +
 .../internals/GlobalStateUpdateTask.java      |  20 +++-
 .../internals/GlobalStreamThread.java         |  25 +---
 .../internals/GlobalStateTaskTest.java        | 111 ++++++++++++++++--
 .../internals/StateConsumerTest.java          |  31 ++---
 .../kafka/test/GlobalStateManagerStub.java    |  10 +-
 .../kafka/streams/TopologyTestDriver.java     |   4 +-
 7 files changed, 147 insertions(+), 56 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateMaintainer.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateMaintainer.java
index 9a8aab6eb3c84..06afb6fde4f2a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateMaintainer.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateMaintainer.java
@@ -34,4 +34,6 @@ interface GlobalStateMaintainer {
     void close(final boolean wipeStateStore) throws IOException;
 
     void update(ConsumerRecord<byte[], byte[]> record);
+
+    void maybeCheckpoint();
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateUpdateTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateUpdateTask.java
index 523228542a8be..da7ebba209a89 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateUpdateTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateUpdateTask.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.streams.errors.DeserializationExceptionHandler;
 import org.apache.kafka.streams.errors.StreamsException;
@@ -45,18 +46,26 @@ public class GlobalStateUpdateTask implements GlobalStateMaintainer {
     private final Map<String, RecordDeserializer> deserializers = new HashMap<>();
     private final GlobalStateManager stateMgr;
     private final DeserializationExceptionHandler deserializationExceptionHandler;
+    private final Time time;
+    private final long flushInterval;
+    private long lastFlush;
 
     public GlobalStateUpdateTask(final LogContext logContext,
                                  final ProcessorTopology topology,
                                  final InternalProcessorContext processorContext,
                                  final GlobalStateManager stateMgr,
-                                 final DeserializationExceptionHandler deserializationExceptionHandler) {
+                                 final DeserializationExceptionHandler deserializationExceptionHandler,
+                                 final Time time,
+                                 final long flushInterval
+                                 ) {
         this.logContext = logContext;
         this.log = logContext.logger(getClass());
         this.topology = topology;
         this.stateMgr = stateMgr;
         this.processorContext = processorContext;
         this.deserializationExceptionHandler = deserializationExceptionHandler;
+        this.time = time;
+        this.flushInterval = flushInterval;
     }
 
     /**
@@ -86,6 +95,7 @@ public Map<TopicPartition, Long> initialize() {
         }
         initTopology();
         processorContext.initialize();
+        lastFlush = time.milliseconds();
         return stateMgr.changelogOffsets();
     }
 
@@ -150,5 +160,13 @@ private void initTopology() {
         }
     }
 
+    @Override
+    public void maybeCheckpoint() {
+        final long now = time.milliseconds();
+        if (now - flushInterval >= lastFlush && StateManagerUtil.checkpointNeeded(false, stateMgr.changelogOffsets(), offsets)) {
+            flushState();
+            lastFlush = now;
+        }
+    }
 
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStreamThread.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStreamThread.java
index 82a0cc51131b2..1ed517b15d493 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStreamThread.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStreamThread.java
@@ -228,25 +228,17 @@ public GlobalStreamThread(final ProcessorTopology topology,
     static class StateConsumer {
         private final Consumer<byte[], byte[]> globalConsumer;
         private final GlobalStateMaintainer stateMaintainer;
-        private final Time time;
         private final Duration pollTime;
-        private final long flushInterval;
         private final Logger log;
 
-        private long lastFlush;
-
         StateConsumer(final LogContext logContext,
                       final Consumer<byte[], byte[]> globalConsumer,
                       final GlobalStateMaintainer stateMaintainer,
-                      final Time time,
-                      final Duration pollTime,
-                      final long flushInterval) {
+                      final Duration pollTime) {
             this.log = logContext.logger(getClass());
             this.globalConsumer = globalConsumer;
             this.stateMaintainer = stateMaintainer;
-            this.time = time;
             this.pollTime = pollTime;
-            this.flushInterval = flushInterval;
         }
 
         /**
@@ -259,7 +251,6 @@ void initialize() {
             for (final Map.Entry<TopicPartition, Long> entry : partitionOffsets.entrySet()) {
                 globalConsumer.seek(entry.getKey(), entry.getValue());
             }
-            lastFlush = time.milliseconds();
         }
 
         void pollAndUpdate() {
@@ -267,11 +258,7 @@ void pollAndUpdate() {
             for (final ConsumerRecord<byte[], byte[]> record : received) {
                 stateMaintainer.update(record);
             }
-            final long now = time.milliseconds();
-            if (now - flushInterval >= lastFlush) {
-                stateMaintainer.flushState();
-                lastFlush = now;
-            }
+            stateMaintainer.maybeCheckpoint();
         }
 
         public void close(final boolean wipeStateStore) throws IOException {
@@ -418,11 +405,11 @@ private StateConsumer initialize() {
                     topology,
                     globalProcessorContext,
                     stateMgr,
-                    config.defaultDeserializationExceptionHandler()
+                    config.defaultDeserializationExceptionHandler(),
+                    time,
+                    config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG)
                 ),
-                time,
-                Duration.ofMillis(config.getLong(StreamsConfig.POLL_MS_CONFIG)),
-                config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG)
+                Duration.ofMillis(config.getLong(StreamsConfig.POLL_MS_CONFIG))
             );
 
             try {
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateTaskTest.java
index 31be9dc2a4db1..af5dc68103c80 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateTaskTest.java
@@ -25,6 +25,7 @@
 import org.apache.kafka.common.serialization.LongSerializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.streams.errors.LogAndContinueExceptionHandler;
 import org.apache.kafka.streams.errors.LogAndFailExceptionHandler;
@@ -46,8 +47,6 @@
 
 import static java.util.Arrays.asList;
 import static org.apache.kafka.streams.processor.internals.testutil.ConsumerRecordUtil.record;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -71,8 +70,12 @@ public class GlobalStateTaskTest {
     private final MockProcessorNode<?, ?, ?, ?> processorTwo = new MockProcessorNode<>();
 
     private final Map<TopicPartition, Long> offsets = new HashMap<>();
-    private File testDirectory = TestUtils.tempDirectory("global-store");
+    private final File testDirectory = TestUtils.tempDirectory("global-store");
     private final NoOpProcessorContext context = new NoOpProcessorContext();
+    private final MockTime time = new MockTime();
+    private final long flushInterval = 1000L;
+    private final long currentOffsetT1 = 50;
+    private final long currentOffsetT2 = 100;
 
     private ProcessorTopology topology;
     private GlobalStateManagerStub stateMgr;
@@ -101,7 +104,9 @@ public void before() {
             topology,
             context,
             stateMgr,
-            new LogAndFailExceptionHandler()
+            new LogAndFailExceptionHandler(),
+            time,
+            flushInterval
         );
     }
 
@@ -188,7 +193,9 @@ public void shouldNotThrowStreamsExceptionWhenKeyDeserializationFailsWithSkipHan
             topology,
             context,
             stateMgr,
-            new LogAndContinueExceptionHandler()
+            new LogAndContinueExceptionHandler(),
+            time,
+            flushInterval
         );
         final byte[] key = new LongSerializer().serialize(topic2, 1L);
         final byte[] recordValue = new IntegerSerializer().serialize(topic2, 10);
@@ -203,7 +210,9 @@ public void shouldNotThrowStreamsExceptionWhenValueDeserializationFails() {
             topology,
             context,
             stateMgr,
-            new LogAndContinueExceptionHandler()
+            new LogAndContinueExceptionHandler(),
+            time,
+            flushInterval
         );
         final byte[] key = new IntegerSerializer().serialize(topic2, 1);
         final byte[] recordValue = new LongSerializer().serialize(topic2, 10L);
@@ -217,10 +226,13 @@ public void shouldFlushStateManagerWithOffsets() {
         final Map<TopicPartition, Long> expectedOffsets = new HashMap<>();
         expectedOffsets.put(t1, 52L);
         expectedOffsets.put(t2, 100L);
+
         globalStateTask.initialize();
-        globalStateTask.update(record(topic1, 1, 51, "foo".getBytes(), "foo".getBytes()));
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 1, "foo".getBytes(), "foo".getBytes()));
         globalStateTask.flushState();
+
         assertEquals(expectedOffsets, stateMgr.changelogOffsets());
+        assertTrue(stateMgr.flushed);
     }
 
     @Test
@@ -228,12 +240,93 @@ public void shouldCheckpointOffsetsWhenStateIsFlushed() {
         final Map<TopicPartition, Long> expectedOffsets = new HashMap<>();
         expectedOffsets.put(t1, 102L);
         expectedOffsets.put(t2, 100L);
+
         globalStateTask.initialize();
-        globalStateTask.update(record(topic1, 1, 101, "foo".getBytes(), "foo".getBytes()));
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 51L, "foo".getBytes(), "foo".getBytes()));
         globalStateTask.flushState();
-        assertThat(stateMgr.changelogOffsets(), equalTo(expectedOffsets));
+
+        assertEquals(expectedOffsets, stateMgr.changelogOffsets());
+        assertTrue(stateMgr.checkpointWritten);
+    }
+
+    @Test
+    public void shouldNotCheckpointIfNotReceivedEnoughRecords() {
+        globalStateTask.initialize();
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 9000L, "foo".getBytes(), "foo".getBytes()));
+        time.sleep(flushInterval); // flush interval elapsed
+        globalStateTask.maybeCheckpoint();
+
+        assertEquals(offsets, stateMgr.changelogOffsets());
+        assertFalse(stateMgr.flushed);
+        assertFalse(stateMgr.checkpointWritten);
+    }
+
+    @Test
+    public void shouldNotCheckpointWhenFlushIntervalHasNotLapsed() {
+        globalStateTask.initialize();
+
+        // offset delta exceeded
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 10000L, "foo".getBytes(), "foo".getBytes()));
+
+        time.sleep(flushInterval / 2);
+        globalStateTask.maybeCheckpoint();
+
+        assertEquals(offsets, stateMgr.changelogOffsets());
+        assertFalse(stateMgr.flushed);
+        assertFalse(stateMgr.checkpointWritten);
+    }
+
+    @Test
+    public void shouldCheckpointIfReceivedEnoughRecordsAndFlushIntervalHasElapsed() {
+        final Map<TopicPartition, Long> expectedOffsets = new HashMap<>();
+        expectedOffsets.put(t1, 10051L); // topic1 advanced with 10001 records
+        expectedOffsets.put(t2, 100L);
+
+        globalStateTask.initialize();
+
+        time.sleep(flushInterval); // flush interval elapsed
+
+        // 10000 records received since last flush => do not flush
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 9999L, "foo".getBytes(), "foo".getBytes()));
+        globalStateTask.maybeCheckpoint();
+
+        assertEquals(offsets, stateMgr.changelogOffsets());
+        assertFalse(stateMgr.flushed);
+        assertFalse(stateMgr.checkpointWritten);
+
+        // 1 more record received => triggers the flush
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 10000L, "foo".getBytes(), "foo".getBytes()));
+        globalStateTask.maybeCheckpoint();
+
+        assertEquals(expectedOffsets, stateMgr.changelogOffsets());
+        assertTrue(stateMgr.flushed);
+        assertTrue(stateMgr.checkpointWritten);
     }
 
+    @Test
+    public void shouldCheckpointIfReceivedEnoughRecordsFromMultipleTopicsAndFlushIntervalElapsed() {
+        final byte[] integerBytes = new IntegerSerializer().serialize(topic2, 1);
+
+        final Map<TopicPartition, Long> expectedOffsets = new HashMap<>();
+        expectedOffsets.put(t1, 9050L); // topic1 advanced with 9000 records
+        expectedOffsets.put(t2, 1101L); // topic2 advanced with 1001 records
+
+        globalStateTask.initialize();
+
+        time.sleep(flushInterval);
+
+        // received 9000 records in topic1
+        globalStateTask.update(record(topic1, 1, currentOffsetT1 + 8999L, "foo".getBytes(), "foo".getBytes()));
+        // received 1001 records in topic2
+        globalStateTask.update(record(topic2, 1, currentOffsetT2 + 1000L, integerBytes, integerBytes));
+        globalStateTask.maybeCheckpoint();
+
+        assertEquals(expectedOffsets, stateMgr.changelogOffsets());
+        assertTrue(stateMgr.flushed);
+        assertTrue(stateMgr.checkpointWritten);
+    }
+
+
     @Test
     public void shouldWipeGlobalStateDirectory() throws Exception {
         assertTrue(stateMgr.baseDir().exists());
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StateConsumerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StateConsumerTest.java
index 1f98eb456d917..5e57939483370 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StateConsumerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StateConsumerTest.java
@@ -21,7 +21,6 @@
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.LogContext;
-import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Utils;
 import org.junit.Before;
 import org.junit.Test;
@@ -32,16 +31,13 @@
 import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 
 public class StateConsumerTest {
 
-    private static final long FLUSH_INTERVAL = 1000L;
     private final TopicPartition topicOne = new TopicPartition("topic-one", 1);
     private final TopicPartition topicTwo = new TopicPartition("topic-two", 1);
-    private final MockTime time = new MockTime();
     private final MockConsumer<byte[], byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
     private final Map<TopicPartition, Long> partitionOffsets = new HashMap<>();
     private final LogContext logContext = new LogContext("test ");
@@ -53,7 +49,7 @@ public void setUp() {
         partitionOffsets.put(topicOne, 20L);
         partitionOffsets.put(topicTwo, 30L);
         stateMaintainer = new TaskStub(partitionOffsets);
-        stateConsumer = new GlobalStreamThread.StateConsumer(logContext, consumer, stateMaintainer, time, Duration.ofMillis(10L), FLUSH_INTERVAL);
+        stateConsumer = new GlobalStreamThread.StateConsumer(logContext, consumer, stateMaintainer, Duration.ofMillis(10L));
     }
 
     @Test
@@ -76,6 +72,7 @@ public void shouldUpdateStateWithReceivedRecordsForPartition() {
         consumer.addRecord(new ConsumerRecord<>("topic-one", 1, 21L, new byte[0], new byte[0]));
         stateConsumer.pollAndUpdate();
         assertEquals(2, stateMaintainer.updatedPartitions.get(topicOne).intValue());
+        assertTrue(stateMaintainer.flushed);
     }
 
     @Test
@@ -87,27 +84,9 @@ public void shouldUpdateStateWithReceivedRecordsForAllTopicPartition() {
         stateConsumer.pollAndUpdate();
         assertEquals(1, stateMaintainer.updatedPartitions.get(topicOne).intValue());
         assertEquals(2, stateMaintainer.updatedPartitions.get(topicTwo).intValue());
-    }
-
-    @Test
-    public void shouldFlushStoreWhenFlushIntervalHasLapsed() {
-        stateConsumer.initialize();
-        consumer.addRecord(new ConsumerRecord<>("topic-one", 1, 20L, new byte[0], new byte[0]));
-        time.sleep(FLUSH_INTERVAL);
-
-        stateConsumer.pollAndUpdate();
         assertTrue(stateMaintainer.flushed);
     }
 
-    @Test
-    public void shouldNotFlushOffsetsWhenFlushIntervalHasNotLapsed() {
-        stateConsumer.initialize();
-        consumer.addRecord(new ConsumerRecord<>("topic-one", 1, 20L, new byte[0], new byte[0]));
-        time.sleep(FLUSH_INTERVAL / 2);
-        stateConsumer.pollAndUpdate();
-        assertFalse(stateMaintainer.flushed);
-    }
-
     @Test
     public void shouldCloseConsumer() throws IOException {
         stateConsumer.close(false);
@@ -161,6 +140,10 @@ public void update(final ConsumerRecord<byte[], byte[]> record) {
             updatedPartitions.put(tp, updatedPartitions.get(tp) + 1);
         }
 
+        @Override
+        public void maybeCheckpoint() {
+            flushState();
+        }
     }
 
-}
\ No newline at end of file
+}
diff --git a/streams/src/test/java/org/apache/kafka/test/GlobalStateManagerStub.java b/streams/src/test/java/org/apache/kafka/test/GlobalStateManagerStub.java
index d34b3c8029f56..3031649944770 100644
--- a/streams/src/test/java/org/apache/kafka/test/GlobalStateManagerStub.java
+++ b/streams/src/test/java/org/apache/kafka/test/GlobalStateManagerStub.java
@@ -35,6 +35,8 @@ public class GlobalStateManagerStub implements GlobalStateManager {
     private final File baseDirectory;
     public boolean initialized;
     public boolean closed;
+    public boolean flushed;
+    public boolean checkpointWritten;
 
     public GlobalStateManagerStub(final Set<String> storeNames,
                                   final Map<TopicPartition, Long> offsets,
@@ -64,7 +66,9 @@ public void registerStore(final StateStore store,
                               final CommitCallback checkpoint) {}
 
     @Override
-    public void flush() {}
+    public void flush() {
+        flushed = true;
+    }
 
     @Override
     public void close() {
@@ -77,7 +81,9 @@ public void updateChangelogOffsets(final Map<TopicPartition, Long> writtenOffset
     }
 
     @Override
-    public void checkpoint() {}
+    public void checkpoint() {
+        checkpointWritten = true;
+    }
 
     @Override
     public StateStore getStore(final String name) {
diff --git a/streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java b/streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java
index 20abaa5407294..5767ed9d20e70 100644
--- a/streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java
+++ b/streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java
@@ -459,7 +459,9 @@ private void setupGlobalTask(final Time mockWallClockTime,
                 globalTopology,
                 globalProcessorContext,
                 globalStateManager,
-                new LogAndContinueExceptionHandler()
+                new LogAndContinueExceptionHandler(),
+                mockWallClockTime,
+                streamsConfig.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG)
             );
             globalStateTask.initialize();
             globalProcessorContext.setRecordContext(null);

From 21a888c4ca0189924a82e4e4ee4909b214c4c9bc Mon Sep 17 00:00:00 2001
From: "Gyeongwon, Do" <dct012@gmail.com>
Date: Mon, 4 Mar 2024 20:36:46 +0900
Subject: [PATCH 102/258] MINOR: Updating comments to match the code (#15388)

This comment was added by #12862

The method with the comment was originally named updateLastSend, but its name was later changed to onSendAttempt.
This method doesn't increment numAttempts.

It seems that the numAttempts is only modified after a Request succeeds or fails.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../apache/kafka/clients/consumer/internals/RequestState.java   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestState.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestState.java
index 321be5e8fe808..a888e7831abeb 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestState.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestState.java
@@ -102,7 +102,7 @@ public boolean requestInFlight() {
     }
 
     public void onSendAttempt(final long currentTimeMs) {
-        // Here we update the timer everytime we try to send a request. Also increment number of attempts.
+        // Here we update the timer everytime we try to send a request.
         this.lastSentMs = currentTimeMs;
     }
 

From aa0443eb607c2d1d3004312f55f7583102127cb8 Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <cadonna@apache.org>
Date: Mon, 4 Mar 2024 12:42:24 +0100
Subject: [PATCH 103/258] KAFKA-16285: Make group metadata available when a new
 assignment is set (#15426)

Currently, in the async Kafka consumer updates to the group metadata
that are received by the heartbeat are propagated to the application thread
in form of an event. Group metadata is updated when a new assignment is
received. The new assignment is directly set in the subscription without
sending an update event from the background thread to the application thread.
That means that there might be a delay between the application thread being
aware of the update to the assignment and the application thread being
aware of the update to the group metadata. This delay can cause stale
group metadata returned by the application thread that then causes
issues when data of the new assignment is committed. A concrete
example is
producer.sendOffsetsToTransaction(offsetsToCommit, groupMetadata)
The offsets to commit might already stem from the new assignment
but the group metadata might relate to the previous assignment.

Reviewers: Kirk True <ktrue@confluent.io>, Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/AsyncKafkaConsumer.java         | 118 ++++----
 .../internals/HeartbeatRequestManager.java    |  18 --
 .../internals/MembershipManagerImpl.java      |   6 +
 .../consumer/internals/RequestManagers.java   |   4 +-
 .../internals/events/BackgroundEvent.java     |   2 +-
 .../events/GroupMetadataUpdateEvent.java      |  53 ----
 .../internals/AsyncKafkaConsumerTest.java     | 271 +++++-------------
 .../HeartbeatRequestManagerTest.java          |  87 ------
 .../internals/RequestManagersTest.java        |  70 +++++
 .../java/org/apache/kafka/test/TestUtils.java |   9 +
 10 files changed, 227 insertions(+), 411 deletions(-)
 delete mode 100644 clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
 create mode 100644 clients/src/test/java/org/apache/kafka/clients/consumer/internals/RequestManagersTest.java

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index e706898b70a40..dafd3293c74c4 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -52,7 +52,6 @@
 import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsEvent;
-import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
 import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseEvent;
 import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
@@ -213,10 +212,6 @@ public void process(final BackgroundEvent event) {
                     process((ErrorEvent) event);
                     break;
 
-                case GROUP_METADATA_UPDATE:
-                    process((GroupMetadataUpdateEvent) event);
-                    break;
-
                 case CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED:
                     process((ConsumerRebalanceListenerCallbackNeededEvent) event);
                     break;
@@ -231,18 +226,6 @@ private void process(final ErrorEvent event) {
             throw event.error();
         }
 
-        private void process(final GroupMetadataUpdateEvent event) {
-            if (AsyncKafkaConsumer.this.groupMetadata.isPresent()) {
-                final ConsumerGroupMetadata currentGroupMetadata = AsyncKafkaConsumer.this.groupMetadata.get();
-                AsyncKafkaConsumer.this.groupMetadata = Optional.of(new ConsumerGroupMetadata(
-                    currentGroupMetadata.groupId(),
-                    event.memberEpoch(),
-                    event.memberId(),
-                    currentGroupMetadata.groupInstanceId()
-                ));
-            }
-        }
-
         private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
             ApplicationEvent invokedEvent = invokeRebalanceCallbacks(
                 rebalanceListenerInvoker,
@@ -256,7 +239,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
 
     private final ApplicationEventHandler applicationEventHandler;
     private final Time time;
-    private Optional<ConsumerGroupMetadata> groupMetadata = Optional.empty();
+    private final AtomicReference<Optional<ConsumerGroupMetadata>> groupMetadata = new AtomicReference<>(Optional.empty());
     private final KafkaConsumerMetrics kafkaConsumerMetrics;
     private Logger log;
     private final String clientId;
@@ -370,6 +353,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
                     fetchMetricsManager,
                     clientTelemetryReporter.map(ClientTelemetryReporter::telemetrySender).orElse(null));
             this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
+            this.groupMetadata.set(initializeGroupMetadata(config, groupRebalanceConfig));
             final Supplier<RequestManagers> requestManagersSupplier = RequestManagers.supplier(time,
                     logContext,
                     backgroundEventHandler,
@@ -383,7 +367,9 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
                     networkClientDelegateSupplier,
                     clientTelemetryReporter,
                     metrics,
-                    offsetCommitCallbackInvoker);
+                    offsetCommitCallbackInvoker,
+                    this::updateGroupMetadata
+            );
             final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(logContext,
                     metadata,
                     applicationEventQueue,
@@ -413,8 +399,6 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
                     config.originals(Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId))
             );
 
-            this.groupMetadata = initializeGroupMetadata(config, groupRebalanceConfig);
-
             // The FetchCollector is only used on the application thread.
             this.fetchCollector = fetchCollectorFactory.build(logContext,
                     metadata,
@@ -426,7 +410,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
 
             this.kafkaConsumerMetrics = new KafkaConsumerMetrics(metrics, CONSUMER_METRIC_GROUP_PREFIX);
 
-            if (groupMetadata.isPresent() &&
+            if (groupMetadata.get().isPresent() &&
                 GroupProtocol.of(config.getString(ConsumerConfig.GROUP_PROTOCOL_CONFIG)) == GroupProtocol.CONSUMER) {
                 config.ignore(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG); // Used by background thread
             }
@@ -478,7 +462,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
                 rebalanceListenerInvoker
         );
         this.metrics = metrics;
-        this.groupMetadata = initializeGroupMetadata(groupId, Optional.empty());
+        this.groupMetadata.set(initializeGroupMetadata(groupId, Optional.empty()));
         this.metadata = metadata;
         this.retryBackoffMs = retryBackoffMs;
         this.defaultApiTimeoutMs = defaultApiTimeoutMs;
@@ -532,7 +516,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
             GroupRebalanceConfig.ProtocolType.CONSUMER
         );
 
-        this.groupMetadata = initializeGroupMetadata(config, groupRebalanceConfig);
+        this.groupMetadata.set(initializeGroupMetadata(config, groupRebalanceConfig));
 
         BlockingQueue<ApplicationEvent> applicationEventQueue = new LinkedBlockingQueue<>();
         BlockingQueue<BackgroundEvent> backgroundEventQueue = new LinkedBlockingQueue<>();
@@ -568,7 +552,8 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
             networkClientDelegateSupplier,
             clientTelemetryReporter,
             metrics,
-            offsetCommitCallbackInvoker
+            offsetCommitCallbackInvoker,
+            this::updateGroupMetadata
         );
         Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(
                 logContext,
@@ -651,17 +636,35 @@ private Optional<ConsumerGroupMetadata> initializeGroupMetadata(final String gro
                 throw new InvalidGroupIdException("The configured " + ConsumerConfig.GROUP_ID_CONFIG
                     + " should not be an empty string or whitespace.");
             } else {
-                return Optional.of(new ConsumerGroupMetadata(
-                    groupId,
-                    JoinGroupRequest.UNKNOWN_GENERATION_ID,
-                    JoinGroupRequest.UNKNOWN_MEMBER_ID,
-                    groupInstanceId
-                ));
+                return Optional.of(initializeConsumerGroupMetadata(groupId, groupInstanceId));
             }
         }
         return Optional.empty();
     }
 
+    private ConsumerGroupMetadata initializeConsumerGroupMetadata(final String groupId,
+                                                                  final Optional<String> groupInstanceId) {
+        return new ConsumerGroupMetadata(
+            groupId,
+            JoinGroupRequest.UNKNOWN_GENERATION_ID,
+            JoinGroupRequest.UNKNOWN_MEMBER_ID,
+            groupInstanceId
+        );
+    }
+
+    private void updateGroupMetadata(final Optional<Integer> memberEpoch, final Optional<String> memberId) {
+        groupMetadata.updateAndGet(
+            oldGroupMetadataOptional -> oldGroupMetadataOptional.map(
+                oldGroupMetadata -> new ConsumerGroupMetadata(
+                    oldGroupMetadata.groupId(),
+                    memberEpoch.orElse(oldGroupMetadata.generationId()),
+                    memberId.orElse(oldGroupMetadata.memberId()),
+                    oldGroupMetadata.groupInstanceId()
+                )
+            )
+        );
+    }
+
     /**
      * poll implementation using {@link ApplicationEventHandler}.
      *  1. Poll for background events. If there's a fetch response event, process the record and return it. If it is
@@ -713,18 +716,14 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
                 wakeupTrigger.maybeTriggerWakeup();
 
                 updateAssignmentMetadataIfNeeded(timer);
-                if (isGenerationKnownOrPartitionsUserAssigned()) {
-                    final Fetch<K, V> fetch = pollForFetches(timer);
-                    if (!fetch.isEmpty()) {
-                        if (fetch.records().isEmpty()) {
-                            log.trace("Returning empty records from `poll()` "
-                                + "since the consumer's position has advanced for at least one topic partition");
-                        }
-
-                        return interceptors.onConsume(new ConsumerRecords<>(fetch.records()));
+                final Fetch<K, V> fetch = pollForFetches(timer);
+                if (!fetch.isEmpty()) {
+                    if (fetch.records().isEmpty()) {
+                        log.trace("Returning empty records from `poll()` "
+                            + "since the consumer's position has advanced for at least one topic partition");
                     }
-                } else {
-                    timer.update();
+
+                    return interceptors.onConsume(new ConsumerRecords<>(fetch.records()));
                 }
                 // We will wait for retryBackoffMs
             } while (timer.notExpired());
@@ -736,13 +735,6 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
         }
     }
 
-    private boolean isGenerationKnownOrPartitionsUserAssigned() {
-        if (subscriptions.hasAutoAssignedPartitions()) {
-            return groupMetadata.filter(g -> g.generationId() != JoinGroupRequest.UNKNOWN_GENERATION_ID).isPresent();
-        }
-        return true;
-    }
-
     /**
      * Commit offsets returned on the last {@link #poll(Duration) poll()} for all the subscribed list of topics and
      * partitions.
@@ -960,7 +952,7 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
     }
 
     private void maybeThrowInvalidGroupIdException() {
-        if (!groupMetadata.isPresent()) {
+        if (!groupMetadata.get().isPresent()) {
             throw new InvalidGroupIdException("To use the group management or offset commit APIs, you must " +
                 "provide a valid " + ConsumerConfig.GROUP_ID_CONFIG + " in the consumer configuration.");
         }
@@ -1186,7 +1178,7 @@ public ConsumerGroupMetadata groupMetadata() {
         acquireAndEnsureOpen();
         try {
             maybeThrowInvalidGroupIdException();
-            return groupMetadata.get();
+            return groupMetadata.get().get();
         } finally {
             release();
         }
@@ -1266,7 +1258,7 @@ private void close(Duration timeout, boolean swallowException) {
      * 3. if partition revocation completes successfully, send leave group
      */
     void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstException) {
-        if (!groupMetadata.isPresent())
+        if (!groupMetadata.get().isPresent())
             return;
         maybeAutoCommitSync(autoCommitEnabled, timer);
         applicationEventHandler.add(new CommitOnCloseEvent());
@@ -1463,7 +1455,7 @@ public void unsubscribe() {
         acquireAndEnsureOpen();
         try {
             fetchBuffer.retainAll(Collections.emptySet());
-            if (groupMetadata.isPresent()) {
+            if (groupMetadata.get().isPresent()) {
                 UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent();
                 applicationEventHandler.add(unsubscribeEvent);
                 log.info("Unsubscribing all topics or patterns and assigned partitions");
@@ -1475,7 +1467,7 @@ public void unsubscribe() {
                 } catch (TimeoutException e) {
                     log.error("Failed while waiting for the unsubscribe event to complete");
                 }
-                groupMetadata = initializeGroupMetadata(groupMetadata.get().groupId(), Optional.empty());
+                resetGroupMetadata();
             }
             subscriptions.unsubscribe();
         } finally {
@@ -1483,6 +1475,16 @@ public void unsubscribe() {
         }
     }
 
+    private void resetGroupMetadata() {
+        groupMetadata.updateAndGet(
+            oldGroupMetadataOptional -> oldGroupMetadataOptional
+                .map(oldGroupMetadata -> initializeConsumerGroupMetadata(
+                    oldGroupMetadata.groupId(),
+                    oldGroupMetadata.groupInstanceId()
+                ))
+        );
+    }
+
     @Override
     @Deprecated
     public ConsumerRecords<K, V> poll(final long timeoutMs) {
@@ -1604,7 +1606,7 @@ private boolean updateFetchPositions(final Timer timer) {
      * according to config {@link CommonClientConfigs#GROUP_ID_CONFIG}
      */
     private boolean isCommittedOffsetsManagementEnabled() {
-        return groupMetadata.isPresent();
+        return groupMetadata.get().isPresent();
     }
 
     /**
@@ -1905,12 +1907,12 @@ public KafkaConsumerMetrics kafkaConsumerMetrics() {
     private void maybeThrowFencedInstanceException() {
         if (offsetCommitCallbackInvoker.hasFencedException()) {
             String groupInstanceId = "unknown";
-            if (!groupMetadata.isPresent()) {
+            if (!groupMetadata.get().isPresent()) {
                 log.error("No group metadata found although a group ID was provided. This is a bug!");
-            } else if (!groupMetadata.get().groupInstanceId().isPresent()) {
+            } else if (!groupMetadata.get().get().groupInstanceId().isPresent()) {
                 log.error("No group instance ID found although the consumer is fenced. This is a bug!");
             } else {
-                groupInstanceId = groupMetadata.get().groupInstanceId().get();
+                groupInstanceId = groupMetadata.get().get().groupInstanceId().get();
             }
             throw new FencedInstanceIdException("Get fenced exception for group.instance.id " + groupInstanceId);
         }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 826774a6a64e0..d2a7205d5f3ff 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -22,7 +22,6 @@
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
-import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.HeartbeatMetricsManager;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
@@ -110,7 +109,6 @@ public class HeartbeatRequestManager implements RequestManager {
      * sending heartbeat until the next poll.
      */
     private final Timer pollTimer;
-    private GroupMetadataUpdateEvent previousGroupMetadataUpdateEvent = null;
 
     /**
      * Holding the heartbeat sensor to measure heartbeat timing and response latency
@@ -328,27 +326,11 @@ private void onResponse(final ConsumerGroupHeartbeatResponse response, long curr
             heartbeatRequestState.onSuccessfulAttempt(currentTimeMs);
             heartbeatRequestState.resetTimer();
             membershipManager.onHeartbeatSuccess(response.data());
-            maybeSendGroupMetadataUpdateEvent();
             return;
         }
         onErrorResponse(response, currentTimeMs);
     }
 
-    private void maybeSendGroupMetadataUpdateEvent() {
-        if (previousGroupMetadataUpdateEvent == null ||
-            !previousGroupMetadataUpdateEvent.memberId().equals(membershipManager.memberId()) ||
-            previousGroupMetadataUpdateEvent.memberEpoch() != membershipManager.memberEpoch()) {
-
-            final GroupMetadataUpdateEvent currentGroupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-                membershipManager.memberEpoch(),
-                previousGroupMetadataUpdateEvent != null && membershipManager.memberId() == null ?
-                    previousGroupMetadataUpdateEvent.memberId() : membershipManager.memberId()
-            );
-            this.backgroundEventHandler.add(currentGroupMetadataUpdateEvent);
-            previousGroupMetadataUpdateEvent = currentGroupMetadataUpdateEvent;
-        }
-    }
-
     private void onErrorResponse(final ConsumerGroupHeartbeatResponse response,
                                  final long currentTimeMs) {
         Errors error = Errors.forCode(response.data().errorCode());
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 35322fb51be33..c2bdd3f860991 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -57,6 +57,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
+import static java.util.Collections.unmodifiableList;
 import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED;
 import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST;
 import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED;
@@ -1491,4 +1492,9 @@ public PollResult poll(final long currentTimeMs) {
         }
         return PollResult.EMPTY;
     }
+
+    // visible for testing
+    List<MemberStateListener> stateListeners() {
+        return unmodifiableList(stateUpdatesListeners);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
index 0b4c043d4a4d6..75d87432db680 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
@@ -123,7 +123,8 @@ public static Supplier<RequestManagers> supplier(final Time time,
                                                      final Supplier<NetworkClientDelegate> networkClientDelegateSupplier,
                                                      final Optional<ClientTelemetryReporter> clientTelemetryReporter,
                                                      final Metrics metrics,
-                                                     final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker
+                                                     final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker,
+                                                     final MemberStateListener applicationThreadMemberStateListener
                                                      ) {
         return new CachedSupplier<RequestManagers>() {
             @Override
@@ -192,6 +193,7 @@ protected RequestManagers create() {
                             time,
                             metrics);
                     membershipManager.registerStateListener(commit);
+                    membershipManager.registerStateListener(applicationThreadMemberStateListener);
                     heartbeatRequestManager = new HeartbeatRequestManager(
                             logContext,
                             time,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
index 9bc3fbebc30a6..4241482bcaa3f 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
@@ -27,7 +27,7 @@
 public abstract class BackgroundEvent {
 
     public enum Type {
-        ERROR, CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED, GROUP_METADATA_UPDATE
+        ERROR, CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED
     }
 
     private final Type type;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
deleted file mode 100644
index 001f549818383..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/GroupMetadataUpdateEvent.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer.internals.events;
-
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.internals.ConsumerNetworkThread;
-
-/**
- * This event is sent by the {@link ConsumerNetworkThread consumer's network thread} to the application thread
- * so that when the user calls the {@link Consumer#groupMetadata()} API, the information is up-to-date. The
- * information for the current state of the group member is managed on the consumer network thread and thus
- * requires this interplay between threads.
- */
-public class GroupMetadataUpdateEvent extends BackgroundEvent {
-
-    private final int memberEpoch;
-    private final String memberId;
-
-    public GroupMetadataUpdateEvent(final int memberEpoch, final String memberId) {
-        super(Type.GROUP_METADATA_UPDATE);
-        this.memberEpoch = memberEpoch;
-        this.memberId = memberId;
-    }
-
-    public int memberEpoch() {
-        return memberEpoch;
-    }
-
-    public String memberId() {
-        return memberId;
-    }
-
-    @Override
-    public String toStringBase() {
-        return super.toStringBase() +
-            ", memberEpoch=" + memberEpoch +
-            ", memberId='" + memberId + '\'';
-    }
-}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 35e742fbd0af8..fe83a417897a6 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -42,7 +42,6 @@
 import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsEvent;
-import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
 import org.apache.kafka.clients.consumer.internals.events.LeaveOnCloseEvent;
 import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.NewTopicsMetadataUpdateRequestEvent;
@@ -52,7 +51,6 @@
 import org.apache.kafka.clients.consumer.internals.events.SyncCommitEvent;
 import org.apache.kafka.clients.consumer.internals.events.UnsubscribeEvent;
 import org.apache.kafka.clients.consumer.internals.events.ValidatePositionsEvent;
-import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.Node;
@@ -80,6 +78,7 @@
 import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
+import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
 import java.time.Duration;
@@ -102,7 +101,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -117,10 +115,12 @@
 import static org.apache.kafka.clients.consumer.internals.MembershipManagerImpl.TOPIC_PARTITION_COMPARATOR;
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
+import static org.apache.kafka.test.TestUtils.requiredConsumerConfig;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -133,6 +133,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -164,19 +165,19 @@ public void resetAll() {
     }
 
     private AsyncKafkaConsumer<String, String> newConsumer() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.GROUP_ID_CONFIG, "group-id");
         return newConsumer(props);
     }
 
     private AsyncKafkaConsumer<String, String> newConsumerWithoutGroupId() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         return newConsumer(props);
     }
 
     @SuppressWarnings("UnusedReturnValue")
     private AsyncKafkaConsumer<String, String> newConsumerWithEmptyGroupId() {
-        final Properties props = requiredConsumerPropertiesAndGroupId("");
+        final Properties props = requiredConsumerConfigAndGroupId("");
         return newConsumer(props);
     }
 
@@ -921,7 +922,7 @@ public void testNoWakeupInCloseCommit() {
 
     @Test
     public void testInterceptorAutoCommitOnClose() {
-        Properties props = requiredConsumerPropertiesAndGroupId("test-id");
+        Properties props = requiredConsumerConfigAndGroupId("test-id");
         props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
         props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
 
@@ -937,7 +938,7 @@ public void testInterceptorAutoCommitOnClose() {
 
     @Test
     public void testInterceptorCommitSync() {
-        Properties props = requiredConsumerPropertiesAndGroupId("test-id");
+        Properties props = requiredConsumerConfigAndGroupId("test-id");
         props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
         props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
 
@@ -952,7 +953,7 @@ public void testInterceptorCommitSync() {
 
     @Test
     public void testNoInterceptorCommitSyncFailed() {
-        Properties props = requiredConsumerPropertiesAndGroupId("test-id");
+        Properties props = requiredConsumerConfigAndGroupId("test-id");
         props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
         props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
 
@@ -968,7 +969,7 @@ public void testNoInterceptorCommitSyncFailed() {
 
     @Test
     public void testInterceptorCommitAsync() {
-        Properties props = requiredConsumerPropertiesAndGroupId("test-id");
+        Properties props = requiredConsumerConfigAndGroupId("test-id");
         props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
         props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
 
@@ -985,7 +986,7 @@ public void testInterceptorCommitAsync() {
 
     @Test
     public void testNoInterceptorCommitAsyncFailed() {
-        Properties props = requiredConsumerPropertiesAndGroupId("test-id");
+        Properties props = requiredConsumerConfigAndGroupId("test-id");
         props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
         props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
 
@@ -1085,7 +1086,7 @@ public void testSubscriptionOnEmptyTopic() {
 
     @Test
     public void testGroupMetadataAfterCreationWithGroupIdIsNull() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         final ConsumerConfig config = new ConsumerConfig(props);
         consumer = newConsumer(config);
 
@@ -1102,7 +1103,7 @@ public void testGroupMetadataAfterCreationWithGroupIdIsNull() {
     @Test
     public void testGroupMetadataAfterCreationWithGroupIdIsNotNull() {
         final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerConfigAndGroupId(groupId));
         consumer = newConsumer(config);
 
         final ConsumerGroupMetadata groupMetadata = consumer.groupMetadata();
@@ -1117,7 +1118,7 @@ public void testGroupMetadataAfterCreationWithGroupIdIsNotNull() {
     public void testGroupMetadataAfterCreationWithGroupIdIsNotNullAndGroupInstanceIdSet() {
         final String groupId = "consumerGroupA";
         final String groupInstanceId = "groupInstanceId1";
-        final Properties props = requiredConsumerPropertiesAndGroupId(groupId);
+        final Properties props = requiredConsumerConfigAndGroupId(groupId);
         props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId);
         final ConsumerConfig config = new ConsumerConfig(props);
         consumer = newConsumer(config);
@@ -1130,164 +1131,65 @@ public void testGroupMetadataAfterCreationWithGroupIdIsNotNullAndGroupInstanceId
         assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
     }
 
-    @Test
-    public void testGroupMetadataUpdateSingleCall() {
+    private MemberStateListener captureGroupMetadataUpdateListener(final MockedStatic<RequestManagers> requestManagers) {
+        ArgumentCaptor<MemberStateListener> applicationThreadMemberStateListener = ArgumentCaptor.forClass(MemberStateListener.class);
+        requestManagers.verify(() -> RequestManagers.supplier(
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            any(),
+            applicationThreadMemberStateListener.capture()
+        ));
+        return applicationThreadMemberStateListener.getValue();
+    }
+
+    @Test
+    public void testGroupMetadataUpdate() {
         final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
-        consumer = newConsumer(config);
-
-        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
-        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
-
-        final int generation = 1;
-        final String memberId = "newMemberId";
-        final ConsumerGroupMetadata expectedGroupMetadata = new ConsumerGroupMetadata(
-            groupId,
-            generation,
-            memberId,
-            Optional.empty()
-        );
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            generation,
-            memberId
-        );
-        backgroundEventQueue.add(groupMetadataUpdateEvent);
-        consumer.assign(singletonList(new TopicPartition("topic", 0)));
-        consumer.poll(Duration.ZERO);
-
-        final ConsumerGroupMetadata actualGroupMetadata = consumer.groupMetadata();
-
-        assertEquals(expectedGroupMetadata, actualGroupMetadata);
-
-        final ConsumerGroupMetadata secondActualGroupMetadataWithoutUpdate = consumer.groupMetadata();
-
-        assertEquals(expectedGroupMetadata, secondActualGroupMetadataWithoutUpdate);
-    }
-
-    @Test
-    public void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsedWithTopics() {
-        testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(() -> {
-            consumer.subscribe(singletonList("topic"));
-        });
-    }
-
-    @Test
-    public void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsedWithPattern() {
-        testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(() -> {
-            when(metadata.fetch()).thenReturn(Cluster.empty());
-            consumer.subscribe(Pattern.compile("topic"));
-        });
-    }
-
-    private void testPollNotReturningRecordsIfGenerationUnknownAndGroupManagementIsUsed(final Runnable subscription) {
-        final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
-        consumer = newConsumer(config);
-        subscription.run();
-
-        consumer.poll(Duration.ZERO);
-
-        verify(fetchCollector, never()).collectFetch(any(FetchBuffer.class));
-    }
-
-    @Test
-    public void testPollReturningRecordsIfGroupIdSetAndGroupManagementIsNotUsed() {
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId("consumerGroupA"));
-        testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(config);
-    }
-
-    @Test
-    public void testPollReturningRecordsIfGroupIdNotSetAndGroupManagementIsNotUsed() {
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerProperties());
-        testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(config);
-    }
-
-    private void testPollReturningRecordsIfGroupMetadataHasUnknownGenerationAndGroupManagementIsNotUsed(final ConsumerConfig config) {
-        final String topic = "topic";
-        final TopicPartition topicPartition = new TopicPartition(topic, 0);
-        consumer = newConsumer(config);
-        consumer.assign(singletonList(topicPartition));
-        final List<ConsumerRecord<String, String>> records = singletonList(
-            new ConsumerRecord<>(topic, 0, 2, "key1", "value1")
-        );
-        when(fetchCollector.collectFetch(any(FetchBuffer.class)))
-            .thenReturn(Fetch.forPartition(topicPartition, records, true));
-        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
-
-        consumer.poll(Duration.ZERO);
-
-        verify(fetchCollector).collectFetch(any(FetchBuffer.class));
-    }
-
-    @Test
-    public void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsedWithTopics() {
-        final String topic = "topic";
-        testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(
-            topic,
-            () -> {
-                consumer.subscribe(singletonList(topic));
-            });
-    }
-
-    @Test
-    public void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsedWithPattern() {
-        final String topic = "topic";
-        testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(
-            topic,
-            () -> {
-                when(metadata.fetch()).thenReturn(Cluster.empty());
-                consumer.subscribe(Pattern.compile(topic));
-            });
-    }
-
-    private void testPollReturningRecordIfGenerationKnownAndGroupManagementIsUsed(final String topic,
-                                                                                  final Runnable subscription) {
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId("consumerGroupA"));
-        final int generation = 1;
-        final String memberId = "newMemberId";
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            generation,
-            memberId
-        );
-        backgroundEventQueue.add(groupMetadataUpdateEvent);
-        final TopicPartition topicPartition = new TopicPartition(topic, 0);
-        final List<ConsumerRecord<String, String>> records = singletonList(
-            new ConsumerRecord<>(topic, 0, 2, "key1", "value1")
-        );
-        when(fetchCollector.collectFetch(any(FetchBuffer.class)))
-            .thenReturn(Fetch.forPartition(topicPartition, records, true));
-        consumer = newConsumer(config);
-        subscription.run();
-
-        consumer.poll(Duration.ZERO);
-
-        verify(fetchCollector).collectFetch(any(FetchBuffer.class));
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerConfigAndGroupId(groupId));
+        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
+            consumer = newConsumer(config);
+            final ConsumerGroupMetadata oldGroupMetadata = consumer.groupMetadata();
+            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
+            final int expectedMemberEpoch = 42;
+            final String expectedMemberId = "memberId";
+            groupMetadataUpdateListener.onMemberEpochUpdated(
+                Optional.of(expectedMemberEpoch),
+                Optional.of(expectedMemberId)
+            );
+            final ConsumerGroupMetadata newGroupMetadata = consumer.groupMetadata();
+            assertEquals(oldGroupMetadata.groupId(), newGroupMetadata.groupId());
+            assertEquals(expectedMemberId, newGroupMetadata.memberId());
+            assertEquals(expectedMemberEpoch, newGroupMetadata.generationId());
+            assertEquals(oldGroupMetadata.groupInstanceId(), newGroupMetadata.groupInstanceId());
+        }
     }
 
     @Test
     public void testGroupMetadataIsResetAfterUnsubscribe() {
         final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
-        consumer = newConsumer(config);
-        consumer.subscribe(singletonList("topic"));
-        final int generation = 1;
-        final String memberId = "newMemberId";
-        final ConsumerGroupMetadata groupMetadataAfterSubscription = new ConsumerGroupMetadata(
-            groupId,
-            generation,
-            memberId,
-            Optional.empty()
-        );
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            generation,
-            memberId
-        );
-        backgroundEventQueue.add(groupMetadataUpdateEvent);
-        when(fetchCollector.collectFetch(any(FetchBuffer.class))).thenReturn(Fetch.empty());
-        consumer.poll(Duration.ZERO);
-
-        assertEquals(groupMetadataAfterSubscription, consumer.groupMetadata());
-
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerConfigAndGroupId(groupId));
+        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
+            consumer = newConsumer(config);
+            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
+            consumer.subscribe(singletonList("topic"));
+            final int memberEpoch = 42;
+            final String memberId = "memberId";
+            groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), Optional.of(memberId));
+            final ConsumerGroupMetadata groupMetadata = consumer.groupMetadata();
+            assertNotEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
+            assertNotEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
+        }
         completeUnsubscribeApplicationEventSuccessfully();
 
         consumer.unsubscribe();
@@ -1298,7 +1200,6 @@ public void testGroupMetadataIsResetAfterUnsubscribe() {
             JoinGroupRequest.UNKNOWN_MEMBER_ID,
             Optional.empty()
         );
-
         assertEquals(groupMetadataAfterUnsubscription, consumer.groupMetadata());
     }
 
@@ -1379,7 +1280,7 @@ private static Stream<Arguments> listenerCallbacksInvokeSource() {
     @Test
     public void testBackgroundError() {
         final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerConfigAndGroupId(groupId));
         consumer = newConsumer(config);
 
         final KafkaException expectedException = new KafkaException("Nobody expects the Spanish Inquisition");
@@ -1394,7 +1295,7 @@ public void testBackgroundError() {
     @Test
     public void testMultipleBackgroundErrors() {
         final String groupId = "consumerGroupA";
-        final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+        final ConsumerConfig config = new ConsumerConfig(requiredConsumerConfigAndGroupId(groupId));
         consumer = newConsumer(config);
 
         final KafkaException expectedException1 = new KafkaException("Nobody expects the Spanish Inquisition");
@@ -1412,7 +1313,7 @@ public void testMultipleBackgroundErrors() {
 
     @Test
     public void testGroupRemoteAssignorUnusedIfGroupIdUndefined() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
         final ConsumerConfig config = new ConsumerConfig(props);
         consumer = newConsumer(config);
@@ -1422,7 +1323,7 @@ public void testGroupRemoteAssignorUnusedIfGroupIdUndefined() {
 
     @Test
     public void testGroupRemoteAssignorUnusedInGenericProtocol() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.GROUP_ID_CONFIG, "consumerGroupA");
         props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT));
         props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
@@ -1434,7 +1335,7 @@ public void testGroupRemoteAssignorUnusedInGenericProtocol() {
 
     @Test
     public void testGroupRemoteAssignorUsedInConsumerProtocol() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.GROUP_ID_CONFIG, "consumerGroupA");
         props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT));
         props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
@@ -1446,7 +1347,7 @@ public void testGroupRemoteAssignorUsedInConsumerProtocol() {
 
     @Test
     public void testGroupIdNull() {
-        final Properties props = requiredConsumerProperties();
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 10000);
         props.put(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, true);
         final ConsumerConfig config = new ConsumerConfig(props);
@@ -1458,7 +1359,7 @@ public void testGroupIdNull() {
 
     @Test
     public void testGroupIdNotNullAndValid() {
-        final Properties props = requiredConsumerPropertiesAndGroupId("consumerGroupA");
+        final Properties props = requiredConsumerConfigAndGroupId("consumerGroupA");
         props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 10000);
         props.put(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, true);
         final ConsumerConfig config = new ConsumerConfig(props);
@@ -1492,7 +1393,6 @@ public void testEnsurePollEventSentOnConsumerPoll() {
         final TopicPartition tp = new TopicPartition("topic", 0);
         final List<ConsumerRecord<String, String>> records = singletonList(
                 new ConsumerRecord<>("topic", 0, 2, "key1", "value1"));
-        backgroundEventQueue.add(new GroupMetadataUpdateEvent(1, "memberId"));
         doAnswer(invocation -> Fetch.forPartition(tp, records, true))
                 .when(fetchCollector)
                 .collectFetch(Mockito.any(FetchBuffer.class));
@@ -1503,7 +1403,7 @@ public void testEnsurePollEventSentOnConsumerPoll() {
     }
 
     private void testInvalidGroupId(final String groupId) {
-        final Properties props = requiredConsumerPropertiesAndGroupId(groupId);
+        final Properties props = requiredConsumerConfigAndGroupId(groupId);
         final ConsumerConfig config = new ConsumerConfig(props);
 
         final Exception exception = assertThrows(
@@ -1514,20 +1414,12 @@ private void testInvalidGroupId(final String groupId) {
         assertEquals("Failed to construct kafka consumer", exception.getMessage());
     }
 
-    private Properties requiredConsumerPropertiesAndGroupId(final String groupId) {
-        final Properties props = requiredConsumerProperties();
+    private Properties requiredConsumerConfigAndGroupId(final String groupId) {
+        final Properties props = requiredConsumerConfig();
         props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
         return props;
     }
 
-    private Properties requiredConsumerProperties() {
-        final Properties props = new Properties();
-        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
-        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
-        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9091");
-        return props;
-    }
-
     private void testUpdateFetchPositionsWithFetchCommittedOffsetsTimeout(boolean committedOffsetsEnabled) {
         completeFetchedCommittedOffsetApplicationEventExceptionally(new TimeoutException());
         doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
@@ -1588,13 +1480,6 @@ public void testLongPollWaitIsLimited() {
             new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
         );
 
-        final int generation = 1;
-        final String memberId = "newMemberId";
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = new GroupMetadataUpdateEvent(
-            generation,
-            memberId
-        );
-        backgroundEventQueue.add(groupMetadataUpdateEvent);
         // On the first iteration, return no data; on the second, return two records
         doAnswer(invocation -> {
             // Mock the subscription being assigned as the first fetch is collected
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 90cfb90bb1bda..8e05e505be471 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -21,7 +21,6 @@
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.Uuid;
@@ -325,91 +324,6 @@ public void testValidateConsumerGroupHeartbeatRequest(final short version) {
         assertEquals(DEFAULT_REMOTE_ASSIGNOR, heartbeatRequest.data().serverAssignor());
     }
 
-    @Test
-    public void testConsumerGroupMetadataFirstUpdate() {
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = makeFirstGroupMetadataUpdate(memberId, memberEpoch);
-        assertEquals(memberEpoch, groupMetadataUpdateEvent.memberEpoch());
-        assertEquals(memberId, groupMetadataUpdateEvent.memberId());
-    }
-
-    @Test
-    public void testConsumerGroupMetadataUpdateWithSameUpdate() {
-        makeFirstGroupMetadataUpdate(memberId, memberEpoch);
-
-        time.sleep(2000);
-        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
-
-        assertEquals(1, result.unsentRequests.size());
-        NetworkClientDelegate.UnsentRequest request = result.unsentRequests.get(0);
-        ClientResponse responseWithSameUpdate = createHeartbeatResponse(request, Errors.NONE);
-        request.handler().onComplete(responseWithSameUpdate);
-        assertEquals(0, backgroundEventQueue.size());
-    }
-
-    @Test
-    public void testConsumerGroupMetadataUpdateWithMemberIdNullButMemberEpochUpdated() {
-        makeFirstGroupMetadataUpdate(memberId, memberEpoch);
-
-        time.sleep(2000);
-        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
-
-        assertEquals(1, result.unsentRequests.size());
-        NetworkClientDelegate.UnsentRequest request = result.unsentRequests.get(0);
-        final int updatedMemberEpoch = 2;
-        ClientResponse responseWithMemberEpochUpdate = createHeartbeatResponseWithMemberIdNull(
-            request,
-            Errors.NONE,
-            updatedMemberEpoch
-        );
-        request.handler().onComplete(responseWithMemberEpochUpdate);
-        assertEquals(1, backgroundEventQueue.size());
-        final BackgroundEvent eventWithUpdatedMemberEpoch = backgroundEventQueue.poll();
-        assertEquals(BackgroundEvent.Type.GROUP_METADATA_UPDATE, eventWithUpdatedMemberEpoch.type());
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = (GroupMetadataUpdateEvent) eventWithUpdatedMemberEpoch;
-        assertEquals(updatedMemberEpoch, groupMetadataUpdateEvent.memberEpoch());
-        assertEquals(memberId, groupMetadataUpdateEvent.memberId());
-    }
-
-    @Test
-    public void testConsumerGroupMetadataUpdateWithMemberIdUpdatedAndMemberEpochSame() {
-        makeFirstGroupMetadataUpdate(memberId, memberEpoch);
-
-        time.sleep(2000);
-        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
-
-        assertEquals(1, result.unsentRequests.size());
-        NetworkClientDelegate.UnsentRequest request = result.unsentRequests.get(0);
-        final String updatedMemberId = "updatedMemberId";
-        ClientResponse responseWithMemberIdUpdate = createHeartbeatResponse(
-            request,
-            Errors.NONE,
-            updatedMemberId,
-            memberEpoch
-        );
-        request.handler().onComplete(responseWithMemberIdUpdate);
-        assertEquals(1, backgroundEventQueue.size());
-        final BackgroundEvent eventWithUpdatedMemberEpoch = backgroundEventQueue.poll();
-        assertEquals(BackgroundEvent.Type.GROUP_METADATA_UPDATE, eventWithUpdatedMemberEpoch.type());
-        final GroupMetadataUpdateEvent groupMetadataUpdateEvent = (GroupMetadataUpdateEvent) eventWithUpdatedMemberEpoch;
-        assertEquals(memberEpoch, groupMetadataUpdateEvent.memberEpoch());
-        assertEquals(updatedMemberId, groupMetadataUpdateEvent.memberId());
-    }
-
-    private GroupMetadataUpdateEvent makeFirstGroupMetadataUpdate(final String memberId, final int memberEpoch) {
-        resetWithZeroHeartbeatInterval(Optional.empty());
-        mockStableMember();
-        when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(new Node(1, "localhost", 9999)));
-        NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
-        assertEquals(1, result.unsentRequests.size());
-        NetworkClientDelegate.UnsentRequest request = result.unsentRequests.get(0);
-        ClientResponse firstResponse = createHeartbeatResponse(request, Errors.NONE, memberId, memberEpoch);
-        request.handler().onComplete(firstResponse);
-        assertEquals(1, backgroundEventQueue.size());
-        final BackgroundEvent event = backgroundEventQueue.poll();
-        assertEquals(BackgroundEvent.Type.GROUP_METADATA_UPDATE, event.type());
-        return (GroupMetadataUpdateEvent) event;
-    }
-
     @ParameterizedTest
     @MethodSource("errorProvider")
     public void testHeartbeatResponseOnErrorHandling(final Errors error, final boolean isFatal) {
@@ -430,7 +344,6 @@ public void testHeartbeatResponseOnErrorHandling(final Errors error, final boole
 
         switch (error) {
             case NONE:
-                verify(backgroundEventHandler).add(any(GroupMetadataUpdateEvent.class));
                 verify(membershipManager, times(2)).onHeartbeatSuccess(mockResponse.data());
                 assertEquals(DEFAULT_HEARTBEAT_INTERVAL_MS, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()));
                 break;
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/RequestManagersTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/RequestManagersTest.java
new file mode 100644
index 0000000000000..640c7e98e4225
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/RequestManagersTest.java
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals;
+
+import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.junit.jupiter.api.Test;
+
+import java.util.Optional;
+import java.util.Properties;
+
+import static org.apache.kafka.test.TestUtils.requiredConsumerConfig;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+public class RequestManagersTest {
+
+    @Test
+    public void testMemberStateListenerRegistered() {
+
+        final MemberStateListener listener = (memberEpoch, memberId) -> { };
+
+        final Properties properties = requiredConsumerConfig();
+        properties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "consumerGroup");
+        final ConsumerConfig config = new ConsumerConfig(properties);
+        final GroupRebalanceConfig groupRebalanceConfig = new GroupRebalanceConfig(
+            config,
+            GroupRebalanceConfig.ProtocolType.CONSUMER
+        );
+        final RequestManagers requestManagers = RequestManagers.supplier(
+            new MockTime(),
+            new LogContext(),
+            mock(BackgroundEventHandler.class),
+            mock(ConsumerMetadata.class),
+            mock(SubscriptionState.class),
+            mock(FetchBuffer.class),
+            config,
+            groupRebalanceConfig,
+            mock(ApiVersions.class),
+            mock(FetchMetricsManager.class),
+            () -> mock(NetworkClientDelegate.class),
+            Optional.empty(),
+            new Metrics(),
+            mock(OffsetCommitCallbackInvoker.class),
+            listener
+        ).get();
+        requestManagers.membershipManager.ifPresent(
+            membershipManager -> assertTrue(((MembershipManagerImpl) membershipManager).stateListeners().contains(listener))
+        );
+    }
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/test/TestUtils.java b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
index db86558f7d0b3..3a7dcfc9305d6 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -37,6 +37,7 @@
 import org.apache.kafka.common.requests.ByteBufferChannel;
 import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata;
 import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
@@ -279,6 +280,14 @@ public static Properties producerConfig(final String bootstrapServers, final Cla
         return producerConfig(bootstrapServers, keySerializer, valueSerializer, new Properties());
     }
 
+    public static Properties requiredConsumerConfig() {
+        final Properties consumerConfig = new Properties();
+        consumerConfig.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9091");
+        consumerConfig.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+        consumerConfig.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+        return consumerConfig;
+    }
+
     public static Properties consumerConfig(final String bootstrapServers,
                                             final String groupId,
                                             final Class<?> keyDeserializer,

From c254b22a4877e70617b2710b95ef44b8cc55ce97 Mon Sep 17 00:00:00 2001
From: PoAn Yang <yangpoan@gmail.com>
Date: Mon, 4 Mar 2024 19:50:56 +0800
Subject: [PATCH 104/258] MINOR: simplify ensure topic exists condition
 (#15458)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../main/java/org/apache/kafka/tools/TopicCommand.java | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/TopicCommand.java b/tools/src/main/java/org/apache/kafka/tools/TopicCommand.java
index 043ae521f5b6c..66650cb9db58f 100644
--- a/tools/src/main/java/org/apache/kafka/tools/TopicCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/TopicCommand.java
@@ -208,9 +208,9 @@ private static Integer getReplicationFactor(TopicPartitionInfo tpi, PartitionRea
      *                           If set to true, the command will throw an exception if the topic with the
      *                           requested name does not exist.
      */
-    private static void ensureTopicExists(List<String> foundTopics, String requestedTopic, Boolean requireTopicExists) {
+    private static void ensureTopicExists(List<String> foundTopics, Optional<String> requestedTopic, Boolean requireTopicExists) {
         // If no topic name was mentioned, do not need to throw exception.
-        if (!(requestedTopic.isEmpty() || !Optional.ofNullable(requestedTopic).isPresent()) && requireTopicExists && foundTopics.isEmpty()) {
+        if (requestedTopic.isPresent() && !requestedTopic.get().isEmpty() && requireTopicExists && foundTopics.isEmpty()) {
             // If given topic doesn't exist then throw exception
             throw new IllegalArgumentException(String.format("Topic '%s' does not exist as expected", requestedTopic));
         }
@@ -490,7 +490,7 @@ public void listTopics(TopicCommandOptions opts) throws ExecutionException, Inte
         public void alterTopic(TopicCommandOptions opts) throws ExecutionException, InterruptedException {
             CommandTopicPartition topic = new CommandTopicPartition(opts);
             List<String> topics = getTopics(opts.topic(), opts.excludeInternalTopics());
-            ensureTopicExists(topics, opts.topic().orElse(""), !opts.ifExists());
+            ensureTopicExists(topics, opts.topic(), !opts.ifExists());
 
             if (!topics.isEmpty()) {
                 Map<String, KafkaFuture<org.apache.kafka.clients.admin.TopicDescription>> topicsInfo = adminClient.describeTopics(topics).topicNameValues();
@@ -556,7 +556,7 @@ public void describeTopic(TopicCommandOptions opts) throws ExecutionException, I
             if (useTopicId) {
                 ensureTopicIdExists(topicIds, inputTopicId.get(), !opts.ifExists());
             } else {
-                ensureTopicExists(topics, opts.topic().orElse(""), !opts.ifExists());
+                ensureTopicExists(topics, opts.topic(), !opts.ifExists());
             }
             List<org.apache.kafka.clients.admin.TopicDescription> topicDescriptions = new ArrayList<>();
 
@@ -632,7 +632,7 @@ numPartitions, getReplicationFactor(firstPartition, reassignment),
 
         public void deleteTopic(TopicCommandOptions opts) throws ExecutionException, InterruptedException {
             List<String> topics = getTopics(opts.topic(), opts.excludeInternalTopics());
-            ensureTopicExists(topics, opts.topic().orElse(""), !opts.ifExists());
+            ensureTopicExists(topics, opts.topic(), !opts.ifExists());
             adminClient.deleteTopics(Collections.unmodifiableList(topics),
                 new DeleteTopicsOptions().retryOnQuotaViolation(false)
             ).all().get();

From 4f92a3f0afda96c04059be81cf7867a0bbc7c276 Mon Sep 17 00:00:00 2001
From: Ayoub Omari <ayoubomari1@outlook.fr>
Date: Tue, 5 Mar 2024 00:56:40 +0100
Subject: [PATCH 105/258] KAFKA-14747: record discarded FK join subscription
 responses (#15395)

A foreign-key-join might drop a "subscription response" message, if the value-hash changed.
This PR adds support to record such event via the existing "dropped records" sensor.

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../ResponseJoinProcessorSupplier.java        | 13 +++++
 .../SubscriptionJoinProcessorSupplier.java    |  6 +-
 ...ForeignTableJoinProcessorSupplierTest.java |  2 +-
 .../ResponseJoinProcessorSupplierTest.java    | 56 +++++++++++++++++--
 4 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplier.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplier.java
index 600b28078b9fc..cbb66f98fa934 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplier.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplier.java
@@ -18,6 +18,7 @@
 package org.apache.kafka.streams.kstream.internals.foreignkeyjoin;
 
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.streams.kstream.ValueJoiner;
 import org.apache.kafka.streams.kstream.internals.KTableValueGetter;
@@ -27,6 +28,8 @@
 import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.ProcessorSupplier;
 import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.processor.internals.InternalProcessorContext;
+import org.apache.kafka.streams.processor.internals.metrics.TaskMetrics;
 import org.apache.kafka.streams.state.ValueAndTimestamp;
 import org.apache.kafka.streams.state.internals.Murmur3;
 import org.slf4j.Logger;
@@ -71,6 +74,8 @@ public Processor<K, SubscriptionResponseWrapper<VO>, K, VR> get() {
             private Serializer<V> runtimeValueSerializer = constructionTimeValueSerializer;
 
             private KTableValueGetter<K, V> valueGetter;
+            private Sensor droppedRecordsSensor;
+
 
             @SuppressWarnings("unchecked")
             @Override
@@ -82,6 +87,13 @@ public void init(final ProcessorContext<K, VR> context) {
                 if (runtimeValueSerializer == null) {
                     runtimeValueSerializer = (Serializer<V>) context.valueSerde().serializer();
                 }
+
+                final InternalProcessorContext<?, ?> internalProcessorContext = (InternalProcessorContext<?, ?>) context;
+                droppedRecordsSensor = TaskMetrics.droppedRecordsSensor(
+                        Thread.currentThread().getName(),
+                        internalProcessorContext.taskId().toString(),
+                        internalProcessorContext.metrics()
+                );
             }
 
             @Override
@@ -112,6 +124,7 @@ public void process(final Record<K, SubscriptionResponseWrapper<VO>> record) {
                     context().forward(record.withValue(result));
                 } else {
                     LOG.trace("Dropping FK-join response due to hash mismatch. Expected {}. Actual {}", messageHash, currentHash);
+                    droppedRecordsSensor.record();
                 }
             }
         };
diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/SubscriptionJoinProcessorSupplier.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/SubscriptionJoinProcessorSupplier.java
index a8677ce2958ae..388b669e988a0 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/SubscriptionJoinProcessorSupplier.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/SubscriptionJoinProcessorSupplier.java
@@ -107,7 +107,11 @@ public void process(final Record<CombinedKey<KO, K>, Change<ValueAndTimestamp<Su
 
                         context().forward(
                             record.withKey(record.key().getPrimaryKey())
-                                .withValue(new SubscriptionResponseWrapper<>(value.getHash(), valueToSend, value.getPrimaryPartition()))
+                                .withValue(new SubscriptionResponseWrapper<>(
+                                        value.getHash(),
+                                        valueToSend,
+                                        value.getPrimaryPartition()
+                                ))
                                 .withTimestamp(resultTimestamp)
                         );
                         break;
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ForeignTableJoinProcessorSupplierTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ForeignTableJoinProcessorSupplierTest.java
index f4f35e6ff05e6..c292dd2e34973 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ForeignTableJoinProcessorSupplierTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ForeignTableJoinProcessorSupplierTest.java
@@ -379,4 +379,4 @@ SubscriptionResponseWrapper<String>> processor(final KTableValueGetterSupplier<S
             new SubscriptionJoinProcessorSupplier<>(valueGetterSupplier);
         return supplier.get();
     }
-}
\ No newline at end of file
+}
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplierTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplierTest.java
index 4c26efe236485..b32c51a3baa4d 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplierTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/foreignkeyjoin/ResponseJoinProcessorSupplierTest.java
@@ -17,6 +17,7 @@
 
 package org.apache.kafka.streams.kstream.internals.foreignkeyjoin;
 
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.streams.kstream.ValueJoiner;
 import org.apache.kafka.streams.kstream.internals.KTableValueGetter;
@@ -25,17 +26,23 @@
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.processor.internals.InternalProcessorContext;
 import org.apache.kafka.streams.state.ValueAndTimestamp;
 import org.apache.kafka.streams.state.internals.Murmur3;
+import org.apache.kafka.test.MockInternalNewProcessorContext;
 import org.junit.Test;
 
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import static org.apache.kafka.common.utils.Utils.mkEntry;
+import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.collection.IsEmptyCollection.empty;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 
 public class ResponseJoinProcessorSupplierTest {
     private static final StringSerializer STRING_SERIALIZER = new StringSerializer();
@@ -88,7 +95,7 @@ public void shouldNotForwardWhenHashDoesNotMatch() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final org.apache.kafka.streams.processor.api.MockProcessorContext<String, String> context = new org.apache.kafka.streams.processor.api.MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -97,6 +104,10 @@ public void shouldNotForwardWhenHashDoesNotMatch() {
         processor.process(new Record<>("lhs1", new SubscriptionResponseWrapper<>(oldHash, "rhsValue", 0), 0));
         final List<MockProcessorContext.CapturedForward<? extends String, ? extends String>> forwarded = context.forwarded();
         assertThat(forwarded, empty());
+
+        // test dropped-records sensors
+        assertEquals(1.0, getDroppedRecordsTotalMetric(context));
+        assertNotEquals(0.0, getDroppedRecordsRateMetric(context));
     }
 
     @Test
@@ -113,7 +124,7 @@ public void shouldIgnoreUpdateWhenLeftHasBecomeNull() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final MockProcessorContext<String, String> context = new MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -122,6 +133,10 @@ public void shouldIgnoreUpdateWhenLeftHasBecomeNull() {
         processor.process(new Record<>("lhs1", new SubscriptionResponseWrapper<>(hash, "rhsValue", 0), 0));
         final List<MockProcessorContext.CapturedForward<? extends String, ? extends String>> forwarded = context.forwarded();
         assertThat(forwarded, empty());
+
+        // test dropped-records sensors
+        assertEquals(1.0, getDroppedRecordsTotalMetric(context));
+        assertNotEquals(0.0, getDroppedRecordsRateMetric(context));
     }
 
     @Test
@@ -138,7 +153,7 @@ public void shouldForwardWhenHashMatches() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final MockProcessorContext<String, String> context = new MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -164,7 +179,7 @@ public void shouldEmitTombstoneForInnerJoinWhenRightIsNull() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final MockProcessorContext<String, String> context = new MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -190,7 +205,7 @@ public void shouldEmitResultForLeftJoinWhenRightIsNull() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final MockProcessorContext<String, String> context = new MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -216,7 +231,7 @@ public void shouldEmitTombstoneForLeftJoinWhenRightIsNullAndLeftIsNull() {
                 leftJoin
             );
         final Processor<String, SubscriptionResponseWrapper<String>, String, String> processor = processorSupplier.get();
-        final MockProcessorContext<String, String> context = new MockProcessorContext<>();
+        final MockInternalNewProcessorContext<String, String> context = new MockInternalNewProcessorContext<>();
         processor.init(context);
         context.setRecordMetadata("topic", 0, 0);
 
@@ -227,4 +242,33 @@ public void shouldEmitTombstoneForLeftJoinWhenRightIsNullAndLeftIsNull() {
         assertThat(forwarded.size(), is(1));
         assertThat(forwarded.get(0).record(), is(new Record<>("lhs1", null, 0)));
     }
+
+    private Object getDroppedRecordsTotalMetric(final InternalProcessorContext<String, String> context) {
+        final MetricName dropTotalMetric = new MetricName(
+            "dropped-records-total",
+            "stream-task-metrics",
+            "The total number of dropped records",
+            mkMap(
+                mkEntry("thread-id", Thread.currentThread().getName()),
+                mkEntry("task-id", "0_0")
+            )
+        );
+
+        return context.metrics().metrics().get(dropTotalMetric).metricValue();
+    }
+
+    private Object getDroppedRecordsRateMetric(final InternalProcessorContext<String, String> context) {
+        final MetricName dropRateMetric = new MetricName(
+            "dropped-records-rate",
+            "stream-task-metrics",
+            "The average number of dropped records per second",
+            mkMap(
+                mkEntry("thread-id", Thread.currentThread().getName()),
+                mkEntry("task-id", "0_0")
+            )
+        );
+
+        return context.metrics().metrics().get(dropRateMetric).metricValue();
+    }
+
 }

From 99e511c706b7da08b559a3ff6a2c207cacd47b86 Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Mon, 4 Mar 2024 15:59:47 -0800
Subject: [PATCH 106/258] KAFKA-16288, KAFKA-16289: Fix Values convertToDecimal
 exception and parseString corruption (#15399)

* KAFKA-16288: Prevent ClassCastExceptions for strings in Values.convertToDecimal
* KAFKA-16289: Values inferred schemas for map and arrays should ignore element order

Signed-off-by: Greg Harris <greg.harris@aiven.io>
Reviewers: Chris Egerton <chrise@aiven.io>
---
 .../org/apache/kafka/connect/data/Values.java |  20 +++-
 .../apache/kafka/connect/data/ValuesTest.java | 110 +++++++++++++++---
 2 files changed, 108 insertions(+), 22 deletions(-)

diff --git a/connect/api/src/main/java/org/apache/kafka/connect/data/Values.java b/connect/api/src/main/java/org/apache/kafka/connect/data/Values.java
index cbaeeae66fb27..7b78c64af0ca7 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/Values.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/Values.java
@@ -424,7 +424,7 @@ protected static Object convertTo(Schema toSchema, Schema fromSchema, Object val
                         return BigDecimal.valueOf(converted);
                     }
                     if (value instanceof String) {
-                        return new BigDecimal(value.toString()).doubleValue();
+                        return new BigDecimal(value.toString());
                     }
                 }
                 if (value instanceof ByteBuffer) {
@@ -802,11 +802,12 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
         try {
             if (parser.canConsume(ARRAY_BEGIN_DELIMITER)) {
                 List<Object> result = new ArrayList<>();
+                boolean compatible = true;
                 Schema elementSchema = null;
                 while (parser.hasNext()) {
                     if (parser.canConsume(ARRAY_END_DELIMITER)) {
                         Schema listSchema;
-                        if (elementSchema != null) {
+                        if (elementSchema != null && compatible) {
                             listSchema = SchemaBuilder.array(elementSchema).schema();
                             result = alignListEntriesWithSchema(listSchema, result);
                         } else {
@@ -821,6 +822,9 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
                     }
                     SchemaAndValue element = parse(parser, true);
                     elementSchema = commonSchemaFor(elementSchema, element);
+                    if (elementSchema == null && element != null && element.schema() != null) {
+                        compatible = false;
+                    }
                     result.add(element != null ? element.value() : null);
 
                     int currentPosition = parser.mark();
@@ -840,15 +844,17 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
 
             if (parser.canConsume(MAP_BEGIN_DELIMITER)) {
                 Map<Object, Object> result = new LinkedHashMap<>();
+                boolean keyCompatible = true;
                 Schema keySchema = null;
+                boolean valueCompatible = true;
                 Schema valueSchema = null;
                 while (parser.hasNext()) {
                     if (parser.canConsume(MAP_END_DELIMITER)) {
                         Schema mapSchema;
-                        if (keySchema != null && valueSchema != null) {
+                        if (keySchema != null && valueSchema != null && keyCompatible && valueCompatible) {
                             mapSchema = SchemaBuilder.map(keySchema, valueSchema).build();
                             result = alignMapKeysAndValuesWithSchema(mapSchema, result);
-                        } else if (keySchema != null) {
+                        } else if (keySchema != null && keyCompatible) {
                             mapSchema = SchemaBuilder.mapWithNullValues(keySchema);
                             result = alignMapKeysWithSchema(mapSchema, result);
                         } else {
@@ -876,7 +882,13 @@ protected static SchemaAndValue parse(Parser parser, boolean embedded) throws No
 
                     parser.canConsume(COMMA_DELIMITER);
                     keySchema = commonSchemaFor(keySchema, key);
+                    if (keySchema == null && key.schema() != null) {
+                        keyCompatible = false;
+                    }
                     valueSchema = commonSchemaFor(valueSchema, value);
+                    if (valueSchema == null && value != null && value.schema() != null) {
+                        valueCompatible = false;
+                    }
                 }
                 // Missing either a comma or an end delimiter
                 if (COMMA_DELIMITER.equals(parser.previous())) {
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java b/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
index 3700a6ee4e6cc..abb6ea4221460 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
@@ -25,6 +25,7 @@
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.text.SimpleDateFormat;
 import java.util.ArrayList;
@@ -363,33 +364,73 @@ public void shouldConvertStringOfListWithOnlyNumericElementTypesIntoListOfLarges
     }
 
     /**
-     * The parsed array has byte values and one int value, so we should return list with single unified type of integers.
+     * We parse into different element types, but cannot infer a common element schema.
+     * This behavior should be independent of the order that the elements appear in the string
      */
     @Test
-    public void shouldConvertStringOfListWithMixedElementTypesIntoListWithDifferentElementTypes() {
-        String str = "[1, 2, \"three\"]";
-        List<?> list = Values.convertToList(Schema.STRING_SCHEMA, str);
-        assertEquals(3, list.size());
-        assertEquals(1, ((Number) list.get(0)).intValue());
-        assertEquals(2, ((Number) list.get(1)).intValue());
-        assertEquals("three", list.get(2));
+    public void shouldParseStringListWithMultipleElementTypes() {
+        assertParseStringArrayWithNoSchema(
+                Arrays.asList((byte) 1, (byte) 2, (short) 300, "four"),
+                "[1, 2, 300, \"four\"]");
+        assertParseStringArrayWithNoSchema(
+                Arrays.asList((byte) 2, (short) 300, "four", (byte) 1),
+                "[2, 300, \"four\", 1]");
+        assertParseStringArrayWithNoSchema(
+                Arrays.asList((short) 300, "four", (byte) 1, (byte) 2),
+                "[300, \"four\", 1, 2]");
+        assertParseStringArrayWithNoSchema(
+                Arrays.asList("four", (byte) 1, (byte) 2, (short) 300),
+                "[\"four\", 1, 2, 300]");
+    }
+
+    private void assertParseStringArrayWithNoSchema(List<Object> expected, String str) {
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.ARRAY, result.schema().type());
+        assertNull(result.schema().valueSchema());
+        List<?> list = (List<?>) result.value();
+        assertEquals(expected, list);
     }
 
     /**
-     * We parse into different element types, but cannot infer a common element schema.
+     * Maps with an inconsistent key type don't find a common type for the keys or the values
+     * This behavior should be independent of the order that the pairs appear in the string
+     */
+    @Test
+    public void shouldParseStringMapWithMultipleKeyTypes() {
+        Map<Object, Object> expected = new HashMap<>();
+        expected.put((byte) 1, (byte) 1);
+        expected.put((byte) 2, (byte) 1);
+        expected.put((short) 300, (short) 300);
+        expected.put("four", (byte) 1);
+        assertParseStringMapWithNoSchema(expected, "{1:1, 2:1, 300:300, \"four\":1}");
+        assertParseStringMapWithNoSchema(expected, "{2:1, 300:300, \"four\":1, 1:1}");
+        assertParseStringMapWithNoSchema(expected, "{300:300, \"four\":1, 1:1, 2:1}");
+        assertParseStringMapWithNoSchema(expected, "{\"four\":1, 1:1, 2:1, 300:300}");
+    }
+
+    /**
+     * Maps with a consistent key type may still not have a common type for the values
+     * This behavior should be independent of the order that the pairs appear in the string
      */
     @Test
-    public void shouldParseStringListWithMultipleElementTypesAndReturnListWithNoSchema() {
-        String str = "[1, 2, 3, \"four\"]";
+    public void shouldParseStringMapWithMultipleValueTypes() {
+        Map<Object, Object> expected = new HashMap<>();
+        expected.put((short) 1, (byte) 1);
+        expected.put((short) 2, (byte) 1);
+        expected.put((short) 300, (short) 300);
+        expected.put((short) 4, "four");
+        assertParseStringMapWithNoSchema(expected, "{1:1, 2:1, 300:300, 4:\"four\"}");
+        assertParseStringMapWithNoSchema(expected, "{2:1, 300:300, 4:\"four\", 1:1}");
+        assertParseStringMapWithNoSchema(expected, "{300:300, 4:\"four\", 1:1, 2:1}");
+        assertParseStringMapWithNoSchema(expected, "{4:\"four\", 1:1, 2:1, 300:300}");
+    }
+
+    private void assertParseStringMapWithNoSchema(Map<Object, Object> expected, String str) {
         SchemaAndValue result = Values.parseString(str);
-        assertEquals(Type.ARRAY, result.schema().type());
+        assertEquals(Type.MAP, result.schema().type());
         assertNull(result.schema().valueSchema());
-        List<?> list = (List<?>) result.value();
-        assertEquals(4, list.size());
-        assertEquals(1, ((Number) list.get(0)).intValue());
-        assertEquals(2, ((Number) list.get(1)).intValue());
-        assertEquals(3, ((Number) list.get(2)).intValue());
-        assertEquals("four", list.get(3));
+        Map<?, ?> list = (Map<?, ?>) result.value();
+        assertEquals(expected, list);
     }
 
     /**
@@ -744,6 +785,39 @@ public void shouldConvertTimestampValues() {
         assertEquals(current, ts4);
     }
 
+    @Test
+    public void shouldConvertDecimalValues() {
+        // Various forms of the same number should all be parsed to the same BigDecimal
+        Number number = 1.0f;
+        String string = number.toString();
+        BigDecimal value = new BigDecimal(string);
+        byte[] bytes = Decimal.fromLogical(Decimal.schema(1), value);
+        ByteBuffer buffer = ByteBuffer.wrap(bytes);
+
+        assertEquals(value, Values.convertToDecimal(null, number, 1));
+        assertEquals(value, Values.convertToDecimal(null, string, 1));
+        assertEquals(value, Values.convertToDecimal(null, value, 1));
+        assertEquals(value, Values.convertToDecimal(null, bytes, 1));
+        assertEquals(value, Values.convertToDecimal(null, buffer, 1));
+    }
+
+    /**
+     * Test parsing distinct number-like types (strings containing numbers, and logical Decimals) in the same list
+     * The parser does not convert Numbers to Decimals, or Strings containing numbers to Numbers automatically.
+     */
+    @Test
+    public void shouldNotConvertArrayValuesToDecimal() {
+        List<Object> decimals = Arrays.asList("\"1.0\"", BigDecimal.valueOf(Long.MAX_VALUE).add(BigDecimal.ONE),
+                BigDecimal.valueOf(Long.MIN_VALUE).subtract(BigDecimal.ONE), (byte) 1, (byte) 1);
+        List<Object> expected = new ArrayList<>(decimals); // most values are directly reproduced with the same type
+        expected.set(0, "1.0"); // The quotes are parsed away, but the value remains a string
+        SchemaAndValue schemaAndValue = Values.parseString(decimals.toString());
+        Schema schema = schemaAndValue.schema();
+        assertEquals(Type.ARRAY, schema.type());
+        assertNull(schema.valueSchema());
+        assertEquals(expected, schemaAndValue.value());
+    }
+
     @Test
     public void canConsume() {
     }

From 47792770a28a0b4a6c321a1822b522f4e81068d6 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Tue, 5 Mar 2024 09:44:51 +0100
Subject: [PATCH 107/258] KAFKA-16169: FencedException in commitAsync not
 propagated without callback (#15437)

The javadocs for commitAsync() (w/o callback) say:

@throws org.apache.kafka.common.errors.FencedInstanceIdException
if this consumer instance gets fenced by broker.
If no callback is passed into commitAsync(), no offset commit callback invocation is submitted. However, we only check for a FencedInstanceIdException when we execute a callback. When the consumer gets fenced by another consumer with the same group.instance.id, and we do not use a callback, we miss the exception.

This change modifies the behavior to propagate the FencedInstanceIdException even if no callback is used. The code is kept very similar to the original consumer.

We also change the order - first try to throw the fenced exception, then execute callbacks. That is the order in the original consumer so it's safer to keep it this way.

For testing, we add a unit test that verifies that the FencedInstanceIdException is thrown in that case.

Reviewers: Philip Nee <pnee@confluent.io>, Matthias J. Sax <matthias@confluent.io>
---
 .../internals/AsyncKafkaConsumer.java         | 15 ++++-
 .../internals/CommitRequestManager.java       |  2 +-
 .../OffsetCommitCallbackInvoker.java          | 10 ---
 .../internals/AsyncKafkaConsumerTest.java     | 67 +++++++++++++++++++
 .../internals/CommitRequestManagerTest.java   |  2 -
 5 files changed, 80 insertions(+), 16 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index dafd3293c74c4..5354503c01619 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -111,6 +111,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -271,6 +272,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
     private boolean cachedSubscriptionHasAllFetchPositions;
     private final WakeupTrigger wakeupTrigger = new WakeupTrigger();
     private final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker;
+    private final AtomicBoolean asyncCommitFenced;
 
     // currentThread holds the threadId of the current thread accessing the AsyncKafkaConsumer
     // and is used to prevent multithreaded access
@@ -353,6 +355,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
                     fetchMetricsManager,
                     clientTelemetryReporter.map(ClientTelemetryReporter::telemetrySender).orElse(null));
             this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
+            this.asyncCommitFenced = new AtomicBoolean(false);
             this.groupMetadata.set(initializeGroupMetadata(config, groupRebalanceConfig));
             final Supplier<RequestManagers> requestManagersSupplier = RequestManagers.supplier(time,
                     logContext,
@@ -473,6 +476,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
         this.clientTelemetryReporter = Optional.empty();
         this.autoCommitEnabled = autoCommitEnabled;
         this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
+        this.asyncCommitFenced = new AtomicBoolean(false);
     }
 
     AsyncKafkaConsumer(LogContext logContext,
@@ -538,6 +542,7 @@ private void process(final ConsumerRebalanceListenerCallbackNeededEvent event) {
             client
         );
         this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
+        this.asyncCommitFenced = new AtomicBoolean(false);
         Supplier<RequestManagers> requestManagersSupplier = RequestManagers.supplier(
             time,
             logContext,
@@ -769,6 +774,10 @@ public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCo
                     offsetCommitCallbackInvoker.enqueueInterceptorInvocation(offsets);
                 }
 
+                if (t instanceof FencedInstanceIdException) {
+                    asyncCommitFenced.set(true);
+                }
+
                 if (callback == null) {
                     if (t != null) {
                         log.error("Offset commit with offsets {} failed", offsets, t);
@@ -784,8 +793,8 @@ public void commitAsync(Map<TopicPartition, OffsetAndMetadata> offsets, OffsetCo
     }
 
     private CompletableFuture<Void> commit(final CommitEvent commitEvent) {
-        maybeInvokeCommitCallbacks();
         maybeThrowFencedInstanceException();
+        maybeInvokeCommitCallbacks();
         maybeThrowInvalidGroupIdException();
 
         Map<TopicPartition, OffsetAndMetadata> offsets = commitEvent.offsets();
@@ -1649,8 +1658,8 @@ private void updateLastSeenEpochIfNewer(TopicPartition topicPartition, OffsetAnd
 
     @Override
     public boolean updateAssignmentMetadataIfNeeded(Timer timer) {
-        maybeInvokeCommitCallbacks();
         maybeThrowFencedInstanceException();
+        maybeInvokeCommitCallbacks();
         backgroundEventProcessor.process();
 
         // Keeping this updateAssignmentMetadataIfNeeded wrapping up the updateFetchPositions as
@@ -1905,7 +1914,7 @@ public KafkaConsumerMetrics kafkaConsumerMetrics() {
     }
 
     private void maybeThrowFencedInstanceException() {
-        if (offsetCommitCallbackInvoker.hasFencedException()) {
+        if (asyncCommitFenced.get()) {
             String groupInstanceId = "unknown";
             if (!groupMetadata.get().isPresent()) {
                 log.error("No group metadata found although a group ID was provided. This is a bug!");
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
index 9206783d561be..f7acbde60bdd5 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
@@ -438,7 +438,7 @@ private Throwable commitSyncExceptionForError(Throwable error) {
 
     private Throwable commitAsyncExceptionForError(Throwable error) {
         if (error instanceof RetriableException) {
-            return new RetriableCommitFailedException(error.getMessage());
+            return new RetriableCommitFailedException(error);
         }
         return error;
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetCommitCallbackInvoker.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetCommitCallbackInvoker.java
index db7770cbda865..47a5df6d1d20b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetCommitCallbackInvoker.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetCommitCallbackInvoker.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.FencedInstanceIdException;
 
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
@@ -33,7 +32,6 @@
  */
 public class OffsetCommitCallbackInvoker {
     private final ConsumerInterceptors<?, ?> interceptors;
-    private boolean hasFencedException = false;
 
     OffsetCommitCallbackInvoker(ConsumerInterceptors<?, ?> interceptors) {
         this.interceptors = interceptors;
@@ -62,19 +60,11 @@ public void executeCallbacks() {
         while (!callbackQueue.isEmpty()) {
             OffsetCommitCallbackTask task = callbackQueue.poll();
             if (task != null) {
-
-                if (task.exception instanceof FencedInstanceIdException)
-                    hasFencedException = true;
-
                 task.callback.onComplete(task.offsets, task.exception);
             }
         }
     }
 
-    public boolean hasFencedException() {
-        return hasFencedException;
-    }
-
     private static class OffsetCommitCallbackTask {
         public final Map<TopicPartition, OffsetAndMetadata> offsets;
         public final Exception exception;
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index fe83a417897a6..5777aa245ab88 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -55,6 +55,7 @@
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.InvalidGroupIdException;
 import org.apache.kafka.common.errors.RetriableException;
@@ -570,6 +571,72 @@ public void testCommitAsyncLeaderEpochUpdate() {
         verify(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
     }
 
+    @Test
+    public void testCommitAsyncTriggersFencedExceptionFromCommitAsync() {
+        final String groupId = "consumerGroupA";
+        final String groupInstanceId = "groupInstanceId1";
+        final Properties props = requiredConsumerConfigAndGroupId(groupId);
+        props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId);
+        final ConsumerConfig config = new ConsumerConfig(props);
+        consumer = newConsumer(config);
+        completeCommitAsyncApplicationEventExceptionally(Errors.FENCED_INSTANCE_ID.exception());
+        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
+        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
+        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        consumer.assign(Collections.singleton(tp));
+        consumer.seek(tp, 20);
+
+        assertDoesNotThrow(() -> consumer.commitAsync());
+
+        Exception e = assertThrows(FencedInstanceIdException.class, () -> consumer.commitAsync());
+        assertEquals("Get fenced exception for group.instance.id groupInstanceId1", e.getMessage());
+    }
+
+    @Test
+    public void testCommitSyncTriggersFencedExceptionFromCommitAsync() {
+        final String groupId = "consumerGroupA";
+        final String groupInstanceId = "groupInstanceId1";
+        final Properties props = requiredConsumerConfigAndGroupId(groupId);
+        props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId);
+        final ConsumerConfig config = new ConsumerConfig(props);
+        consumer = newConsumer(config);
+        completeCommitAsyncApplicationEventExceptionally(Errors.FENCED_INSTANCE_ID.exception());
+        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
+        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
+        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        consumer.assign(Collections.singleton(tp));
+        consumer.seek(tp, 20);
+
+        assertDoesNotThrow(() -> consumer.commitAsync());
+
+        Exception e =  assertThrows(FencedInstanceIdException.class, () -> consumer.commitSync());
+        assertEquals("Get fenced exception for group.instance.id groupInstanceId1", e.getMessage());
+    }
+
+    @Test
+    public void testPollTriggersFencedExceptionFromCommitAsync() {
+        final String groupId = "consumerGroupA";
+        final String groupInstanceId = "groupInstanceId1";
+        final Properties props = requiredConsumerConfigAndGroupId(groupId);
+        props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId);
+        final ConsumerConfig config = new ConsumerConfig(props);
+        consumer = newConsumer(config);
+        completeCommitAsyncApplicationEventExceptionally(Errors.FENCED_INSTANCE_ID.exception());
+        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
+        completeFetchedCommittedOffsetApplicationEventSuccessfully(mkMap());
+        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        consumer.assign(Collections.singleton(tp));
+        consumer.seek(tp, 20);
+
+        assertDoesNotThrow(() -> consumer.commitAsync());
+
+        Exception e = assertThrows(FencedInstanceIdException.class, () -> consumer.poll(Duration.ZERO));
+        assertEquals("Get fenced exception for group.instance.id groupInstanceId1", e.getMessage());
+    }
+
     @Test
     public void testEnsurePollExecutedCommitAsyncCallbacks() {
         consumer = newConsumer();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
index c27494d69a038..7e18924e7b136 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
@@ -75,7 +75,6 @@
 import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_GROUP_ID;
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_GROUP_INSTANCE_ID;
-import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
 import static org.apache.kafka.test.TestUtils.assertFutureThrows;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -96,7 +95,6 @@ public class CommitRequestManagerTest {
 
     private long retryBackoffMs = 100;
     private long retryBackoffMaxMs = 1000;
-    private String consumerMetricGroupPrefix = CONSUMER_METRIC_GROUP_PREFIX;
     private static final String CONSUMER_COORDINATOR_METRICS = "consumer-coordinator-metrics";
     private Node mockedNode = new Node(1, "host1", 9092);
     private SubscriptionState subscriptionState;

From eea369af947dcff567f849183ba2217ac5e9a2ba Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Tue, 5 Mar 2024 13:11:56 +0300
Subject: [PATCH 108/258] KAFKA-14588 Log cleaner configuration move to
 CleanerConfig (#15387)

In order to move ConfigCommand to tools we must move all it's dependencies which includes KafkaConfig and other core classes to java. This PR moves log cleaner configuration to CleanerConfig class of storage module.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle                                  |  2 +
 checkstyle/import-control-core.xml            |  1 +
 checkstyle/import-control.xml                 |  2 +
 .../util/clusters/EmbeddedKafkaCluster.java   |  3 +-
 .../src/main/scala/kafka/log/LogCleaner.scala | 12 +--
 .../main/scala/kafka/server/KafkaConfig.scala | 76 ++++++-------------
 .../kafka/testkit/KafkaClusterTestKit.java    |  3 +-
 .../api/PlaintextAdminIntegrationTest.scala   | 12 +--
 .../DynamicBrokerReconfigurationTest.scala    | 32 ++++----
 ...gCleanerParameterizedIntegrationTest.scala | 12 +--
 .../scala/unit/kafka/log/LogCleanerTest.scala |  8 +-
 .../kafka/server/ControllerApisTest.scala     | 13 ++--
 .../server/DynamicBrokerConfigTest.scala      | 18 ++---
 .../unit/kafka/server/KafkaConfigTest.scala   | 18 ++---
 .../scala/unit/kafka/utils/TestUtils.scala    |  2 +-
 .../config/ServerTopicConfigSynonyms.java     |  2 +-
 .../apache/kafka/server/config/Defaults.java  |  7 --
 .../storage/internals/log/CleanerConfig.java  | 42 ++++++++++
 .../utils/EmbeddedKafkaCluster.java           |  3 +-
 19 files changed, 141 insertions(+), 127 deletions(-)

diff --git a/build.gradle b/build.gradle
index f2ba22e86b39d..7dc8f50342fdd 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2177,6 +2177,7 @@ project(':streams') {
     testImplementation project(':core')
     testImplementation project(':tools')
     testImplementation project(':core').sourceSets.test.output
+    testImplementation project(':storage')
     testImplementation project(':server-common')
     testImplementation project(':server-common').sourceSets.test.output
     testImplementation libs.log4j
@@ -2975,6 +2976,7 @@ project(':connect:runtime') {
     testImplementation project(':metadata')
     testImplementation project(':core').sourceSets.test.output
     testImplementation project(':server-common')
+    testImplementation project(':storage')
     testImplementation project(':connect:test-plugins')
 
     testImplementation libs.easymock
diff --git a/checkstyle/import-control-core.xml b/checkstyle/import-control-core.xml
index 3f9a21fffc6c8..782d2fe4617e9 100644
--- a/checkstyle/import-control-core.xml
+++ b/checkstyle/import-control-core.xml
@@ -56,6 +56,7 @@
     <allow pkg="org.apache.kafka.metalog" />
     <allow pkg="org.apache.kafka.server.common" />
     <allow pkg="org.apache.kafka.server.fault" />
+    <allow class="org.apache.kafka.storage.internals.log.CleanerConfig" />
   </subpackage>
 
   <subpackage name="tools">
diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index caf1fe5ebe146..a52b3d94e3240 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -410,6 +410,7 @@
       <allow pkg="com.fasterxml.jackson" />
       <allow pkg="org.apache.kafka.tools" />
       <allow pkg="org.apache.kafka.server.config" />
+      <allow class="org.apache.kafka.storage.internals.log.CleanerConfig" />
     </subpackage>
 
     <subpackage name="test">
@@ -601,6 +602,7 @@
         <allow pkg="com.fasterxml.jackson.core.type" />
         <allow pkg="org.apache.kafka.metadata" />
         <allow pkg="org.eclipse.jetty.client"/>
+        <allow class="org.apache.kafka.storage.internals.log.CleanerConfig" />
       </subpackage>
     </subpackage>
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
index e3694cae2911a..c15aa27ae5964 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
@@ -55,6 +55,7 @@
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.metadata.BrokerState;
+import org.apache.kafka.storage.internals.log.CleanerConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -160,7 +161,7 @@ private void doStart() {
         putIfAbsent(brokerConfig, KafkaConfig.OffsetsTopicReplicationFactorProp(), (short) brokers.length);
         putIfAbsent(brokerConfig, KafkaConfig.AutoCreateTopicsEnableProp(), false);
         // reduce the size of the log cleaner map to reduce test memory usage
-        putIfAbsent(brokerConfig, KafkaConfig.LogCleanerDedupeBufferSizeProp(), 2 * 1024 * 1024L);
+        putIfAbsent(brokerConfig, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, 2 * 1024 * 1024L);
 
         Object listenerConfig = brokerConfig.get(KafkaConfig.InterBrokerListenerNameProp());
         if (listenerConfig == null)
diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala
index 8098ea237e06d..b653f40b287a8 100644
--- a/core/src/main/scala/kafka/log/LogCleaner.scala
+++ b/core/src/main/scala/kafka/log/LogCleaner.scala
@@ -499,13 +499,13 @@ class LogCleaner(initialConfig: CleanerConfig,
 
 object LogCleaner {
   val ReconfigurableConfigs: Set[String] = Set(
-    KafkaConfig.LogCleanerThreadsProp,
-    KafkaConfig.LogCleanerDedupeBufferSizeProp,
-    KafkaConfig.LogCleanerDedupeBufferLoadFactorProp,
-    KafkaConfig.LogCleanerIoBufferSizeProp,
+    CleanerConfig.LOG_CLEANER_THREADS_PROP,
+    CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP,
+    CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP,
+    CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_PROP,
     KafkaConfig.MessageMaxBytesProp,
-    KafkaConfig.LogCleanerIoMaxBytesPerSecondProp,
-    KafkaConfig.LogCleanerBackoffMsProp
+    CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP,
+    CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP
   )
 
   def cleanerConfig(config: KafkaConfig): CleanerConfig = {
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index c6f51a000e2a5..eaddb047da8d5 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -49,7 +49,7 @@ import org.apache.kafka.server.config.{Defaults, ServerTopicConfigSynonyms}
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.record.BrokerCompressionType
 import org.apache.kafka.server.util.Csv
-import org.apache.kafka.storage.internals.log.{LogConfig, ProducerStateManagerConfig}
+import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig, ProducerStateManagerConfig}
 import org.apache.kafka.storage.internals.log.LogConfig.MessageFormatVersion
 import org.apache.zookeeper.client.ZKClientConfig
 
@@ -211,17 +211,6 @@ object KafkaConfig {
   val LogRetentionBytesProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.RETENTION_BYTES_CONFIG)
   val LogCleanupIntervalMsProp = LogConfigPrefix + "retention.check.interval.ms"
   val LogCleanupPolicyProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.CLEANUP_POLICY_CONFIG)
-  val LogCleanerThreadsProp = LogConfigPrefix + "cleaner.threads"
-  val LogCleanerIoMaxBytesPerSecondProp = LogConfigPrefix + "cleaner.io.max.bytes.per.second"
-  val LogCleanerDedupeBufferSizeProp = LogConfigPrefix + "cleaner.dedupe.buffer.size"
-  val LogCleanerIoBufferSizeProp = LogConfigPrefix + "cleaner.io.buffer.size"
-  val LogCleanerDedupeBufferLoadFactorProp = LogConfigPrefix + "cleaner.io.buffer.load.factor"
-  val LogCleanerBackoffMsProp = LogConfigPrefix + "cleaner.backoff.ms"
-  val LogCleanerMinCleanRatioProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MIN_CLEANABLE_DIRTY_RATIO_CONFIG)
-  val LogCleanerEnableProp = LogConfigPrefix + "cleaner.enable"
-  val LogCleanerDeleteRetentionMsProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.DELETE_RETENTION_MS_CONFIG)
-  val LogCleanerMinCompactionLagMsProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MIN_COMPACTION_LAG_MS_CONFIG)
-  val LogCleanerMaxCompactionLagMsProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MAX_COMPACTION_LAG_MS_CONFIG)
   val LogIndexSizeMaxBytesProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.SEGMENT_INDEX_BYTES_CONFIG)
   val LogIndexIntervalBytesProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.INDEX_INTERVAL_BYTES_CONFIG)
   val LogFlushIntervalMessagesProp = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.FLUSH_MESSAGES_INTERVAL_CONFIG)
@@ -652,25 +641,6 @@ object KafkaConfig {
   val LogRetentionBytesDoc = "The maximum size of the log before deleting it"
   val LogCleanupIntervalMsDoc = "The frequency in milliseconds that the log cleaner checks whether any log is eligible for deletion"
   val LogCleanupPolicyDoc = "The default cleanup policy for segments beyond the retention window. A comma separated list of valid policies. Valid policies are: \"delete\" and \"compact\""
-  val LogCleanerThreadsDoc = "The number of background threads to use for log cleaning"
-  val LogCleanerIoMaxBytesPerSecondDoc = "The log cleaner will be throttled so that the sum of its read and write i/o will be less than this value on average"
-  val LogCleanerDedupeBufferSizeDoc = "The total memory used for log deduplication across all cleaner threads"
-  val LogCleanerIoBufferSizeDoc = "The total memory used for log cleaner I/O buffers across all cleaner threads"
-  val LogCleanerDedupeBufferLoadFactorDoc = "Log cleaner dedupe buffer load factor. The percentage full the dedupe buffer can become. A higher value " +
-  "will allow more log to be cleaned at once but will lead to more hash collisions"
-  val LogCleanerBackoffMsDoc = "The amount of time to sleep when there are no logs to clean"
-  val LogCleanerMinCleanRatioDoc = "The minimum ratio of dirty log to total log for a log to eligible for cleaning. " +
-    "If the " + LogCleanerMaxCompactionLagMsProp + " or the " + LogCleanerMinCompactionLagMsProp +
-    " configurations are also specified, then the log compactor considers the log eligible for compaction " +
-    "as soon as either: (i) the dirty ratio threshold has been met and the log has had dirty (uncompacted) " +
-    "records for at least the " + LogCleanerMinCompactionLagMsProp + " duration, or (ii) if the log has had " +
-    "dirty (uncompacted) records for at most the " + LogCleanerMaxCompactionLagMsProp + " period."
-  val LogCleanerEnableDoc = "Enable the log cleaner process to run on the server. Should be enabled if using any topics with a cleanup.policy=compact including the internal offsets topic. If disabled those topics will not be compacted and continually grow in size."
-  val LogCleanerDeleteRetentionMsDoc = "The amount of time to retain tombstone message markers for log compacted topics. This setting also gives a bound " +
-    "on the time in which a consumer must complete a read if they begin from offset 0 to ensure that they get a valid snapshot of the final stage (otherwise  " +
-    "tombstones messages may be collected before a consumer completes their scan)."
-  val LogCleanerMinCompactionLagMsDoc = "The minimum time a message will remain uncompacted in the log. Only applicable for logs that are being compacted."
-  val LogCleanerMaxCompactionLagMsDoc = "The maximum time a message will remain ineligible for compaction in the log. Only applicable for logs that are being compacted."
   val LogIndexSizeMaxBytesDoc = "The maximum size in bytes of the offset index"
   val LogIndexIntervalBytesDoc = "The interval with which we add an entry to the offset index."
   val LogFlushIntervalMessagesDoc = "The number of messages accumulated on a log partition before messages are flushed to disk."
@@ -1075,17 +1045,17 @@ object KafkaConfig {
       .define(LogRetentionBytesProp, LONG, LogConfig.DEFAULT_RETENTION_BYTES, HIGH, LogRetentionBytesDoc)
       .define(LogCleanupIntervalMsProp, LONG, Defaults.LOG_CLEANUP_INTERVAL_MS, atLeast(1), MEDIUM, LogCleanupIntervalMsDoc)
       .define(LogCleanupPolicyProp, LIST, LogConfig.DEFAULT_CLEANUP_POLICY, ValidList.in(TopicConfig.CLEANUP_POLICY_COMPACT, TopicConfig.CLEANUP_POLICY_DELETE), MEDIUM, LogCleanupPolicyDoc)
-      .define(LogCleanerThreadsProp, INT, Defaults.LOG_CLEANER_THREADS, atLeast(0), MEDIUM, LogCleanerThreadsDoc)
-      .define(LogCleanerIoMaxBytesPerSecondProp, DOUBLE, Defaults.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND, MEDIUM, LogCleanerIoMaxBytesPerSecondDoc)
-      .define(LogCleanerDedupeBufferSizeProp, LONG, Defaults.LOG_CLEANER_DEDUPE_BUFFER_SIZE, MEDIUM, LogCleanerDedupeBufferSizeDoc)
-      .define(LogCleanerIoBufferSizeProp, INT, Defaults.LOG_CLEANER_IO_BUFFER_SIZE, atLeast(0), MEDIUM, LogCleanerIoBufferSizeDoc)
-      .define(LogCleanerDedupeBufferLoadFactorProp, DOUBLE, Defaults.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR, MEDIUM, LogCleanerDedupeBufferLoadFactorDoc)
-      .define(LogCleanerBackoffMsProp, LONG, Defaults.LOG_CLEANER_BACKOFF_MS, atLeast(0), MEDIUM, LogCleanerBackoffMsDoc)
-      .define(LogCleanerMinCleanRatioProp, DOUBLE, LogConfig.DEFAULT_MIN_CLEANABLE_DIRTY_RATIO, between(0, 1), MEDIUM, LogCleanerMinCleanRatioDoc)
-      .define(LogCleanerEnableProp, BOOLEAN, Defaults.LOG_CLEANER_ENABLE, MEDIUM, LogCleanerEnableDoc)
-      .define(LogCleanerDeleteRetentionMsProp, LONG, LogConfig.DEFAULT_DELETE_RETENTION_MS, atLeast(0), MEDIUM, LogCleanerDeleteRetentionMsDoc)
-      .define(LogCleanerMinCompactionLagMsProp, LONG, LogConfig.DEFAULT_MIN_COMPACTION_LAG_MS, atLeast(0), MEDIUM, LogCleanerMinCompactionLagMsDoc)
-      .define(LogCleanerMaxCompactionLagMsProp, LONG, LogConfig.DEFAULT_MAX_COMPACTION_LAG_MS, atLeast(1), MEDIUM, LogCleanerMaxCompactionLagMsDoc)
+      .define(CleanerConfig.LOG_CLEANER_THREADS_PROP, INT, CleanerConfig.LOG_CLEANER_THREADS, atLeast(0), MEDIUM, CleanerConfig.LOG_CLEANER_THREADS_DOC)
+      .define(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP, DOUBLE, CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND, MEDIUM, CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_DOC)
+      .define(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, LONG, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE, MEDIUM, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_DOC)
+      .define(CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_PROP, INT, CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE, atLeast(0), MEDIUM, CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_DOC)
+      .define(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP, DOUBLE, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR, MEDIUM, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_DOC)
+      .define(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP, LONG, CleanerConfig.LOG_CLEANER_BACKOFF_MS, atLeast(0), MEDIUM, CleanerConfig.LOG_CLEANER_BACKOFF_MS_DOC)
+      .define(CleanerConfig.LOG_CLEANER_MIN_CLEAN_RATIO_PROP, DOUBLE, LogConfig.DEFAULT_MIN_CLEANABLE_DIRTY_RATIO, between(0, 1), MEDIUM, CleanerConfig.LOG_CLEANER_MIN_CLEAN_RATIO_DOC)
+      .define(CleanerConfig.LOG_CLEANER_ENABLE_PROP, BOOLEAN, CleanerConfig.LOG_CLEANER_ENABLE, MEDIUM, CleanerConfig.LOG_CLEANER_ENABLE_DOC)
+      .define(CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP, LONG, LogConfig.DEFAULT_DELETE_RETENTION_MS, atLeast(0), MEDIUM, CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_DOC)
+      .define(CleanerConfig.LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP, LONG, LogConfig.DEFAULT_MIN_COMPACTION_LAG_MS, atLeast(0), MEDIUM, CleanerConfig.LOG_CLEANER_MIN_COMPACTION_LAG_MS_DOC)
+      .define(CleanerConfig.LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP, LONG, LogConfig.DEFAULT_MAX_COMPACTION_LAG_MS, atLeast(1), MEDIUM, CleanerConfig.LOG_CLEANER_MAX_COMPACTION_LAG_MS_DOC)
       .define(LogIndexSizeMaxBytesProp, INT, LogConfig.DEFAULT_SEGMENT_INDEX_BYTES, atLeast(4), MEDIUM, LogIndexSizeMaxBytesDoc)
       .define(LogIndexIntervalBytesProp, INT, LogConfig.DEFAULT_INDEX_INTERVAL_BYTES, atLeast(0), MEDIUM, LogIndexIntervalBytesDoc)
       .define(LogFlushIntervalMessagesProp, LONG, LogConfig.DEFAULT_FLUSH_MESSAGES_INTERVAL, atLeast(1), HIGH, LogFlushIntervalMessagesDoc)
@@ -1652,7 +1622,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
   val logDirs = CoreUtils.parseCsvList(Option(getString(KafkaConfig.LogDirsProp)).getOrElse(getString(KafkaConfig.LogDirProp)))
   def logSegmentBytes = getInt(KafkaConfig.LogSegmentBytesProp)
   def logFlushIntervalMessages = getLong(KafkaConfig.LogFlushIntervalMessagesProp)
-  val logCleanerThreads = getInt(KafkaConfig.LogCleanerThreadsProp)
+  val logCleanerThreads = getInt(CleanerConfig.LOG_CLEANER_THREADS_PROP)
   def numRecoveryThreadsPerDataDir = getInt(KafkaConfig.NumRecoveryThreadsPerDataDirProp)
   val logFlushSchedulerIntervalMs = getLong(KafkaConfig.LogFlushSchedulerIntervalMsProp)
   val logFlushOffsetCheckpointIntervalMs = getInt(KafkaConfig.LogFlushOffsetCheckpointIntervalMsProp).toLong
@@ -1662,16 +1632,16 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
   val offsetsRetentionMinutes = getInt(KafkaConfig.OffsetsRetentionMinutesProp)
   val offsetsRetentionCheckIntervalMs = getLong(KafkaConfig.OffsetsRetentionCheckIntervalMsProp)
   def logRetentionBytes = getLong(KafkaConfig.LogRetentionBytesProp)
-  val logCleanerDedupeBufferSize = getLong(KafkaConfig.LogCleanerDedupeBufferSizeProp)
-  val logCleanerDedupeBufferLoadFactor = getDouble(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp)
-  val logCleanerIoBufferSize = getInt(KafkaConfig.LogCleanerIoBufferSizeProp)
-  val logCleanerIoMaxBytesPerSecond = getDouble(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp)
-  def logCleanerDeleteRetentionMs = getLong(KafkaConfig.LogCleanerDeleteRetentionMsProp)
-  def logCleanerMinCompactionLagMs = getLong(KafkaConfig.LogCleanerMinCompactionLagMsProp)
-  def logCleanerMaxCompactionLagMs = getLong(KafkaConfig.LogCleanerMaxCompactionLagMsProp)
-  val logCleanerBackoffMs = getLong(KafkaConfig.LogCleanerBackoffMsProp)
-  def logCleanerMinCleanRatio = getDouble(KafkaConfig.LogCleanerMinCleanRatioProp)
-  val logCleanerEnable = getBoolean(KafkaConfig.LogCleanerEnableProp)
+  val logCleanerDedupeBufferSize = getLong(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP)
+  val logCleanerDedupeBufferLoadFactor = getDouble(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP)
+  val logCleanerIoBufferSize = getInt(CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_PROP)
+  val logCleanerIoMaxBytesPerSecond = getDouble(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP)
+  def logCleanerDeleteRetentionMs = getLong(CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP)
+  def logCleanerMinCompactionLagMs = getLong(CleanerConfig.LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP)
+  def logCleanerMaxCompactionLagMs = getLong(CleanerConfig.LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP)
+  val logCleanerBackoffMs = getLong(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP)
+  def logCleanerMinCleanRatio = getDouble(CleanerConfig.LOG_CLEANER_MIN_CLEAN_RATIO_PROP)
+  val logCleanerEnable = getBoolean(CleanerConfig.LOG_CLEANER_ENABLE_PROP)
   def logIndexSizeMaxBytes = getInt(KafkaConfig.LogIndexSizeMaxBytesProp)
   def logIndexIntervalBytes = getInt(KafkaConfig.LogIndexIntervalBytesProp)
   def logDeleteDelayMs = getLong(KafkaConfig.LogDeleteDelayMsProp)
diff --git a/core/src/test/java/kafka/testkit/KafkaClusterTestKit.java b/core/src/test/java/kafka/testkit/KafkaClusterTestKit.java
index fa11fa9b186f1..4c6d1f9ef4adc 100644
--- a/core/src/test/java/kafka/testkit/KafkaClusterTestKit.java
+++ b/core/src/test/java/kafka/testkit/KafkaClusterTestKit.java
@@ -41,6 +41,7 @@
 import org.apache.kafka.server.common.ApiMessageAndVersion;
 import org.apache.kafka.server.fault.FaultHandler;
 import org.apache.kafka.server.fault.MockFaultHandler;
+import org.apache.kafka.storage.internals.log.CleanerConfig;
 import org.apache.kafka.test.TestUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -192,7 +193,7 @@ private KafkaConfig createNodeConfig(TestKitNode node) {
             props.put(RaftConfig.QUORUM_VOTERS_CONFIG, uninitializedQuorumVotersString);
 
             // reduce log cleaner offset map memory usage
-            props.put(KafkaConfig$.MODULE$.LogCleanerDedupeBufferSizeProp(), "2097152");
+            props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, "2097152");
 
             // Add associated broker node property overrides
             if (brokerNode != null) {
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
index 54be3337625d4..e084454f5ff7b 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -45,7 +45,7 @@ import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{ConsumerGroupState, ElectionType, TopicCollection, TopicPartition, TopicPartitionInfo, TopicPartitionReplica, Uuid}
 import org.apache.kafka.controller.ControllerRequestContextUtil.ANONYMOUS_CONTEXT
 import org.apache.kafka.server.config.Defaults
-import org.apache.kafka.storage.internals.log.LogConfig
+import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig}
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Disabled, TestInfo}
 import org.junit.jupiter.params.ParameterizedTest
@@ -447,7 +447,7 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
       configs.get(brokerResource2).entries.size)
     assertEquals(brokers(2).config.brokerId.toString, configs.get(brokerResource2).get(KafkaConfig.BrokerIdProp).value)
     assertEquals(brokers(2).config.logCleanerThreads.toString,
-      configs.get(brokerResource2).get(KafkaConfig.LogCleanerThreadsProp).value)
+      configs.get(brokerResource2).get(CleanerConfig.LOG_CLEANER_THREADS_PROP).value)
 
     checkValidAlterConfigs(client, this, topicResource1, topicResource2)
   }
@@ -2532,7 +2532,7 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
       .all().get(15, TimeUnit.SECONDS)
 
     val newLogCleanerDeleteRetention = new Properties
-    newLogCleanerDeleteRetention.put(KafkaConfig.LogCleanerDeleteRetentionMsProp, "34")
+    newLogCleanerDeleteRetention.put(CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP, "34")
     TestUtils.incrementalAlterConfigs(brokers, client, newLogCleanerDeleteRetention, perBrokerConfig = true)
       .all().get(15, TimeUnit.SECONDS)
 
@@ -2543,14 +2543,14 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
         controllerServer.config.nodeId.toString)
       controllerServer.controller.incrementalAlterConfigs(ANONYMOUS_CONTEXT,
         Collections.singletonMap(controllerNodeResource,
-          Collections.singletonMap(KafkaConfig.LogCleanerDeleteRetentionMsProp,
+          Collections.singletonMap(CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP,
             new SimpleImmutableEntry(AlterConfigOp.OpType.SET, "34"))), false).get()
       ensureConsistentKRaftMetadata()
     }
 
     waitUntilTrue(() => brokers.forall(_.config.originals.getOrDefault(
-      KafkaConfig.LogCleanerDeleteRetentionMsProp, "").toString.equals("34")),
-      s"Timed out waiting for change to ${KafkaConfig.LogCleanerDeleteRetentionMsProp}",
+      CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP, "").toString.equals("34")),
+      s"Timed out waiting for change to ${CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP}",
       waitTimeMs = 60000L)
 
     waitUntilTrue(() => brokers.forall(_.config.originals.getOrDefault(
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
index 4be668aa32113..c984eae0272c2 100644
--- a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -254,7 +254,7 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
     expectedProps.setProperty(KafkaConfig.LogRetentionTimeMillisProp, "1680000000")
     expectedProps.setProperty(KafkaConfig.LogRetentionTimeHoursProp, "168")
     expectedProps.setProperty(KafkaConfig.LogRollTimeHoursProp, "168")
-    expectedProps.setProperty(KafkaConfig.LogCleanerThreadsProp, "1")
+    expectedProps.setProperty(CleanerConfig.LOG_CLEANER_THREADS_PROP, "1")
     val logRetentionMs = configEntry(configDesc, KafkaConfig.LogRetentionTimeMillisProp)
     verifyConfig(KafkaConfig.LogRetentionTimeMillisProp, logRetentionMs,
       isSensitive = false, isReadOnly = false, expectedProps)
@@ -264,8 +264,8 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
     val logRollHours = configEntry(configDesc, KafkaConfig.LogRollTimeHoursProp)
     verifyConfig(KafkaConfig.LogRollTimeHoursProp, logRollHours,
       isSensitive = false, isReadOnly = true, expectedProps)
-    val logCleanerThreads = configEntry(configDesc, KafkaConfig.LogCleanerThreadsProp)
-    verifyConfig(KafkaConfig.LogCleanerThreadsProp, logCleanerThreads,
+    val logCleanerThreads = configEntry(configDesc, CleanerConfig.LOG_CLEANER_THREADS_PROP)
+    verifyConfig(CleanerConfig.LOG_CLEANER_THREADS_PROP, logCleanerThreads,
       isSensitive = false, isReadOnly = false, expectedProps)
 
     def synonymsList(configEntry: ConfigEntry): List[(String, ConfigSource)] =
@@ -278,7 +278,7 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
       (KafkaConfig.LogRetentionTimeHoursProp, ConfigSource.DEFAULT_CONFIG)),
       synonymsList(logRetentionHours))
     assertEquals(List((KafkaConfig.LogRollTimeHoursProp, ConfigSource.DEFAULT_CONFIG)), synonymsList(logRollHours))
-    assertEquals(List((KafkaConfig.LogCleanerThreadsProp, ConfigSource.DEFAULT_CONFIG)), synonymsList(logCleanerThreads))
+    assertEquals(List((CleanerConfig.LOG_CLEANER_THREADS_PROP, ConfigSource.DEFAULT_CONFIG)), synonymsList(logCleanerThreads))
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
@@ -536,19 +536,19 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
     verifyThreads("kafka-log-cleaner-thread-", countPerBroker = 1)
 
     val props = new Properties
-    props.put(KafkaConfig.LogCleanerThreadsProp, "2")
-    props.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, "20000000")
-    props.put(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp, "0.8")
-    props.put(KafkaConfig.LogCleanerIoBufferSizeProp, "300000")
+    props.put(CleanerConfig.LOG_CLEANER_THREADS_PROP, "2")
+    props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, "20000000")
+    props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP, "0.8")
+    props.put(CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_PROP, "300000")
     props.put(KafkaConfig.MessageMaxBytesProp, "40000")
-    props.put(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, "50000000")
-    props.put(KafkaConfig.LogCleanerBackoffMsProp, "6000")
+    props.put(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP, "50000000")
+    props.put(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP, "6000")
 
     // Verify cleaner config was updated. Wait for one of the configs to be updated and verify
     // that all other others were updated at the same time since they are reconfigured together
     var newCleanerConfig: CleanerConfig = null
     TestUtils.waitUntilTrue(() => {
-      reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogCleanerThreadsProp, "2"))
+      reconfigureServers(props, perBrokerConfig = false, (CleanerConfig.LOG_CLEANER_THREADS_PROP, "2"))
       newCleanerConfig = servers.head.logManager.cleaner.currentConfig
       newCleanerConfig.numThreads == 2
     }, "Log cleaner not reconfigured", 60000)
@@ -566,8 +566,8 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
     def cleanerThreads = Thread.getAllStackTraces.keySet.asScala.filter(_.getName.startsWith("kafka-log-cleaner-thread-"))
     cleanerThreads.take(2).foreach(_.interrupt())
     TestUtils.waitUntilTrue(() => cleanerThreads.size == (2 * numServers) - 2, "Threads did not exit")
-    props.put(KafkaConfig.LogCleanerBackoffMsProp, "8000")
-    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogCleanerBackoffMsProp, "8000"))
+    props.put(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP, "8000")
+    reconfigureServers(props, perBrokerConfig = false, (CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP, "8000"))
     verifyThreads("kafka-log-cleaner-thread-", countPerBroker = 2)
 
     // Verify that produce/consume worked throughout this test without any retries in producer
@@ -635,10 +635,10 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
     props.put(KafkaConfig.LogRetentionTimeMillisProp, TimeUnit.DAYS.toMillis(1).toString)
     props.put(KafkaConfig.MessageMaxBytesProp, "100000")
     props.put(KafkaConfig.LogIndexIntervalBytesProp, "10000")
-    props.put(KafkaConfig.LogCleanerDeleteRetentionMsProp, TimeUnit.DAYS.toMillis(1).toString)
-    props.put(KafkaConfig.LogCleanerMinCompactionLagMsProp, "60000")
+    props.put(CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP, TimeUnit.DAYS.toMillis(1).toString)
+    props.put(CleanerConfig.LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP, "60000")
     props.put(KafkaConfig.LogDeleteDelayMsProp, "60000")
-    props.put(KafkaConfig.LogCleanerMinCleanRatioProp, "0.3")
+    props.put(CleanerConfig.LOG_CLEANER_MIN_CLEAN_RATIO_PROP, "0.3")
     props.put(KafkaConfig.LogCleanupPolicyProp, "delete")
     props.put(KafkaConfig.UncleanLeaderElectionEnableProp, "false")
     props.put(KafkaConfig.MinInSyncReplicasProp, "2")
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
index 8729045db7e79..49e518ac2af89 100755
--- a/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
@@ -251,13 +251,13 @@ class LogCleanerParameterizedIntegrationTest extends AbstractLogCleanerIntegrati
 
     def kafkaConfigWithCleanerConfig(cleanerConfig: CleanerConfig): KafkaConfig = {
       val props = TestUtils.createBrokerConfig(0, "localhost:2181")
-      props.put(KafkaConfig.LogCleanerThreadsProp, cleanerConfig.numThreads.toString)
-      props.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, cleanerConfig.dedupeBufferSize.toString)
-      props.put(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp, cleanerConfig.dedupeBufferLoadFactor.toString)
-      props.put(KafkaConfig.LogCleanerIoBufferSizeProp, cleanerConfig.ioBufferSize.toString)
+      props.put(CleanerConfig.LOG_CLEANER_THREADS_PROP, cleanerConfig.numThreads.toString)
+      props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, cleanerConfig.dedupeBufferSize.toString)
+      props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP, cleanerConfig.dedupeBufferLoadFactor.toString)
+      props.put(CleanerConfig.LOG_CLEANER_IO_BUFFER_SIZE_PROP, cleanerConfig.ioBufferSize.toString)
       props.put(KafkaConfig.MessageMaxBytesProp, cleanerConfig.maxMessageSize.toString)
-      props.put(KafkaConfig.LogCleanerBackoffMsProp, cleanerConfig.backoffMs.toString)
-      props.put(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, cleanerConfig.maxIoBytesPerSecond.toString)
+      props.put(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP, cleanerConfig.backoffMs.toString)
+      props.put(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP, cleanerConfig.maxIoBytesPerSecond.toString)
       KafkaConfig.fromProps(props)
     }
 
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
index 161493457e1f3..b3a7259219620 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
@@ -1959,7 +1959,7 @@ class LogCleanerTest extends Logging {
   @Test
   def testReconfigureLogCleanerIoMaxBytesPerSecond(): Unit = {
     val oldKafkaProps = TestUtils.createBrokerConfig(1, "localhost:2181")
-    oldKafkaProps.setProperty(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, "10000000")
+    oldKafkaProps.setProperty(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP, "10000000")
 
     val logCleaner = new LogCleaner(LogCleaner.cleanerConfig(new KafkaConfig(oldKafkaProps)),
       logDirs = Array(TestUtils.tempDir()),
@@ -1973,14 +1973,14 @@ class LogCleanerTest extends Logging {
     }
 
     try {
-      assertEquals(10000000, logCleaner.throttler.desiredRatePerSec, s"Throttler.desiredRatePerSec should be initialized from initial `${KafkaConfig.LogCleanerIoMaxBytesPerSecondProp}` config.")
+      assertEquals(10000000, logCleaner.throttler.desiredRatePerSec, s"Throttler.desiredRatePerSec should be initialized from initial `${CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP}` config.")
 
       val newKafkaProps = TestUtils.createBrokerConfig(1, "localhost:2181")
-      newKafkaProps.setProperty(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, "20000000")
+      newKafkaProps.setProperty(CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP, "20000000")
 
       logCleaner.reconfigure(new KafkaConfig(oldKafkaProps), new KafkaConfig(newKafkaProps))
 
-      assertEquals(20000000, logCleaner.throttler.desiredRatePerSec, s"Throttler.desiredRatePerSec should be updated with new `${KafkaConfig.LogCleanerIoMaxBytesPerSecondProp}` config.")
+      assertEquals(20000000, logCleaner.throttler.desiredRatePerSec, s"Throttler.desiredRatePerSec should be updated with new `${CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP}` config.")
     } finally {
       logCleaner.shutdown()
     }
diff --git a/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala b/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
index 00777e9677bb2..135de6e51441f 100644
--- a/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ControllerApisTest.scala
@@ -55,6 +55,7 @@ import org.apache.kafka.image.publisher.ControllerRegistrationsPublisher
 import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult, Authorizer}
 import org.apache.kafka.server.common.{ApiMessageAndVersion, Features, MetadataVersion, ProducerIdsBlock}
 import org.apache.kafka.server.util.FutureUtils
+import org.apache.kafka.storage.internals.log.CleanerConfig
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, Test}
 import org.junit.jupiter.params.ParameterizedTest
@@ -318,19 +319,19 @@ class ControllerApisTest {
           setResourceName("1").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new OldAlterableConfigCollection(util.Arrays.asList(new OldAlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000")).iterator())),
         new OldAlterConfigsResource().
           setResourceName("2").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new OldAlterableConfigCollection(util.Arrays.asList(new OldAlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000")).iterator())),
         new OldAlterConfigsResource().
           setResourceName("2").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new OldAlterableConfigCollection(util.Arrays.asList(new OldAlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000")).iterator())),
         new OldAlterConfigsResource().
           setResourceName("baz").
@@ -472,7 +473,7 @@ class ControllerApisTest {
           setResourceName("1").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new AlterableConfigCollection(util.Arrays.asList(new AlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000").
             setConfigOperation(AlterConfigOp.OpType.SET.id())).iterator())),
         new AlterConfigsResource().
@@ -536,14 +537,14 @@ class ControllerApisTest {
           setResourceName("3").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new AlterableConfigCollection(util.Arrays.asList(new AlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000").
             setConfigOperation(AlterConfigOp.OpType.SET.id())).iterator())),
         new AlterConfigsResource().
           setResourceName("3").
           setResourceType(ConfigResource.Type.BROKER.id()).
           setConfigs(new AlterableConfigCollection(util.Arrays.asList(new AlterableConfig().
-            setName(KafkaConfig.LogCleanerBackoffMsProp).
+            setName(CleanerConfig.LOG_CLEANER_BACKOFF_MS_PROP).
             setValue("100000").
             setConfigOperation(AlterConfigOp.OpType.SET.id())).iterator())),
         new AlterConfigsResource().
diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
index be065f5b8caef..4e1c8eed2769d 100755
--- a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
@@ -38,7 +38,7 @@ import org.apache.kafka.server.config.Defaults
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.kafka.server.util.KafkaScheduler
-import org.apache.kafka.storage.internals.log.{LogConfig, ProducerStateManagerConfig}
+import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig, ProducerStateManagerConfig}
 import org.apache.kafka.test.MockMetricsReporter
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
@@ -215,7 +215,7 @@ class DynamicBrokerConfigTest {
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, nonDynamicProps)
 
     // Test update of configs with invalid type
-    val invalidProps = Map(KafkaConfig.LogCleanerThreadsProp -> "invalid")
+    val invalidProps = Map(CleanerConfig.LOG_CLEANER_THREADS_PROP -> "invalid")
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, invalidProps)
 
     val excludedTopicConfig = Map(KafkaConfig.LogMessageFormatVersionProp -> "0.10.2")
@@ -225,21 +225,21 @@ class DynamicBrokerConfigTest {
   @Test
   def testConfigUpdateWithReconfigurableValidationFailure(): Unit = {
     val origProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
-    origProps.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, "100000000")
+    origProps.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, "100000000")
     val config = KafkaConfig(origProps)
     config.dynamicConfig.initialize(None, None)
 
     val validProps = Map.empty[String, String]
-    val invalidProps = Map(KafkaConfig.LogCleanerThreadsProp -> "20")
+    val invalidProps = Map(CleanerConfig.LOG_CLEANER_THREADS_PROP -> "20")
 
     def validateLogCleanerConfig(configs: util.Map[String, _]): Unit = {
-      val cleanerThreads = configs.get(KafkaConfig.LogCleanerThreadsProp).toString.toInt
+      val cleanerThreads = configs.get(CleanerConfig.LOG_CLEANER_THREADS_PROP).toString.toInt
       if (cleanerThreads <=0 || cleanerThreads >= 5)
         throw new ConfigException(s"Invalid cleaner threads $cleanerThreads")
     }
     val reconfigurable = new Reconfigurable {
       override def configure(configs: util.Map[String, _]): Unit = {}
-      override def reconfigurableConfigs(): util.Set[String] = Set(KafkaConfig.LogCleanerThreadsProp).asJava
+      override def reconfigurableConfigs(): util.Set[String] = Set(CleanerConfig.LOG_CLEANER_THREADS_PROP).asJava
       override def validateReconfiguration(configs: util.Map[String, _]): Unit = validateLogCleanerConfig(configs)
       override def reconfigure(configs: util.Map[String, _]): Unit = {}
     }
@@ -248,7 +248,7 @@ class DynamicBrokerConfigTest {
     config.dynamicConfig.removeReconfigurable(reconfigurable)
 
     val brokerReconfigurable = new BrokerReconfigurable {
-      override def reconfigurableConfigs: collection.Set[String] = Set(KafkaConfig.LogCleanerThreadsProp)
+      override def reconfigurableConfigs: collection.Set[String] = Set(CleanerConfig.LOG_CLEANER_THREADS_PROP)
       override def validateReconfiguration(newConfig: KafkaConfig): Unit = validateLogCleanerConfig(newConfig.originals)
       override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {}
     }
@@ -260,8 +260,8 @@ class DynamicBrokerConfigTest {
   def testReconfigurableValidation(): Unit = {
     val origProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     val config = KafkaConfig(origProps)
-    val invalidReconfigurableProps = Set(KafkaConfig.LogCleanerThreadsProp, KafkaConfig.BrokerIdProp, "some.prop")
-    val validReconfigurableProps = Set(KafkaConfig.LogCleanerThreadsProp, KafkaConfig.LogCleanerDedupeBufferSizeProp, "some.prop")
+    val invalidReconfigurableProps = Set(CleanerConfig.LOG_CLEANER_THREADS_PROP, KafkaConfig.BrokerIdProp, "some.prop")
+    val validReconfigurableProps = Set(CleanerConfig.LOG_CLEANER_THREADS_PROP, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, "some.prop")
 
     def createReconfigurable(configs: Set[String]) = new Reconfigurable {
       override def configure(configs: util.Map[String, _]): Unit = {}
diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
index e9426245b143f..a5d4d961fe123 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
@@ -40,7 +40,7 @@ import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.common.MetadataVersion.{IBP_0_8_2, IBP_3_0_IV1}
 import org.apache.kafka.server.config.ServerTopicConfigSynonyms
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
-import org.apache.kafka.storage.internals.log.LogConfig
+import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig}
 import org.junit.jupiter.api.function.Executable
 
 import scala.annotation.nowarn
@@ -845,14 +845,14 @@ class KafkaConfigTest {
         case KafkaConfig.LogRetentionBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.LogCleanupIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
         case KafkaConfig.LogCleanupPolicyProp => assertPropertyInvalid(baseProperties, name, "unknown_policy", "0")
-        case KafkaConfig.LogCleanerIoMaxBytesPerSecondProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.LogCleanerDedupeBufferSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "1024")
-        case KafkaConfig.LogCleanerDedupeBufferLoadFactorProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.LogCleanerEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case KafkaConfig.LogCleanerDeleteRetentionMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.LogCleanerMinCompactionLagMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.LogCleanerMaxCompactionLagMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.LogCleanerMinCleanRatioProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number", "1024")
+        case CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case CleanerConfig.LOG_CLEANER_DELETE_RETENTION_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case CleanerConfig.LOG_CLEANER_MIN_CLEAN_RATIO_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.LogIndexSizeMaxBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "3")
         case KafkaConfig.LogFlushIntervalMessagesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
         case KafkaConfig.LogFlushSchedulerIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 7b1ce80cbed3b..17ee3139a51ff 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -373,7 +373,7 @@ object TestUtils extends Logging {
     props.put(KafkaConfig.DeleteTopicEnableProp, enableDeleteTopic.toString)
     props.put(KafkaConfig.LogDeleteDelayMsProp, "1000")
     props.put(KafkaConfig.ControlledShutdownRetryBackoffMsProp, "100")
-    props.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, "2097152")
+    props.put(CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, "2097152")
     props.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
     if (!props.containsKey(KafkaConfig.OffsetsTopicPartitionsProp))
       props.put(KafkaConfig.OffsetsTopicPartitionsProp, "5")
diff --git a/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java b/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java
index 320de2db6b9c1..02cdc1cc875f8 100644
--- a/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java
+++ b/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java
@@ -30,7 +30,7 @@
 
 public final class ServerTopicConfigSynonyms {
     private static final String LOG_PREFIX = "log.";
-    private static final String LOG_CLEANER_PREFIX = LOG_PREFIX + "cleaner.";
+    public static final String LOG_CLEANER_PREFIX = LOG_PREFIX + "cleaner.";
 
     /**
      * Maps topic configurations to their equivalent broker configurations.
diff --git a/server/src/main/java/org/apache/kafka/server/config/Defaults.java b/server/src/main/java/org/apache/kafka/server/config/Defaults.java
index e55f641790f6b..0c7ca123b87c1 100644
--- a/server/src/main/java/org/apache/kafka/server/config/Defaults.java
+++ b/server/src/main/java/org/apache/kafka/server/config/Defaults.java
@@ -106,13 +106,6 @@ public class Defaults {
     public static final int NUM_PARTITIONS = 1;
     public static final String LOG_DIR = "/tmp/kafka-logs";
     public static final long LOG_CLEANUP_INTERVAL_MS = 5 * 60 * 1000L;
-    public static final int LOG_CLEANER_THREADS = 1;
-    public static final double LOG_CLEANER_IO_MAX_BYTES_PER_SECOND = Double.MAX_VALUE;
-    public static final long LOG_CLEANER_DEDUPE_BUFFER_SIZE = 128 * 1024 * 1024L;
-    public static final int LOG_CLEANER_IO_BUFFER_SIZE = 512 * 1024;
-    public static final double LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR = 0.9d;
-    public static final int LOG_CLEANER_BACKOFF_MS = 15 * 1000;
-    public static final boolean LOG_CLEANER_ENABLE = true;
     public static final int LOG_FLUSH_OFFSET_CHECKPOINT_INTERVAL_MS = 60000;
     public static final int LOG_FLUSH_START_OFFSET_CHECKPOINT_INTERVAL_MS = 60000;
     public static final int NUM_RECOVERY_THREADS_PER_DATA_DIR = 1;
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/CleanerConfig.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/CleanerConfig.java
index 9e38d0c02e96b..8168197fe08ab 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/CleanerConfig.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/CleanerConfig.java
@@ -16,11 +16,53 @@
  */
 package org.apache.kafka.storage.internals.log;
 
+import org.apache.kafka.common.config.TopicConfig;
+import org.apache.kafka.server.config.ServerTopicConfigSynonyms;
+
 /**
  * Configuration parameters for the log cleaner.
  */
 public class CleanerConfig {
     public static final String HASH_ALGORITHM = "MD5";
+    public static final int LOG_CLEANER_THREADS = 1;
+    public static final double LOG_CLEANER_IO_MAX_BYTES_PER_SECOND = Double.MAX_VALUE;
+    public static final long LOG_CLEANER_DEDUPE_BUFFER_SIZE = 128 * 1024 * 1024L;
+    public static final int LOG_CLEANER_IO_BUFFER_SIZE = 512 * 1024;
+    public static final double LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR = 0.9d;
+    public static final int LOG_CLEANER_BACKOFF_MS = 15 * 1000;
+    public static final boolean LOG_CLEANER_ENABLE = true;
+
+    public static final String LOG_CLEANER_THREADS_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "threads";
+    public static final String LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "io.max.bytes.per.second";
+    public static final String LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "dedupe.buffer.size";
+    public static final String LOG_CLEANER_IO_BUFFER_SIZE_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "io.buffer.size";
+    public static final String LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "io.buffer.load.factor";
+    public static final String LOG_CLEANER_BACKOFF_MS_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "backoff.ms";
+    public static final String LOG_CLEANER_MIN_CLEAN_RATIO_PROP = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MIN_CLEANABLE_DIRTY_RATIO_CONFIG);
+    public static final String LOG_CLEANER_ENABLE_PROP = ServerTopicConfigSynonyms.LOG_CLEANER_PREFIX + "enable";
+    public static final String LOG_CLEANER_DELETE_RETENTION_MS_PROP = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.DELETE_RETENTION_MS_CONFIG);
+    public static final String LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MIN_COMPACTION_LAG_MS_CONFIG);
+    public static final String LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.MAX_COMPACTION_LAG_MS_CONFIG);
+
+    public static final String LOG_CLEANER_MIN_CLEAN_RATIO_DOC = "The minimum ratio of dirty log to total log for a log to eligible for cleaning. " +
+            "If the " + LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP + " or the " + LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP +
+            " configurations are also specified, then the log compactor considers the log eligible for compaction " +
+            "as soon as either: (i) the dirty ratio threshold has been met and the log has had dirty (uncompacted) " +
+            "records for at least the " + LOG_CLEANER_MIN_COMPACTION_LAG_MS_PROP + " duration, or (ii) if the log has had " +
+            "dirty (uncompacted) records for at most the " + LOG_CLEANER_MAX_COMPACTION_LAG_MS_PROP + " period.";
+    public static final String LOG_CLEANER_THREADS_DOC = "The number of background threads to use for log cleaning";
+    public static final String LOG_CLEANER_IO_MAX_BYTES_PER_SECOND_DOC = "The log cleaner will be throttled so that the sum of its read and write i/o will be less than this value on average";
+    public static final String LOG_CLEANER_DEDUPE_BUFFER_SIZE_DOC = "The total memory used for log deduplication across all cleaner threads";
+    public static final String LOG_CLEANER_IO_BUFFER_SIZE_DOC = "The total memory used for log cleaner I/O buffers across all cleaner threads";
+    public static final String LOG_CLEANER_DEDUPE_BUFFER_LOAD_FACTOR_DOC = "Log cleaner dedupe buffer load factor. The percentage full the dedupe buffer can become. A higher value " +
+            "will allow more log to be cleaned at once but will lead to more hash collisions";
+    public static final String LOG_CLEANER_BACKOFF_MS_DOC = "The amount of time to sleep when there are no logs to clean";
+    public static final String LOG_CLEANER_ENABLE_DOC = "Enable the log cleaner process to run on the server. Should be enabled if using any topics with a cleanup.policy=compact including the internal offsets topic. If disabled those topics will not be compacted and continually grow in size.";
+    public static final String LOG_CLEANER_DELETE_RETENTION_MS_DOC = "The amount of time to retain tombstone message markers for log compacted topics. This setting also gives a bound " +
+            "on the time in which a consumer must complete a read if they begin from offset 0 to ensure that they get a valid snapshot of the final stage (otherwise  " +
+            "tombstones messages may be collected before a consumer completes their scan).";
+    public static final String LOG_CLEANER_MIN_COMPACTION_LAG_MS_DOC = "The minimum time a message will remain uncompacted in the log. Only applicable for logs that are being compacted.";
+    public static final String LOG_CLEANER_MAX_COMPACTION_LAG_MS_DOC = "The maximum time a message will remain ineligible for compaction in the log. Only applicable for logs that are being compacted.";
 
     public final int numThreads;
     public final long dedupeBufferSize;
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
index 3f3e9a243bd6c..4232e1d74c978 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
@@ -24,6 +24,7 @@
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.server.config.ConfigType;
 import org.apache.kafka.server.util.MockTime;
+import org.apache.kafka.storage.internals.log.CleanerConfig;
 import org.apache.kafka.test.TestCondition;
 import org.apache.kafka.test.TestUtils;
 import org.slf4j.Logger;
@@ -111,7 +112,7 @@ public void start() throws IOException {
         brokerConfig.put(KafkaConfig.ZkConnectProp(), zKConnectString());
         putIfAbsent(brokerConfig, KafkaConfig.ListenersProp(), "PLAINTEXT://localhost:" + DEFAULT_BROKER_PORT);
         putIfAbsent(brokerConfig, KafkaConfig.DeleteTopicEnableProp(), true);
-        putIfAbsent(brokerConfig, KafkaConfig.LogCleanerDedupeBufferSizeProp(), 2 * 1024 * 1024L);
+        putIfAbsent(brokerConfig, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, 2 * 1024 * 1024L);
         putIfAbsent(brokerConfig, KafkaConfig.GroupMinSessionTimeoutMsProp(), 0);
         putIfAbsent(brokerConfig, KafkaConfig.GroupInitialRebalanceDelayMsProp(), 0);
         putIfAbsent(brokerConfig, KafkaConfig.OffsetsTopicReplicationFactorProp(), (short) 1);

From 2f401ff4c85f6797391b8a3dd57d651f4de3d6ad Mon Sep 17 00:00:00 2001
From: Dongnuo Lyu <139248811+dongnuo123@users.noreply.github.com>
Date: Tue, 5 Mar 2024 09:50:58 -0500
Subject: [PATCH 109/258] MINOR: parameterize group-id in
 GroupMetadataManagerTestContext (#15467)

This pr parameterize some group ids in GroupMetadataManagerTestContext that are now constant strings.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../group/GroupMetadataManagerTestContext.java       | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
index d2d00c1582405..ad70107506ccb 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -643,7 +643,7 @@ public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteRebalance
 
         JoinGroupResponseData leaderJoinResponse =
             joinClassicGroupAsDynamicMemberAndCompleteJoin(new JoinGroupRequestBuilder()
-                .withGroupId("group-id")
+                .withGroupId(groupId)
                 .withMemberId(UNKNOWN_MEMBER_ID)
                 .withDefaultProtocolTypeAndProtocols()
                 .withRebalanceTimeoutMs(10000)
@@ -654,7 +654,7 @@ public JoinGroupResponseData joinClassicGroupAsDynamicMemberAndCompleteRebalance
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
         SyncResult syncResult = sendClassicGroupSync(new SyncGroupRequestBuilder()
-            .withGroupId("group-id")
+            .withGroupId(groupId)
             .withMemberId(leaderJoinResponse.memberId())
             .withGenerationId(leaderJoinResponse.generationId())
             .build());
@@ -804,7 +804,7 @@ public RebalanceResult staticMembersJoinAndRebalance(
         int rebalanceTimeoutMs,
         int sessionTimeoutMs
     ) throws Exception {
-        ClassicGroup group = createClassicGroup("group-id");
+        ClassicGroup group = createClassicGroup(groupId);
 
         JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
             .withGroupId(groupId)
@@ -901,7 +901,7 @@ public RebalanceResult staticMembersJoinAndRebalance(
     public PendingMemberGroupResult setupGroupWithPendingMember(ClassicGroup group) throws Exception {
         // Add the first member
         JoinGroupRequestData joinRequest = new JoinGroupRequestBuilder()
-            .withGroupId("group-id")
+            .withGroupId(group.groupId())
             .withMemberId(UNKNOWN_MEMBER_ID)
             .withDefaultProtocolTypeAndProtocols()
             .withRebalanceTimeoutMs(10000)
@@ -914,7 +914,7 @@ public PendingMemberGroupResult setupGroupWithPendingMember(ClassicGroup group)
         List<SyncGroupRequestData.SyncGroupRequestAssignment> assignment = new ArrayList<>();
         assignment.add(new SyncGroupRequestData.SyncGroupRequestAssignment().setMemberId(leaderJoinResponse.memberId()));
         SyncGroupRequestData syncRequest = new SyncGroupRequestBuilder()
-            .withGroupId("group-id")
+            .withGroupId(group.groupId())
             .withMemberId(leaderJoinResponse.memberId())
             .withGenerationId(leaderJoinResponse.generationId())
             .withAssignment(assignment)
@@ -1191,7 +1191,7 @@ public void verifyDescribeGroupsReturnsDeadGroup(String groupId) {
 
         assertEquals(
             Collections.singletonList(new DescribeGroupsResponseData.DescribedGroup()
-                .setGroupId("group-id")
+                .setGroupId(groupId)
                 .setGroupState(DEAD.toString())
             ),
             describedGroups

From 2c1943d836d828755af382b9d7996ff092854fe2 Mon Sep 17 00:00:00 2001
From: Colin Patrick McCabe <cmccabe@apache.org>
Date: Tue, 5 Mar 2024 12:02:19 -0800
Subject: [PATCH 110/258] MINOR: remove test constructor for
 PartitionAssignment (#15435)

Remove the test constructor for PartitionAssignment and remove the TODO.
Also add KRaftClusterTest.testCreatePartitions to get more coverage for
createPartitions.

Reviewers: David Arthur <mumrah@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/server/KRaftClusterTest.scala       | 35 ++++++++++++++++++-
 .../placement/PartitionAssignment.java        |  6 ----
 .../PartitionChangeBuilderTest.java           |  6 ++--
 .../PartitionReassignmentReplicasTest.java    | 22 ++++++------
 .../ReplicationControlManagerTest.java        | 18 +++++-----
 .../placement/PartitionAssignmentTest.java    | 12 ++++---
 .../placement/StripedReplicaPlacerTest.java   | 19 +++++-----
 .../placement/TopicAssignmentTest.java        |  9 ++---
 8 files changed, 80 insertions(+), 47 deletions(-)

diff --git a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
index 130d0e5642ea4..3be1b400ab564 100644
--- a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
@@ -26,7 +26,7 @@ import org.apache.kafka.clients.admin._
 import org.apache.kafka.common.acl.{AclBinding, AclBindingFilter}
 import org.apache.kafka.common.config.{ConfigException, ConfigResource}
 import org.apache.kafka.common.config.ConfigResource.Type
-import org.apache.kafka.common.errors.{PolicyViolationException, UnsupportedVersionException}
+import org.apache.kafka.common.errors.{InvalidPartitionsException,PolicyViolationException, UnsupportedVersionException}
 import org.apache.kafka.common.message.DescribeClusterRequestData
 import org.apache.kafka.common.metadata.{ConfigRecord, FeatureLevelRecord}
 import org.apache.kafka.common.metrics.Metrics
@@ -792,6 +792,39 @@ class KRaftClusterTest {
     }
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = Array("3.7-IV0", "3.7-IV2"))
+  def testCreatePartitions(metadataVersionString: String): Unit = {
+    val cluster = new KafkaClusterTestKit.Builder(
+      new TestKitNodes.Builder().
+        setNumBrokerNodes(4).
+        setBootstrapMetadataVersion(MetadataVersion.fromVersionString(metadataVersionString)).
+        setNumControllerNodes(3).build()).
+      build()
+    try {
+      cluster.format()
+      cluster.startup()
+      cluster.waitForReadyBrokers()
+      val admin = Admin.create(cluster.clientProperties())
+      try {
+        val createResults = admin.createTopics(Arrays.asList(
+          new NewTopic("foo", 1, 3.toShort),
+          new NewTopic("bar", 2, 3.toShort))).values()
+        createResults.get("foo").get()
+        createResults.get("bar").get()
+        val increaseResults = admin.createPartitions(Map(
+          "foo" -> NewPartitions.increaseTo(3),
+          "bar" -> NewPartitions.increaseTo(2)).asJava).values()
+        increaseResults.get("foo").get()
+        assertEquals(classOf[InvalidPartitionsException], assertThrows(
+          classOf[ExecutionException], () => increaseResults.get("bar").get()).getCause.getClass)
+      } finally {
+        admin.close()
+      }
+    } finally {
+      cluster.close()
+    }
+  }
   private def clusterImage(
     cluster: KafkaClusterTestKit,
     brokerId: Int
diff --git a/metadata/src/main/java/org/apache/kafka/metadata/placement/PartitionAssignment.java b/metadata/src/main/java/org/apache/kafka/metadata/placement/PartitionAssignment.java
index 177d5311afdda..a7012d1505c03 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/placement/PartitionAssignment.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/placement/PartitionAssignment.java
@@ -17,7 +17,6 @@
 
 package org.apache.kafka.metadata.placement;
 
-import org.apache.kafka.common.DirectoryId;
 import org.apache.kafka.common.Uuid;
 
 import java.util.ArrayList;
@@ -39,11 +38,6 @@ public class PartitionAssignment {
     private final List<Integer> replicas;
     private final List<Uuid> directories;
 
-    // TODO remove -- just here for testing
-    public PartitionAssignment(List<Integer> replicas) {
-        this(replicas, brokerId -> DirectoryId.UNASSIGNED);
-    }
-
     public PartitionAssignment(List<Integer> replicas, DefaultDirProvider defaultDirProvider) {
         this.replicas = Collections.unmodifiableList(new ArrayList<>(replicas));
         Uuid[] directories = new Uuid[replicas.size()];
diff --git a/metadata/src/test/java/org/apache/kafka/controller/PartitionChangeBuilderTest.java b/metadata/src/test/java/org/apache/kafka/controller/PartitionChangeBuilderTest.java
index efc9bd2a24f0e..b956a7f2ef92d 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/PartitionChangeBuilderTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/PartitionChangeBuilderTest.java
@@ -27,7 +27,6 @@
 import org.apache.kafka.metadata.PartitionRegistration;
 import org.apache.kafka.metadata.Replicas;
 import org.apache.kafka.metadata.placement.DefaultDirProvider;
-import org.apache.kafka.metadata.placement.PartitionAssignment;
 import org.apache.kafka.server.common.ApiMessageAndVersion;
 import org.apache.kafka.server.common.MetadataVersion;
 import org.junit.jupiter.api.Test;
@@ -49,6 +48,7 @@
 import static org.apache.kafka.controller.PartitionChangeBuilder.changeRecordIsNoOp;
 import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER;
 import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER_CHANGE;
+import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -502,7 +502,7 @@ public void testRevertReassignment(short version) {
     @MethodSource("partitionChangeRecordVersions")
     public void testRemovingReplicaReassignment(short version) {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Replicas.toList(FOO.replicas)), new PartitionAssignment(Arrays.asList(1, 2)));
+            partitionAssignment(Replicas.toList(FOO.replicas)), partitionAssignment(Arrays.asList(1, 2)));
         assertEquals(Collections.singletonList(3), replicas.removing());
         assertEquals(Collections.emptyList(), replicas.adding());
         assertEquals(Arrays.asList(1, 2, 3), replicas.replicas());
@@ -527,7 +527,7 @@ public void testRemovingReplicaReassignment(short version) {
     @MethodSource("partitionChangeRecordVersions")
     public void testAddingReplicaReassignment(short version) {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Replicas.toList(FOO.replicas)), new PartitionAssignment(Arrays.asList(1, 2, 3, 4)));
+            partitionAssignment(Replicas.toList(FOO.replicas)), partitionAssignment(Arrays.asList(1, 2, 3, 4)));
         assertEquals(Collections.emptyList(), replicas.removing());
         assertEquals(Collections.singletonList(4), replicas.adding());
         assertEquals(Arrays.asList(1, 2, 3, 4), replicas.replicas());
diff --git a/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java b/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
index b2bc540bda5f6..17be98d47f0b1 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/PartitionReassignmentReplicasTest.java
@@ -24,10 +24,10 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.metadata.LeaderRecoveryState;
 import org.apache.kafka.metadata.PartitionRegistration;
-import org.apache.kafka.metadata.placement.PartitionAssignment;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
+import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -38,7 +38,7 @@ public class PartitionReassignmentReplicasTest {
     @Test
     public void testNoneAddedOrRemoved() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(3, 2, 1)), new PartitionAssignment(Arrays.asList(3, 2, 1)));
+            partitionAssignment(Arrays.asList(3, 2, 1)), partitionAssignment(Arrays.asList(3, 2, 1)));
         assertEquals(Collections.emptyList(), replicas.removing());
         assertEquals(Collections.emptyList(), replicas.adding());
         assertEquals(Arrays.asList(3, 2, 1), replicas.replicas());
@@ -47,7 +47,7 @@ public void testNoneAddedOrRemoved() {
     @Test
     public void testAdditions() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(3, 2, 1)), new PartitionAssignment(Arrays.asList(3, 6, 2, 1, 5)));
+            partitionAssignment(Arrays.asList(3, 2, 1)), partitionAssignment(Arrays.asList(3, 6, 2, 1, 5)));
         assertEquals(Collections.emptyList(), replicas.removing());
         assertEquals(Arrays.asList(5, 6), replicas.adding());
         assertEquals(Arrays.asList(3, 6, 2, 1, 5), replicas.replicas());
@@ -56,7 +56,7 @@ public void testAdditions() {
     @Test
     public void testRemovals() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(3, 2, 1, 0)), new PartitionAssignment(Arrays.asList(3, 1)));
+            partitionAssignment(Arrays.asList(3, 2, 1, 0)), partitionAssignment(Arrays.asList(3, 1)));
         assertEquals(Arrays.asList(0, 2), replicas.removing());
         assertEquals(Collections.emptyList(), replicas.adding());
         assertEquals(Arrays.asList(3, 1, 0, 2), replicas.replicas());
@@ -65,7 +65,7 @@ public void testRemovals() {
     @Test
     public void testAdditionsAndRemovals() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(3, 2, 1, 0)), new PartitionAssignment(Arrays.asList(7, 3, 1, 9)));
+            partitionAssignment(Arrays.asList(3, 2, 1, 0)), partitionAssignment(Arrays.asList(7, 3, 1, 9)));
         assertEquals(Arrays.asList(0, 2), replicas.removing());
         assertEquals(Arrays.asList(7, 9), replicas.adding());
         assertEquals(Arrays.asList(7, 3, 1, 9, 0, 2), replicas.replicas());
@@ -74,7 +74,7 @@ public void testAdditionsAndRemovals() {
     @Test
     public void testRearrangement() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(3, 2, 1, 0)), new PartitionAssignment(Arrays.asList(0, 1, 3, 2)));
+            partitionAssignment(Arrays.asList(3, 2, 1, 0)), partitionAssignment(Arrays.asList(0, 1, 3, 2)));
         assertEquals(Collections.emptyList(), replicas.removing());
         assertEquals(Collections.emptyList(), replicas.adding());
         assertEquals(Arrays.asList(0, 1, 3, 2), replicas.replicas());
@@ -83,7 +83,7 @@ public void testRearrangement() {
     @Test
     public void testDoesNotCompleteReassignment() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(3, 4, 5)));
+            partitionAssignment(Arrays.asList(0, 1, 2)), partitionAssignment(Arrays.asList(3, 4, 5)));
         assertTrue(replicas.isReassignmentInProgress());
         Optional<PartitionReassignmentReplicas.CompletedReassignment> reassignmentOptional =
             replicas.maybeCompleteReassignment(Arrays.asList(0, 1, 2, 3, 4));
@@ -107,7 +107,7 @@ public void testDoesNotCompleteReassignmentIfNoneOngoing() {
     @Test
     public void testDoesCompleteReassignmentAllNewReplicas() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(3, 4, 5)));
+            partitionAssignment(Arrays.asList(0, 1, 2)), partitionAssignment(Arrays.asList(3, 4, 5)));
         assertTrue(replicas.isReassignmentInProgress());
         Optional<PartitionReassignmentReplicas.CompletedReassignment> reassignmentOptional =
             replicas.maybeCompleteReassignment(Arrays.asList(0, 1, 2, 3, 4, 5));
@@ -120,7 +120,7 @@ public void testDoesCompleteReassignmentAllNewReplicas() {
     @Test
     public void testDoesCompleteReassignmentSomeNewReplicas() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(0, 1, 3)));
+            partitionAssignment(Arrays.asList(0, 1, 2)), partitionAssignment(Arrays.asList(0, 1, 3)));
         assertTrue(replicas.isReassignmentInProgress());
         Optional<PartitionReassignmentReplicas.CompletedReassignment> reassignmentOptional =
             replicas.maybeCompleteReassignment(Arrays.asList(0, 1, 2, 3));
@@ -199,7 +199,7 @@ public void testIsReassignmentInProgress() {
     @Test
     public void testDoesNotCompleteReassignmentIfIsrDoesNotHaveAllTargetReplicas() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(0, 1, 3)));
+            partitionAssignment(Arrays.asList(0, 1, 2)), partitionAssignment(Arrays.asList(0, 1, 3)));
         assertTrue(replicas.isReassignmentInProgress());
         Optional<PartitionReassignmentReplicas.CompletedReassignment> reassignmentOptional =
             replicas.maybeCompleteReassignment(Arrays.asList(3));
@@ -209,7 +209,7 @@ public void testDoesNotCompleteReassignmentIfIsrDoesNotHaveAllTargetReplicas() {
     @Test
     public void testOriginalReplicas() {
         PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
-            new PartitionAssignment(Arrays.asList(0, 1, 2)), new PartitionAssignment(Arrays.asList(0, 1, 3)));
+            partitionAssignment(Arrays.asList(0, 1, 2)), partitionAssignment(Arrays.asList(0, 1, 3)));
         assertEquals(Arrays.asList(0, 1, 2), replicas.originalReplicas());
     }
 }
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
index 3d54720be921f..e9885fa8e94b5 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
@@ -87,7 +87,6 @@
 import org.apache.kafka.metadata.PartitionRegistration;
 import org.apache.kafka.metadata.RecordTestUtils;
 import org.apache.kafka.metadata.Replicas;
-import org.apache.kafka.metadata.placement.PartitionAssignment;
 import org.apache.kafka.metadata.placement.StripedReplicaPlacer;
 import org.apache.kafka.metadata.placement.UsableBroker;
 import org.apache.kafka.server.common.ApiMessageAndVersion;
@@ -148,6 +147,7 @@
 import static org.apache.kafka.controller.ControllerRequestContextUtil.anonymousContextFor;
 import static org.apache.kafka.controller.ControllerRequestContextUtil.anonymousContextWithMutationQuotaExceededFor;
 import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER;
+import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -1584,13 +1584,13 @@ public void testCreatePartitionsISRInvariants() throws Exception {
     public void testValidateGoodManualPartitionAssignments() throws Exception {
         ReplicationControlTestContext ctx = new ReplicationControlTestContext.Builder().build();
         ctx.registerBrokers(1, 2, 3);
-        ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1)),
+        ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1)),
             OptionalInt.of(1));
-        ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1)),
+        ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1)),
             OptionalInt.empty());
-        ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1, 2, 3)),
+        ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1, 2, 3)),
             OptionalInt.of(3));
-        ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1, 2, 3)),
+        ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1, 2, 3)),
             OptionalInt.empty());
     }
 
@@ -1600,20 +1600,20 @@ public void testValidateBadManualPartitionAssignments() throws Exception {
         ctx.registerBrokers(1, 2);
         assertEquals("The manual partition assignment includes an empty replica list.",
             assertThrows(InvalidReplicaAssignmentException.class, () ->
-                ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList()),
+                ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList()),
                     OptionalInt.empty())).getMessage());
         assertEquals("The manual partition assignment includes broker 3, but no such " +
             "broker is registered.", assertThrows(InvalidReplicaAssignmentException.class, () ->
-                ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1, 2, 3)),
+                ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1, 2, 3)),
                     OptionalInt.empty())).getMessage());
         assertEquals("The manual partition assignment includes the broker 2 more than " +
             "once.", assertThrows(InvalidReplicaAssignmentException.class, () ->
-                ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1, 2, 2)),
+                ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1, 2, 2)),
                     OptionalInt.empty())).getMessage());
         assertEquals("The manual partition assignment includes a partition with 2 " +
             "replica(s), but this is not consistent with previous partitions, which have " +
                 "3 replica(s).", assertThrows(InvalidReplicaAssignmentException.class, () ->
-                    ctx.replicationControl.validateManualPartitionAssignment(new PartitionAssignment(asList(1, 2)),
+                    ctx.replicationControl.validateManualPartitionAssignment(partitionAssignment(asList(1, 2)),
                         OptionalInt.of(3))).getMessage());
     }
 
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/PartitionAssignmentTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/PartitionAssignmentTest.java
index 06cf5ae50d5dc..6dca18b4dd575 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/placement/PartitionAssignmentTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/PartitionAssignmentTest.java
@@ -20,6 +20,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 
+import org.apache.kafka.common.DirectoryId;
 import org.apache.kafka.common.Uuid;
 import org.junit.jupiter.api.Test;
 
@@ -27,20 +28,23 @@
 import java.util.List;
 
 public class PartitionAssignmentTest {
+    public static PartitionAssignment partitionAssignment(List<Integer> replicas) {
+        return new PartitionAssignment(replicas, __ -> DirectoryId.MIGRATING);
+    }
 
     @Test
     public void testPartitionAssignmentReplicas() {
         List<Integer> replicas = Arrays.asList(0, 1, 2);
-        assertEquals(replicas, new PartitionAssignment(replicas).replicas());
+        assertEquals(replicas, partitionAssignment(replicas).replicas());
     }
 
     @Test
     public void testConsistentEqualsAndHashCode() {
         List<PartitionAssignment> partitionAssignments = Arrays.asList(
-            new PartitionAssignment(
+            partitionAssignment(
                 Arrays.asList(0, 1, 2)
             ),
-            new PartitionAssignment(
+            partitionAssignment(
                 Arrays.asList(1, 2, 0)
             )
         );
@@ -49,7 +53,7 @@ public void testConsistentEqualsAndHashCode() {
             for (int j = 0; j < partitionAssignments.size(); j++) {
                 if (i == j) {
                     assertEquals(partitionAssignments.get(i), partitionAssignments.get(j));
-                    assertEquals(partitionAssignments.get(i), new PartitionAssignment(partitionAssignments.get(i).replicas()));
+                    assertEquals(partitionAssignments.get(i), partitionAssignment(partitionAssignments.get(i).replicas()));
                     assertEquals(partitionAssignments.get(i).hashCode(), partitionAssignments.get(j).hashCode());
                 } else {
                     assertNotEquals(partitionAssignments.get(i), partitionAssignments.get(j));
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java
index 8b02416d2be61..924fcdb7559aa 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java
@@ -34,6 +34,7 @@
 import java.util.Map;
 import java.util.Optional;
 
+import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -105,7 +106,7 @@ public Iterator<UsableBroker> usableBrokers() {
 
             @Override
             public Uuid defaultDir(int brokerId) {
-                return DirectoryId.UNASSIGNED;
+                return DirectoryId.MIGRATING;
             }
         });
     }
@@ -118,9 +119,9 @@ public Uuid defaultDir(int brokerId) {
     public void testMultiPartitionTopicPlacementOnSingleUnfencedBroker() {
         MockRandom random = new MockRandom();
         StripedReplicaPlacer placer = new StripedReplicaPlacer(random);
-        assertEquals(new TopicAssignment(Arrays.asList(new PartitionAssignment(Arrays.asList(0)),
-                new PartitionAssignment(Arrays.asList(0)),
-                new PartitionAssignment(Arrays.asList(0)))),
+        assertEquals(new TopicAssignment(Arrays.asList(partitionAssignment(Arrays.asList(0)),
+                partitionAssignment(Arrays.asList(0)),
+                partitionAssignment(Arrays.asList(0)))),
                 place(placer, 0, 3, (short) 1, Arrays.asList(
                         new UsableBroker(0, Optional.empty(), false),
                         new UsableBroker(1, Optional.empty(), true))));
@@ -224,11 +225,11 @@ public void testNonPositiveReplicationFactor() {
     public void testSuccessfulPlacement() {
         MockRandom random = new MockRandom();
         StripedReplicaPlacer placer = new StripedReplicaPlacer(random);
-        assertEquals(new TopicAssignment(Arrays.asList(new PartitionAssignment(Arrays.asList(2, 3, 0)),
-                new PartitionAssignment(Arrays.asList(3, 0, 1)),
-                new PartitionAssignment(Arrays.asList(0, 1, 2)),
-                new PartitionAssignment(Arrays.asList(1, 2, 3)),
-                new PartitionAssignment(Arrays.asList(1, 0, 2)))),
+        assertEquals(new TopicAssignment(Arrays.asList(partitionAssignment(Arrays.asList(2, 3, 0)),
+                partitionAssignment(Arrays.asList(3, 0, 1)),
+                partitionAssignment(Arrays.asList(0, 1, 2)),
+                partitionAssignment(Arrays.asList(1, 2, 3)),
+                partitionAssignment(Arrays.asList(1, 0, 2)))),
             place(placer, 0, 5, (short) 3, Arrays.asList(
                 new UsableBroker(0, Optional.empty(), false),
                 new UsableBroker(3, Optional.empty(), false),
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/TopicAssignmentTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/TopicAssignmentTest.java
index 7b5a24c3b8409..26f8841d834f5 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/placement/TopicAssignmentTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/TopicAssignmentTest.java
@@ -17,6 +17,7 @@
 
 package org.apache.kafka.metadata.placement;
 
+import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 
@@ -33,8 +34,8 @@ public void testTopicAssignmentReplicas() {
         List<Integer> replicasP0 = Arrays.asList(0, 1, 2);
         List<Integer> replicasP1 = Arrays.asList(1, 2, 0);
         List<PartitionAssignment> partitionAssignments = Arrays.asList(
-            new PartitionAssignment(replicasP0),
-            new PartitionAssignment(replicasP1)
+            partitionAssignment(replicasP0),
+            partitionAssignment(replicasP1)
         );
         assertEquals(partitionAssignments, new TopicAssignment(partitionAssignments).assignments());
     }
@@ -44,14 +45,14 @@ public void testConsistentEqualsAndHashCode() {
         List<TopicAssignment> topicAssignments = Arrays.asList(
             new TopicAssignment(
                 Arrays.asList(
-                    new PartitionAssignment(
+                    partitionAssignment(
                         Arrays.asList(0, 1, 2)
                     )
                 )
             ),
             new TopicAssignment(
                 Arrays.asList(
-                    new PartitionAssignment(
+                    partitionAssignment(
                         Arrays.asList(1, 2, 0)
                     )
                 )

From 1ca939128521459ef921d3ec71dda2a507ba1f15 Mon Sep 17 00:00:00 2001
From: "Gyeongwon, Do" <dct012@gmail.com>
Date: Wed, 6 Mar 2024 06:16:35 +0900
Subject: [PATCH 111/258] MINOR: Remove controlPlaneRequestProcessor in
 BrokerServer (#15245)

It seems likely that BrokerServer was built upon the KafkaServer codebase.(#10113)
KafkaServer, using Zookeeper, separates controlPlane and dataPlane to implement KIP-291.
In KRaft, the roles of DataPlane and ControlPlane in KafkaServer seem to be divided into BrokerServer and ControllerServer.

It appears that the initial implementation of BrokerServer initialized and used the controlPlaneRequestProcessor, but it seems to have been removed, except for the code used in the shutdown method, through subsequent modifications.(#10931)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 core/src/main/scala/kafka/server/BrokerServer.scala | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/core/src/main/scala/kafka/server/BrokerServer.scala b/core/src/main/scala/kafka/server/BrokerServer.scala
index dc72ceb35ca7a..8ac568a5e3a7e 100644
--- a/core/src/main/scala/kafka/server/BrokerServer.scala
+++ b/core/src/main/scala/kafka/server/BrokerServer.scala
@@ -96,7 +96,6 @@ class BrokerServer(
   var status: ProcessStatus = SHUTDOWN
 
   @volatile var dataPlaneRequestProcessor: KafkaApis = _
-  var controlPlaneRequestProcessor: KafkaApis = _
 
   var authorizer: Option[Authorizer] = None
   @volatile var socketServer: SocketServer = _
@@ -655,8 +654,6 @@ class BrokerServer(
         CoreUtils.swallow(dataPlaneRequestHandlerPool.shutdown(), this)
       if (dataPlaneRequestProcessor != null)
         CoreUtils.swallow(dataPlaneRequestProcessor.close(), this)
-      if (controlPlaneRequestProcessor != null)
-        CoreUtils.swallow(controlPlaneRequestProcessor.close(), this)
       CoreUtils.swallow(authorizer.foreach(_.close()), this)
 
       /**

From 554fa57af85ec337a556f35cbc6d2970ff252dc4 Mon Sep 17 00:00:00 2001
From: John Yu <54207775+chiacyu@users.noreply.github.com>
Date: Wed, 6 Mar 2024 09:00:58 +0800
Subject: [PATCH 112/258] KAFKA-16209 : fetchSnapshot might return null if
 topic is created before v2.8 (#15444)

Change the function with a better way to deal with the NULL pointer exception.

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../scala/unit/kafka/log/ProducerStateManagerTest.scala     | 6 ++++++
 .../kafka/storage/internals/log/ProducerStateManager.java   | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
index 03aa847ded73d..810e0b1e4ee08 100644
--- a/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
@@ -579,6 +579,12 @@ class ProducerStateManagerTest {
     assertTrue(logDir.list().head.nonEmpty, "Snapshot file is empty")
   }
 
+  @Test
+  def testFetchSnapshotEmptySnapShot(): Unit = {
+    val offset = 1
+    assertEquals(Optional.empty(), stateManager.fetchSnapshot(offset))
+  }
+
   @Test
   def testRecoverFromSnapshotUnfinishedTransaction(): Unit = {
     val epoch = 0.toShort
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
index 270aa0a42f94c..da56ddaccc864 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/ProducerStateManager.java
@@ -605,7 +605,7 @@ public void deleteSnapshotsBefore(long offset) throws IOException {
     }
 
     public Optional<File> fetchSnapshot(long offset) {
-        return Optional.of(snapshots.get(offset)).map(x -> x.file());
+        return Optional.ofNullable(snapshots.get(offset)).map(x -> x.file());
     }
 
     private Optional<SnapshotFile> oldestSnapshotFile() {

From e81379d3fea956dd8900b7f4b68e0c1328401871 Mon Sep 17 00:00:00 2001
From: Victor van den Hoven <victor.vanden.hoven@alliander.com>
Date: Wed, 6 Mar 2024 02:06:20 +0100
Subject: [PATCH 113/258] KAFKA-15417: flip joinSpuriousLookBackTimeMs and emit
 non-joined items (#14426)

Kafka Streams support asymmetric join windows. Depending on the window configuration
we need to compute window close time etc differently.

This PR flips `joinSpuriousLookBackTimeMs`, because they were not correct, and
introduced the `windowsAfterIntervalMs`-field that is used to find if emitting records can be skipped.

Reviewers: Hao Li <hli@confluent.io>, Guozhang Wang <guozhang.wang.us@gmail.com>, Matthias J. Sax <matthias@confluent.io>
---
 .../kstream/internals/KStreamKStreamJoin.java |  38 +++-
 .../KStreamKStreamIntegrationTest.java        |   2 +
 .../internals/KStreamKStreamLeftJoinTest.java | 183 +++++++++++++++++-
 .../KStreamKStreamOuterJoinTest.java          |  41 ++--
 4 files changed, 236 insertions(+), 28 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
index 603e1e8255091..124386b9bc3ae 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
@@ -51,7 +51,8 @@ class KStreamKStreamJoin<K, V1, V2, VOut> implements ProcessorSupplier<K, V1, K,
     private final long joinAfterMs;
     private final long joinGraceMs;
     private final boolean enableSpuriousResultFix;
-    private final long joinSpuriousLookBackTimeMs;
+    private final long windowsBeforeMs;
+    private final long windowsAfterMs;
 
     private final boolean outer;
     private final boolean isLeftSide;
@@ -72,12 +73,12 @@ class KStreamKStreamJoin<K, V1, V2, VOut> implements ProcessorSupplier<K, V1, K,
         if (isLeftSide) {
             this.joinBeforeMs = windows.beforeMs;
             this.joinAfterMs = windows.afterMs;
-            this.joinSpuriousLookBackTimeMs = windows.beforeMs;
         } else {
             this.joinBeforeMs = windows.afterMs;
             this.joinAfterMs = windows.beforeMs;
-            this.joinSpuriousLookBackTimeMs = windows.afterMs;
         }
+        this.windowsAfterMs = windows.afterMs;
+        this.windowsBeforeMs = windows.beforeMs;
         this.joinGraceMs = windows.gracePeriodMs();
         this.enableSpuriousResultFix = windows.spuriousResultFixEnabled();
         this.joiner = joiner;
@@ -136,11 +137,12 @@ public void process(final Record<K, V1> record) {
                 return;
             }
 
-            boolean needOuterJoin = outer;
             // Emit all non-joined records which window has closed
             if (inputRecordTimestamp == sharedTimeTracker.streamTime) {
                 outerJoinStore.ifPresent(store -> emitNonJoinedOuterRecords(store, record));
             }
+
+            boolean needOuterJoin = outer;
             try (final WindowStoreIterator<V2> iter = otherWindowStore.fetch(record.key(), timeFrom, timeTo)) {
                 while (iter.hasNext()) {
                     needOuterJoin = false;
@@ -200,7 +202,7 @@ private void emitNonJoinedOuterRecords(
             // to reduce runtime cost, we try to avoid paying those cost
 
             // only try to emit left/outer join results if there _might_ be any result records
-            if (sharedTimeTracker.minTime >= sharedTimeTracker.streamTime - joinSpuriousLookBackTimeMs - joinGraceMs) {
+            if (sharedTimeTracker.minTime + joinAfterMs + joinGraceMs >= sharedTimeTracker.streamTime) {
                 return;
             }
             // throttle the emit frequency to a (configurable) interval;
@@ -222,6 +224,8 @@ private void emitNonJoinedOuterRecords(
                 TimestampedKeyAndJoinSide<K> prevKey = null;
 
                 while (it.hasNext()) {
+                    boolean outerJoinLeftBreak = false;
+                    boolean outerJoinRightBreak = false;
                     final KeyValue<TimestampedKeyAndJoinSide<K>, LeftOrRightValue<V1, V2>> next = it.next();
                     final TimestampedKeyAndJoinSide<K> timestampedKeyAndJoinSide = next.key;
                     final LeftOrRightValue<V1, V2> value = next.value;
@@ -230,8 +234,19 @@ private void emitNonJoinedOuterRecords(
                     sharedTimeTracker.minTime = timestamp;
 
                     // Skip next records if window has not closed
-                    if (timestamp + joinAfterMs + joinGraceMs >= sharedTimeTracker.streamTime) {
-                        break;
+                    final long outerJoinLookBackTimeMs = getOuterJoinLookBackTimeMs(timestampedKeyAndJoinSide);
+                    if (sharedTimeTracker.minTime + outerJoinLookBackTimeMs + joinGraceMs >= sharedTimeTracker.streamTime) {
+                        if (timestampedKeyAndJoinSide.isLeftSide()) {
+                            outerJoinLeftBreak = true; // there are no more candidates to emit on left-outerJoin-side
+                        } else {
+                            outerJoinRightBreak = true; // there are no more candidates to emit on right-outerJoin-side
+                        }
+                        if (outerJoinLeftBreak && outerJoinRightBreak) {
+                            break; // there are no more candidates to emit on left-outerJoin-side and
+                                    // right-outerJoin-side
+                        } else {
+                            continue; // there are possibly candidates left on the other outerJoin-side
+                        }
                     }
 
                     final VOut nullJoinedValue;
@@ -268,6 +283,15 @@ private void emitNonJoinedOuterRecords(
             }
         }
 
+        private long getOuterJoinLookBackTimeMs(final TimestampedKeyAndJoinSide<K> timestampedKeyAndJoinSide) {
+            // depending on the JoinSide we fill in the outerJoinLookBackTimeMs
+            if (timestampedKeyAndJoinSide.isLeftSide()) {
+                return windowsAfterMs; // On the left-JoinSide we look back in time
+            } else {
+                return windowsBeforeMs; // On the right-JoinSide we look forward in time
+            }
+        }
+
         @Override
         public void close() {
             sharedTimeTrackerSupplier.remove(context().taskId());
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamKStreamIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamKStreamIntegrationTest.java
index 1d9a77b5bf495..10ab37cee0714 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamKStreamIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamKStreamIntegrationTest.java
@@ -27,6 +27,7 @@
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.StreamsConfig.InternalConfig;
 import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
 import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
 import org.apache.kafka.streams.kstream.JoinWindows;
@@ -99,6 +100,7 @@ public void before(final TestInfo testInfo) throws IOException {
         final String safeTestName = safeUniqueTestName(testInfo);
         streamsConfig = getStreamsConfig(safeTestName);
         streamsConfig.put(StreamsConfig.STATE_DIR_CONFIG, stateDirBasePath);
+        streamsConfig.put(InternalConfig.EMIT_INTERVAL_MS_KSTREAMS_OUTER_JOIN_SPURIOUS_RESULTS_FIX, 0L);
     }
 
     @AfterEach
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamLeftJoinTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamLeftJoinTest.java
index 156b553455d47..fd36b241b27b0 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamLeftJoinTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamLeftJoinTest.java
@@ -436,6 +436,184 @@ public void testRightNonJoinedRecordsAreNeverEmittedByTheRightProcessor() {
         }
     }
 
+    @Test
+    public void testLeftJoinedRecordsWithZeroAfterAreEmitted() {
+        final StreamsBuilder builder = new StreamsBuilder();
+
+        final int[] expectedKeys = new int[] {0, 1, 2, 3};
+
+        final KStream<Integer, String> stream1;
+        final KStream<Integer, String> stream2;
+        final KStream<Integer, String> joined;
+        final MockApiProcessorSupplier<Integer, String, Void, Void> supplier = new MockApiProcessorSupplier<>();
+        stream1 = builder.stream(topic1, consumed);
+        stream2 = builder.stream(topic2, consumed);
+        
+        joined = stream1.leftJoin(
+            stream2,
+            MockValueJoiner.TOSTRING_JOINER,
+            JoinWindows.ofTimeDifferenceWithNoGrace(ofMillis(100)).after(ZERO),
+            StreamJoined.with(Serdes.Integer(),
+                Serdes.String(),
+                Serdes.String())
+        );
+        joined.process(supplier);
+
+        final Collection<Set<String>> copartitionGroups =
+            TopologyWrapper.getInternalTopologyBuilder(builder.build()).copartitionGroups();
+
+        assertEquals(1, copartitionGroups.size());
+        assertEquals(new HashSet<>(Arrays.asList(topic1, topic2)), copartitionGroups.iterator().next());
+
+        try (final TopologyTestDriver driver = new TopologyTestDriver(builder.build(), PROPS)) {
+            final TestInputTopic<Integer, String> inputTopic1 =
+                    driver.createInputTopic(topic1, new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+            final TestInputTopic<Integer, String> inputTopic2 =
+                    driver.createInputTopic(topic2, new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+            final MockApiProcessor<Integer, String, Void, Void> processor = supplier.theCapturedProcessor();
+
+            processor.init(null);
+            
+            // push four items with increasing timestamps to the primary stream; the other window is empty; 
+            // this should emit the first three left-joined items;
+            // A3 is not triggered yet
+            // w1 = {}
+            // w2 = {}
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = {}
+            long time = 1000L;
+            for (int i = 0; i < expectedKeys.length; i++) {
+                inputTopic1.pipeInput(expectedKeys[i], "A" + expectedKeys[i], time + i);
+            }
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(0, "A0+null", 1000L),
+                    new KeyValueTimestamp<>(1, "A1+null", 1001L),
+                    new KeyValueTimestamp<>(2, "A2+null", 1002L)
+            );
+            
+            // push four items smaller timestamps (out of window) to the secondary stream; 
+            // this should produce four joined items
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = {}
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999) }
+            time = 1000L - 1L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "a" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(0, "A0+a0", 1000L),
+                    new KeyValueTimestamp<>(1, "A1+a1", 1001L),
+                    new KeyValueTimestamp<>(2, "A2+a2", 1002L),
+                    new KeyValueTimestamp<>(3, "A3+a3", 1003L)
+            );
+
+            // push four items with increased timestamps to the secondary stream; 
+            // this should produce four joined item
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999) }
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //            0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000) }
+            time += 1L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "b" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(0, "A0+b0", 1000L),
+                    new KeyValueTimestamp<>(1, "A1+b1", 1001L),
+                    new KeyValueTimestamp<>(2, "A2+b2", 1002L),
+                    new KeyValueTimestamp<>(3, "A3+b3", 1003L)
+            );
+
+            // push four items with increased timestamps to the secondary stream; 
+            // this should produce only three joined items;
+            // c0 arrives too late to be joined with A0
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //        0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000) }
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //            0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //            0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001) }
+            time += 1L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "c" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(1, "A1+c1", 1001L),
+                    new KeyValueTimestamp<>(2, "A2+c2", 1002L),
+                    new KeyValueTimestamp<>(3, "A3+c3", 1003L)
+            );
+
+            // push four items with increased timestamps to the secondary stream;
+            // this should produce only two joined items;
+            // d0 and d1 arrive too late to be joined with A0 and A1
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //        0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //        0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001) }
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //            0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //            0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001),
+            //            0:d0 (ts: 1002), 1:d1 (ts: 1002), 2:d2 (ts: 1002), 3:d3 (ts: 1002) }
+            time += 1L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "d" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(2, "A2+d2", 1002L),
+                    new KeyValueTimestamp<>(3, "A3+d3", 1003L)
+            );
+
+            // push four items with increased timestamps to the secondary stream; 
+            // this should produce one joined item;
+            // only e3 can be joined with A3;
+            // e0, e1 and e2 arrive too late to be joined with A0, A1 and A2
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //        0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //        0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001),
+            //        0:d0 (ts: 1002), 1:d1 (ts: 1002), 2:d2 (ts: 1002), 3:d3 (ts: 1002) }
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //            0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //            0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001),
+            //            0:d0 (ts: 1002), 1:d1 (ts: 1002), 2:d2 (ts: 1002), 3:d3 (ts: 1002),
+            //            0:e0 (ts: 1003), 1:e1 (ts: 1003), 2:e2 (ts: 1003), 3:e3 (ts: 1003) }
+            time += 1L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "e" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult(
+                new KeyValueTimestamp<>(3, "A3+e3", 1003L)
+            );
+
+            // push four items with larger timestamps to the secondary stream;
+            // no (non-)joined items can be produced
+            // 
+            // w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            // w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //        0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //        0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001),
+            //        0:d0 (ts: 1002), 1:d1 (ts: 1002), 2:d2 (ts: 1002), 3:d3 (ts: 1002),
+            //        0:e0 (ts: 1003), 1:e1 (ts: 1003), 2:e2 (ts: 1003), 3:e3 (ts: 1003) }
+            // --> w1 = { 0:A0 (ts: 1000), 1:A1 (ts: 1001), 2:A2 (ts: 1002), 3:A3 (ts: 1003) }
+            //     w2 = { 0:a0 (ts: 999), 1:a1 (ts: 999), 2:a2 (ts: 999), 3:a3 (ts: 999),
+            //            0:b0 (ts: 1000), 1:b1 (ts: 1000), 2:b2 (ts: 1000), 3:b3 (ts: 1000),
+            //            0:c0 (ts: 1001), 1:c1 (ts: 1001), 2:c2 (ts: 1001), 3:c3 (ts: 1001),
+            //            0:d0 (ts: 1002), 1:d1 (ts: 1002), 2:d2 (ts: 1002), 3:d3 (ts: 1002),
+            //            0:e0 (ts: 1003), 1:e1 (ts: 1003), 2:e2 (ts: 1003), 3:e3 (ts: 1003),
+            //            0:f0 (ts: 1100), 1:f1 (ts: 1100), 2:f2 (ts: 1100), 3:f3 (ts: 1100) }
+            time = 1000 + 100L;
+            for (final int expectedKey : expectedKeys) {
+                inputTopic2.pipeInput(expectedKey, "f" + expectedKey, time);
+            }
+            processor.checkAndClearProcessResult();
+        }
+    }
+
     @Test
     public void testLeftJoinWithInMemoryCustomSuppliers() {
         final JoinWindows joinWindows = JoinWindows.ofTimeDifferenceAndGrace(ofMillis(100L), ofMillis(0L));
@@ -609,8 +787,9 @@ public void testOrdering() {
             inputTopic1.pipeInput(1, "A1", 100L);
             processor.checkAndClearProcessResult();
 
-            // push one item to the other window that has a join; this should produce non-joined records with a closed window first, then
-            // the joined records
+            // push one item to the other window that has a join; 
+            // this should produce the joined record first;
+            // then non-joined record with a closed window
             // by the time they were produced before
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 100) }
             // w2 = { }
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
index 099dc5b0c83fa..28a5f1488fbce 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
@@ -108,11 +108,11 @@ public void testOuterJoinDuplicatesWithFixDisabledOldApi() {
             inputTopic2.pipeInput(1, "b1", 0L);
 
             processor.checkAndClearProcessResult(
-                    new KeyValueTimestamp<>(0, "A0+null", 0L),
-                    new KeyValueTimestamp<>(0, "A0-0+null", 0L),
-                    new KeyValueTimestamp<>(0, "A0+a0", 0L),
-                    new KeyValueTimestamp<>(0, "A0-0+a0", 0L),
-                    new KeyValueTimestamp<>(1, "null+b1", 0L)
+                new KeyValueTimestamp<>(0, "A0+null", 0L),
+                new KeyValueTimestamp<>(0, "A0-0+null", 0L),
+                new KeyValueTimestamp<>(0, "A0+a0", 0L),
+                new KeyValueTimestamp<>(0, "A0-0+a0", 0L),
+                new KeyValueTimestamp<>(1, "null+b1", 0L)
             );
         }
     }
@@ -438,13 +438,13 @@ public void testOrdering() {
             inputTopic1.pipeInput(1, "A1", 100L);
             processor.checkAndClearProcessResult();
 
-            // push one item to the other window that has a join; this should produce non-joined records with a closed window first, then
-            // the joined records
-            // by the time they were produced before
+            // push one item to the other window that has a join;
+            // this should produce the not-joined record first;
+            // then the joined record
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 100) }
             // w2 = { }
-            // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 0) }
-            // --> w2 = { 0:a0 (ts: 110) }
+            // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 100) }
+            // --> w2 = { 1:a1 (ts: 110) }
             inputTopic2.pipeInput(1, "a1", 110L);
             processor.checkAndClearProcessResult(
                 new KeyValueTimestamp<>(0, "A0+null", 0L),
@@ -788,7 +788,7 @@ public void shouldNotEmitLeftJoinResultForAsymmetricBeforeWindow() {
                 new KeyValueTimestamp<>(1, "A1+null", 1L)
             );
 
-            // push one item to the other stream; this should not produce any items
+            // push one item to the other stream; this should produce one right-join item
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
             // w2 = { 0:a0 (ts: 100), 1:a1 (ts: 102) }
             // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
@@ -841,7 +841,8 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
             final MockApiProcessor<Integer, String, Void, Void> processor = supplier.theCapturedProcessor();
             long time = 0L;
 
-            // push two items to the primary stream; the other window is empty; this should not produce any item
+            // push two items to the primary stream; the other window is empty; 
+            // this should produce one left-joined item
             // w1 = {}
             // w2 = {}
             // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
@@ -849,7 +850,9 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
             for (int i = 0; i < 2; i++) {
                 inputTopic1.pipeInput(expectedKeys[i], "A" + expectedKeys[i], time + i);
             }
-            processor.checkAndClearProcessResult();
+            processor.checkAndClearProcessResult(
+                new KeyValueTimestamp<>(0, "A0+null", 0L)
+            );
 
             // push one item to the other stream; this should produce one full-join item
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
@@ -863,7 +866,8 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
                 new KeyValueTimestamp<>(1, "A1+a1", 1L)
             );
 
-            // push one item to the other stream; this should produce one left-join item
+            // push one item to the other stream;
+            // this should not produce any item
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
             // w2 = { 1:a1 (ts: 1) }
             // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
@@ -871,9 +875,7 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
             time += 100;
             inputTopic2.pipeInput(expectedKeys[2], "a" + expectedKeys[2], time);
 
-            processor.checkAndClearProcessResult(
-                new KeyValueTimestamp<>(0, "A0+null", 0L)
-            );
+            processor.checkAndClearProcessResult();
 
             // push one item to the other stream; this should not produce any item
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
@@ -884,11 +886,12 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
 
             processor.checkAndClearProcessResult();
 
-            // push one item to the first stream; this should produce one full-join item
+            // push one item to the first stream;
+            // this should produce one inner-join item;
             // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1) }
             // w2 = { 1:a1 (ts: 1), 2:a2 (ts: 101), 3:a3 (ts: 101) }
             // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 1), 2:A2 (ts: 201) }
-            // --> w2 = { 1:a1 (ts: 1), 2:a2 (ts: 101), 3:a3 (ts: 101 }
+            // --> w2 = { 1:a1 (ts: 1), 2:a2 (ts: 101), 3:a3 (ts: 101) }
             time += 100;
             inputTopic1.pipeInput(expectedKeys[2], "A" + expectedKeys[2], time);
 

From ae047bbe56e8dc37ac18472ee631a14e1b35be82 Mon Sep 17 00:00:00 2001
From: "Cheng-Kai, Zhang" <kevin.zhang.tw@gmail.com>
Date: Wed, 6 Mar 2024 16:34:46 +0800
Subject: [PATCH 114/258] KAFKA-16347: Upgrade zookeeper 3.8.3 -> 3.8.4
 (#15480)

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 LICENSE-binary             | 4 ++--
 gradle/dependencies.gradle | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/LICENSE-binary b/LICENSE-binary
index 69361f878e73a..1031f83686410 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -270,8 +270,8 @@ scala-java8-compat_2.12-1.0.2
 scala-java8-compat_2.13-1.0.2
 snappy-java-1.1.10.5
 swagger-annotations-2.2.8
-zookeeper-3.8.3
-zookeeper-jute-3.8.3
+zookeeper-3.8.4
+zookeeper-jute-3.8.4
 
 ===============================================================================
 This product bundles various third-party components under other open source
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 5ac978f9b00d4..6503b51e192e9 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -164,7 +164,7 @@ versions += [
   snappy: "1.1.10.5",
   spotbugs: "4.8.0",
   zinc: "1.9.2",
-  zookeeper: "3.8.3",
+  zookeeper: "3.8.4",
   zstd: "1.5.5-11"
 ]
 

From f6198bc075f7e8e6af9e1fb53875e92de5057872 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Wed, 6 Mar 2024 12:13:39 +0300
Subject: [PATCH 115/258] KAFKA-14589 [3/4] Tests of ConsoleGroupCommand
 rewritten in java  (#15365)

Is contains some of ConsoleGroupCommand tests rewritten in java.
Intention of separate PR is to reduce changes and simplify review.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 checkstyle/import-control.xml                 |   2 +
 .../AbstractAuthorizerIntegrationTest.scala   | 130 +++++++++++++
 .../kafka/api/AbstractSaslTest.scala          |  22 +++
 .../kafka/api/AuthorizerIntegrationTest.scala | 118 +-----------
 ...aslClientsWithInvalidCredentialsTest.scala |  42 +----
 .../group/AuthorizerIntegrationTest.java      |  60 ++++++
 ...SaslClientsWithInvalidCredentialsTest.java | 177 ++++++++++++++++++
 7 files changed, 395 insertions(+), 156 deletions(-)
 create mode 100644 core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
 create mode 100644 core/src/test/scala/integration/kafka/api/AbstractSaslTest.scala
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java

diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index a52b3d94e3240..8bbf572821210 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -326,6 +326,8 @@
 
       <subpackage name="group">
         <allow pkg="kafka.api"/>
+        <allow pkg="kafka.security"/>
+        <allow pkg="kafka.zk" />
         <allow pkg="org.apache.kafka.tools"/>
         <allow pkg="org.apache.kafka.server.config" />
         <allow pkg="scala"/>
diff --git a/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
new file mode 100644
index 0000000000000..b3e1ba9a64d3b
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
@@ -0,0 +1,130 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+ * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+ * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package kafka.api
+
+import kafka.security.authorizer.AclAuthorizer
+import kafka.security.authorizer.AclEntry.WildcardHost
+import kafka.server.{BaseRequestTest, KafkaConfig}
+import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.clients.producer.ProducerConfig
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.acl.AccessControlEntry
+import org.apache.kafka.common.acl.AclOperation.CLUSTER_ACTION
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.resource.PatternType.LITERAL
+import org.apache.kafka.common.resource.{Resource, ResourcePattern}
+import org.apache.kafka.common.resource.ResourceType.{CLUSTER, GROUP, TOPIC, TRANSACTIONAL_ID}
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal}
+import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder
+import org.apache.kafka.metadata.authorizer.StandardAuthorizer
+import org.junit.jupiter.api.{BeforeEach, TestInfo}
+
+import java.util.Properties
+
+object AbstractAuthorizerIntegrationTest {
+  val BrokerPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "broker")
+  val ClientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
+
+  val BrokerListenerName = "BROKER"
+  val ClientListenerName = "CLIENT"
+  val ControllerListenerName = "CONTROLLER"
+
+  class PrincipalBuilder extends DefaultKafkaPrincipalBuilder(null, null) {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      context.listenerName match {
+        case BrokerListenerName | ControllerListenerName => BrokerPrincipal
+        case ClientListenerName => ClientPrincipal
+        case listenerName => throw new IllegalArgumentException(s"No principal mapped to listener $listenerName")
+      }
+    }
+  }
+}
+
+/**
+ * Abstract authorizer test to be used both in scala and java tests of authorizer.
+ */
+class AbstractAuthorizerIntegrationTest extends BaseRequestTest {
+  import AbstractAuthorizerIntegrationTest._
+
+  override def interBrokerListenerName: ListenerName = new ListenerName(BrokerListenerName)
+  override def listenerName: ListenerName = new ListenerName(ClientListenerName)
+  override def brokerCount: Int = 1
+
+  def clientPrincipal: KafkaPrincipal = ClientPrincipal
+  def brokerPrincipal: KafkaPrincipal = BrokerPrincipal
+
+  val clientPrincipalString: String = clientPrincipal.toString
+
+  val brokerId: Integer = 0
+  val topic = "topic"
+  val topicPattern = "topic.*"
+  val transactionalId = "transactional.id"
+  val producerId = 83392L
+  val part = 0
+  val correlationId = 0
+  val clientId = "client-Id"
+  val tp = new TopicPartition(topic, part)
+  val logDir = "logDir"
+  val group = "my-group"
+  val protocolType = "consumer"
+  val protocolName = "consumer-range"
+  val clusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
+  val topicResource = new ResourcePattern(TOPIC, topic, LITERAL)
+  val groupResource = new ResourcePattern(GROUP, group, LITERAL)
+  val transactionalIdResource = new ResourcePattern(TRANSACTIONAL_ID, transactionalId, LITERAL)
+
+  producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "1")
+  producerConfig.setProperty(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "false")
+  producerConfig.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, "50000")
+  consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, group)
+
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.BrokerIdProp, brokerId.toString)
+    addNodeProperties(properties)
+  }
+
+  override def kraftControllerConfigs(): collection.Seq[Properties] = {
+    val controllerConfigs = super.kraftControllerConfigs()
+    controllerConfigs.foreach(addNodeProperties)
+    controllerConfigs
+  }
+
+  private def addNodeProperties(properties: Properties): Unit = {
+    if (isKRaftTest()) {
+      properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[StandardAuthorizer].getName)
+      properties.put(StandardAuthorizer.SUPER_USERS_CONFIG, BrokerPrincipal.toString)
+    } else {
+      properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
+    }
+
+    properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+    properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
+    properties.put(KafkaConfig.TransactionsTopicPartitionsProp, "1")
+    properties.put(KafkaConfig.TransactionsTopicReplicationFactorProp, "1")
+    properties.put(KafkaConfig.TransactionsTopicMinISRProp, "1")
+    properties.put(KafkaConfig.UnstableApiVersionsEnableProp, "true")
+    properties.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[PrincipalBuilder].getName)
+  }
+
+  @BeforeEach
+  override def setUp(testInfo: TestInfo): Unit = {
+    doSetup(testInfo, createOffsetsTopic = false)
+
+    // Allow inter-broker communication
+    addAndVerifyAcls(Set(new AccessControlEntry(brokerPrincipal.toString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
+
+    createOffsetsTopic(listenerName = interBrokerListenerName)
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/AbstractSaslTest.scala b/core/src/test/scala/integration/kafka/api/AbstractSaslTest.scala
new file mode 100644
index 0000000000000..4bef3422a246d
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/AbstractSaslTest.scala
@@ -0,0 +1,22 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+ * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+ * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package kafka.api
+
+/**
+ * Abstract Sasl test to be used both in scala and java tests of SASL.
+ * Separate class required to overcome issues related to usage of scala trait in java code.
+ * @see SaslClientsWithInvalidCredentialsTest
+ */
+abstract class AbstractSaslTest extends IntegrationTestHarness with SaslSetup {
+}
+
diff --git a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
index c399b46450978..ded470eba80d4 100644
--- a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
@@ -18,10 +18,8 @@ import java.util
 import java.util.concurrent.ExecutionException
 import java.util.regex.Pattern
 import java.util.{Collections, Optional, Properties}
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.security.authorizer.AclEntry
 import kafka.security.authorizer.AclEntry.WildcardHost
-import kafka.server.{BaseRequestTest, KafkaConfig}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import kafka.utils.TestUtils.waitUntilTrue
 import org.apache.kafka.clients.admin.{Admin, AlterConfigOp, NewTopic}
@@ -30,7 +28,6 @@ import org.apache.kafka.clients.producer._
 import org.apache.kafka.common.acl.AclOperation._
 import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
 import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBindingFilter, AclOperation, AclPermissionType}
-import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
 import org.apache.kafka.common.config.{ConfigResource, LogLevelConfig, TopicConfig}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic.GROUP_METADATA_TOPIC_NAME
@@ -53,14 +50,11 @@ import org.apache.kafka.common.requests._
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.resource.ResourceType._
 import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourcePatternFilter, ResourceType}
-import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{ElectionType, IsolationLevel, KafkaException, Node, TopicPartition, Uuid, requests}
-import org.apache.kafka.metadata.authorizer.StandardAuthorizer
 import org.apache.kafka.test.{TestUtils => JTestUtils}
 import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.{BeforeEach, TestInfo}
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.{CsvSource, ValueSource}
 
@@ -72,55 +66,7 @@ import scala.annotation.nowarn
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
-object AuthorizerIntegrationTest {
-  val BrokerPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "broker")
-  val ClientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
-
-  val BrokerListenerName = "BROKER"
-  val ClientListenerName = "CLIENT"
-  val ControllerListenerName = "CONTROLLER"
-
-  class PrincipalBuilder extends DefaultKafkaPrincipalBuilder(null, null) {
-    override def build(context: AuthenticationContext): KafkaPrincipal = {
-      context.listenerName match {
-        case BrokerListenerName | ControllerListenerName => BrokerPrincipal
-        case ClientListenerName => ClientPrincipal
-        case listenerName => throw new IllegalArgumentException(s"No principal mapped to listener $listenerName")
-      }
-    }
-  }
-}
-
-class AuthorizerIntegrationTest extends BaseRequestTest {
-  import AuthorizerIntegrationTest._
-
-  override def interBrokerListenerName: ListenerName = new ListenerName(BrokerListenerName)
-  override def listenerName: ListenerName = new ListenerName(ClientListenerName)
-  override def brokerCount: Int = 1
-
-  def clientPrincipal: KafkaPrincipal = ClientPrincipal
-  def brokerPrincipal: KafkaPrincipal = BrokerPrincipal
-
-  val clientPrincipalString: String = clientPrincipal.toString
-
-  val brokerId: Integer = 0
-  val topic = "topic"
-  val topicPattern = "topic.*"
-  val transactionalId = "transactional.id"
-  val producerId = 83392L
-  val part = 0
-  val correlationId = 0
-  val clientId = "client-Id"
-  val tp = new TopicPartition(topic, part)
-  val logDir = "logDir"
-  val group = "my-group"
-  val protocolType = "consumer"
-  val protocolName = "consumer-range"
-  val clusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
-  val topicResource = new ResourcePattern(TOPIC, topic, LITERAL)
-  val groupResource = new ResourcePattern(GROUP, group, LITERAL)
-  val transactionalIdResource = new ResourcePattern(TRANSACTIONAL_ID, transactionalId, LITERAL)
-
+class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   val groupReadAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)))
   val groupDescribeAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
   val groupDeleteAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)))
@@ -143,40 +89,6 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   val numRecords = 1
 
-  producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "1")
-  producerConfig.setProperty(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "false")
-  producerConfig.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, "50000")
-  consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, group)
-
-  override def brokerPropertyOverrides(properties: Properties): Unit = {
-    properties.put(KafkaConfig.BrokerIdProp, brokerId.toString)
-    addNodeProperties(properties)
-  }
-
-  override def kraftControllerConfigs(): collection.Seq[Properties] = {
-    val controllerConfigs = super.kraftControllerConfigs()
-    controllerConfigs.foreach(addNodeProperties)
-    controllerConfigs
-  }
-
-  private def addNodeProperties(properties: Properties): Unit = {
-    if (isKRaftTest()) {
-      properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[StandardAuthorizer].getName)
-      properties.put(StandardAuthorizer.SUPER_USERS_CONFIG, BrokerPrincipal.toString)
-    } else {
-      properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
-    }
-
-    properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
-    properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
-    properties.put(KafkaConfig.TransactionsTopicPartitionsProp, "1")
-    properties.put(KafkaConfig.TransactionsTopicReplicationFactorProp, "1")
-    properties.put(KafkaConfig.TransactionsTopicMinISRProp, "1")
-    properties.put(KafkaConfig.UnstableApiVersionsEnableProp, "true")
-    properties.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[PrincipalBuilder].getName)
-  }
-
-
   val requestKeyToError = (topicNames: Map[Uuid, String], version: Short) => Map[ApiKeys, Nothing => Errors](
     ApiKeys.METADATA -> ((resp: requests.MetadataResponse) => resp.errors.asScala.find(_._1 == topic).getOrElse(("test", Errors.NONE))._2),
     ApiKeys.PRODUCE -> ((resp: requests.ProduceResponse) => {
@@ -333,16 +245,6 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.DESCRIBE_TRANSACTIONS -> transactionalIdDescribeAcl
   )
 
-  @BeforeEach
-  override def setUp(testInfo: TestInfo): Unit = {
-    doSetup(testInfo, createOffsetsTopic = false)
-
-    // Allow inter-broker communication
-    addAndVerifyAcls(Set(new AccessControlEntry(brokerPrincipal.toString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
-
-    createOffsetsTopic(listenerName = interBrokerListenerName)
-  }
-
   private def createMetadataRequest(allowAutoTopicCreation: Boolean) = {
     new requests.MetadataRequest.Builder(List(topic).asJava, allowAutoTopicCreation).build()
   }
@@ -1703,20 +1605,6 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     createAdminClient().describeConsumerGroups(Seq(group).asJava).describedGroups().get(group).get()
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDescribeGroupCliWithGroupDescribe(quorum: String): Unit = {
-    createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupService = new ConsumerGroupService(opts)
-    consumerGroupService.describeGroups()
-    consumerGroupService.close()
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testListGroupApiWithAndWithoutListGroupAcls(quorum: String): Unit = {
diff --git a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
index eba51c968f9ae..11afbd89d2434 100644
--- a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
@@ -23,7 +23,6 @@ import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.SaslAuthenticationException
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
 import org.junit.jupiter.api.Assertions._
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
 import kafka.server.KafkaConfig
 import kafka.utils.{JaasTestUtils, TestUtils}
 import kafka.zk.ConfigEntityChangeNotificationZNode
@@ -31,7 +30,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 
-class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with SaslSetup {
+class SaslClientsWithInvalidCredentialsTest extends AbstractSaslTest {
   private val kafkaClientSaslMechanism = "SCRAM-SHA-256"
   private val kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
   override protected val securityProtocol = SecurityProtocol.SASL_PLAINTEXT
@@ -166,45 +165,6 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
     }
   }
 
-  @Test
-  def testConsumerGroupServiceWithAuthenticationFailure(): Unit = {
-    val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
-
-    val consumer = createConsumer()
-    try {
-      consumer.subscribe(List(topic).asJava)
-
-      verifyAuthenticationException(consumerGroupService.listGroups())
-    } finally consumerGroupService.close()
-  }
-
-  @Test
-  def testConsumerGroupServiceWithAuthenticationSuccess(): Unit = {
-    createClientCredential()
-    val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
-
-    val consumer = createConsumer()
-    try {
-      consumer.subscribe(List(topic).asJava)
-
-      verifyWithRetry(consumer.poll(Duration.ofMillis(1000)))
-      assertEquals(1, consumerGroupService.listConsumerGroups().size)
-    }
-    finally consumerGroupService.close()
-  }
-
-  private def prepareConsumerGroupService = {
-    val propsFile = TestUtils.tempPropertiesFile(Map("security.protocol" -> "SASL_PLAINTEXT", "sasl.mechanism" -> kafkaClientSaslMechanism))
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(),
-                        "--describe",
-                        "--group", "test.group",
-                        "--command-config", propsFile.getAbsolutePath)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupService = new ConsumerGroupService(opts)
-    consumerGroupService
-  }
-
   private def createClientCredential(): Unit = {
     createScramCredentialsViaPrivilegedAdminClient(JaasTestUtils.KafkaScramUser2, JaasTestUtils.KafkaScramPassword2)
   }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
new file mode 100644
index 0000000000000..3094572b04921
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import kafka.admin.ConsumerGroupCommand;
+import kafka.api.AbstractAuthorizerIntegrationTest;
+import kafka.security.authorizer.AclEntry;
+import org.apache.kafka.common.acl.AccessControlEntry;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import scala.collection.immutable.Map$;
+
+import java.util.Collections;
+import java.util.Properties;
+
+import static org.apache.kafka.common.acl.AclOperation.DESCRIBE;
+import static org.apache.kafka.common.acl.AclPermissionType.ALLOW;
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME;
+import static org.apache.kafka.tools.consumer.group.ConsumerGroupCommandTest.set;
+
+public class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDescribeGroupCliWithGroupDescribe(String quorum) {
+        addAndVerifyAcls(set(Collections.singleton(new AccessControlEntry(ClientPrincipal().toString(), AclEntry.WildcardHost(), DESCRIBE, ALLOW))), groupResource());
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group()};
+        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(cgcArgs);
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupService = new ConsumerGroupCommand.ConsumerGroupService(opts, Map$.MODULE$.empty());
+        consumerGroupService.describeGroups();
+        consumerGroupService.close();
+    }
+
+    private void createTopicWithBrokerPrincipal(String topic) {
+        // Note the principal builder implementation maps all connections on the
+        // inter-broker listener to the broker principal.
+        createTopic(
+            topic,
+            1,
+            1,
+            new Properties(),
+            interBrokerListenerName(),
+            new Properties()
+        );
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java
new file mode 100644
index 0000000000000..06b727a59a8fd
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import kafka.admin.ConsumerGroupCommand;
+import kafka.api.AbstractSaslTest;
+import kafka.api.Both$;
+import kafka.utils.JaasTestUtils;
+import kafka.zk.ConfigEntityChangeNotificationZNode;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
+import org.junit.jupiter.api.function.Executable;
+import scala.Option;
+import scala.Some$;
+import scala.collection.JavaConverters;
+import scala.collection.Seq;
+import scala.collection.immutable.Map$;
+
+import java.io.File;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Properties;
+
+import static org.apache.kafka.tools.consumer.group.ConsumerGroupCommandTest.seq;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class SaslClientsWithInvalidCredentialsTest extends AbstractSaslTest {
+    private static final String TOPIC = "topic";
+    public static final int NUM_PARTITIONS = 1;
+    public static final int BROKER_COUNT = 1;
+    public static final String KAFKA_CLIENT_SASL_MECHANISM = "SCRAM-SHA-256";
+    private static final Seq<String> KAFKA_SERVER_SASL_MECHANISMS = seq(Collections.singletonList(KAFKA_CLIENT_SASL_MECHANISM));
+
+    @SuppressWarnings({"deprecation"})
+    private Consumer<byte[], byte[]> createConsumer() {
+        return createConsumer(
+            new ByteArrayDeserializer(),
+            new ByteArrayDeserializer(),
+            new Properties(),
+            JavaConverters.asScalaSet(Collections.<String>emptySet()).toList()
+        );
+    }
+
+    @Override
+    public SecurityProtocol securityProtocol() {
+        return SecurityProtocol.SASL_PLAINTEXT;
+    }
+
+    @Override
+    public Option<Properties> serverSaslProperties() {
+        return Some$.MODULE$.apply(kafkaServerSaslProperties(KAFKA_SERVER_SASL_MECHANISMS, KAFKA_CLIENT_SASL_MECHANISM));
+    }
+
+    @Override
+    public Option<Properties> clientSaslProperties() {
+        return Some$.MODULE$.apply(kafkaClientSaslProperties(KAFKA_CLIENT_SASL_MECHANISM, false));
+    }
+
+    @Override
+    public int brokerCount() {
+        return 1;
+    }
+
+    @Override
+    public void configureSecurityBeforeServersStart(TestInfo testInfo) {
+        super.configureSecurityBeforeServersStart(testInfo);
+        zkClient().makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path());
+        // Create broker credentials before starting brokers
+        createScramCredentials(zkConnect(), JaasTestUtils.KafkaScramAdmin(), JaasTestUtils.KafkaScramAdminPassword());
+    }
+
+    @Override
+    public Admin createPrivilegedAdminClient() {
+        return createAdminClient(bootstrapServers(listenerName()), securityProtocol(), trustStoreFile(), clientSaslProperties(),
+            KAFKA_CLIENT_SASL_MECHANISM, JaasTestUtils.KafkaScramAdmin(), JaasTestUtils.KafkaScramAdminPassword());
+    }
+
+    @BeforeEach
+    @Override
+    public void setUp(TestInfo testInfo) {
+        startSasl(jaasSections(KAFKA_SERVER_SASL_MECHANISMS, Some$.MODULE$.apply(KAFKA_CLIENT_SASL_MECHANISM), Both$.MODULE$,
+            JaasTestUtils.KafkaServerContextName()));
+        super.setUp(testInfo);
+        createTopic(
+            TOPIC,
+            NUM_PARTITIONS,
+            BROKER_COUNT,
+            new Properties(),
+            listenerName(),
+            new Properties());
+    }
+
+    @AfterEach
+    @Override
+    public void tearDown() {
+        super.tearDown();
+        closeSasl();
+    }
+
+    @Test
+    public void testConsumerGroupServiceWithAuthenticationFailure() throws Exception {
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupService = prepareConsumerGroupService();
+        try (Consumer<byte[], byte[]> consumer = createConsumer()) {
+            consumer.subscribe(Collections.singletonList(TOPIC));
+
+            verifyAuthenticationException(consumerGroupService::listGroups);
+        } finally {
+            consumerGroupService.close();
+        }
+    }
+
+    @Test
+    public void testConsumerGroupServiceWithAuthenticationSuccess() throws Exception {
+        createScramCredentialsViaPrivilegedAdminClient(JaasTestUtils.KafkaScramUser2(), JaasTestUtils.KafkaScramPassword2());
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupService = prepareConsumerGroupService();
+        try (Consumer<byte[], byte[]> consumer = createConsumer()) {
+            consumer.subscribe(Collections.singletonList(TOPIC));
+
+            TestUtils.waitForCondition(() -> {
+                try {
+                    consumer.poll(Duration.ofMillis(1000));
+                    return true;
+                } catch (SaslAuthenticationException ignored) {
+                    return false;
+                }
+            }, "failed to poll data with authentication");
+            assertEquals(1, consumerGroupService.listConsumerGroups().size());
+        } finally {
+            consumerGroupService.close();
+        }
+    }
+
+    private ConsumerGroupCommand.ConsumerGroupService prepareConsumerGroupService() throws IOException {
+        File propsFile = TestUtils.tempFile(
+            "security.protocol=SASL_PLAINTEXT\n" +
+            "sasl.mechanism=" + KAFKA_CLIENT_SASL_MECHANISM);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()),
+            "--describe",
+            "--group", "test.group",
+            "--command-config", propsFile.getAbsolutePath()};
+        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(cgcArgs);
+        return new ConsumerGroupCommand.ConsumerGroupService(opts, Map$.MODULE$.empty());
+    }
+
+    private void verifyAuthenticationException(Executable action) {
+        long startMs = System.currentTimeMillis();
+        assertThrows(Exception.class, action);
+        long elapsedMs = System.currentTimeMillis() - startMs;
+        assertTrue(elapsedMs <= 5000, "Poll took too long, elapsed=" + elapsedMs);
+    }
+}

From 8f5bb8cfb306513d215a5258a3dbcf8d0bad9c94 Mon Sep 17 00:00:00 2001
From: "Cheng-Kai, Zhang" <kevin.zhang.tw@gmail.com>
Date: Wed, 6 Mar 2024 17:19:54 +0800
Subject: [PATCH 116/258] KAFKA-16252: Fix the documentation and adjust the
 format (#15473)

Currently, there are few document files generated automatically like the task genConnectMetricsDocs
However, the unwanted log information also added into it.
And the format is not aligned with other which has Mbean located of the third column.

I modified the code logic so the format could follow other section in ops.html
Also close the log since we take everything from the std as a documentation

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle                                  |  2 +-
 .../apache/kafka/common/metrics/Metrics.java  | 32 +++++++------------
 2 files changed, 12 insertions(+), 22 deletions(-)

diff --git a/build.gradle b/build.gradle
index 7dc8f50342fdd..cba4381aae5e0 100644
--- a/build.gradle
+++ b/build.gradle
@@ -3048,7 +3048,7 @@ project(':connect:runtime') {
   }
 
   task genConnectMetricsDocs(type: JavaExec) {
-    classpath = sourceSets.test.runtimeClasspath
+    classpath = sourceSets.main.runtimeClasspath
     mainClass = 'org.apache.kafka.connect.runtime.ConnectMetrics'
     if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() }
     standardOutput = new File(generatedDocsDir, "connect_metrics.html").newOutputStream()
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/Metrics.java b/clients/src/main/java/org/apache/kafka/common/metrics/Metrics.java
index 6447cdb5c7e4e..b52285dac63f4 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/Metrics.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/Metrics.java
@@ -277,40 +277,30 @@ public static String toHtmlTable(String domain, Iterable<MetricNameTemplate> all
                 }
             }
         }
-        
         StringBuilder b = new StringBuilder();
-        b.append("<table class=\"data-table\"><tbody>\n");
-    
+        b.append("<table class=\"data-table\">\n<tbody>\n");
+        b.append("<tr>\n");
+        b.append("<th>Metric/Attribute name</th>\n");
+        b.append("<th>Description</th>\n");
+        b.append("<th>Mbean name</th>\n");
+        b.append("</tr>\n");
         for (Entry<String, Map<String, String>> e : beansAndAttributes.entrySet()) {
-            b.append("<tr>\n");
-            b.append("<td colspan=3 class=\"mbeanName\" style=\"background-color:#ccc; font-weight: bold;\">");
-            b.append(e.getKey());
-            b.append("</td>");
-            b.append("</tr>\n");
-            
-            b.append("<tr>\n");
-            b.append("<th style=\"width: 90px\"></th>\n");
-            b.append("<th>Attribute name</th>\n");
-            b.append("<th>Description</th>\n");
-            b.append("</tr>\n");
-            
             for (Entry<String, String> e2 : e.getValue().entrySet()) {
                 b.append("<tr>\n");
-                b.append("<td></td>");
                 b.append("<td>");
                 b.append(e2.getKey());
-                b.append("</td>");
+                b.append("</td>\n");
                 b.append("<td>");
                 b.append(e2.getValue());
-                b.append("</td>");
+                b.append("</td>\n");
+                b.append("<td>");
+                b.append(e.getKey());
+                b.append("</td>\n");
                 b.append("</tr>\n");
             }
-    
         }
         b.append("</tbody></table>");
-    
         return b.toString();
-    
     }
 
     public MetricConfig config() {

From bc0c73e944af69966c43d51d20417b441aa0f3f5 Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Wed, 6 Mar 2024 19:39:34 +0800
Subject: [PATCH 117/258] KAFKA-16322 upgrade jline from 3.22.0 to 3.25.1
 (#15464)

An issue in the component "GroovyEngine.execute" of jline-groovy versions through 3.24.1 allows attackers to cause an OOM (OutofMemory) error. Please refer to https://devhub.checkmarx.com/cve-details/CVE-2023-50572 for more details

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 LICENSE-binary             | 2 +-
 gradle/dependencies.gradle | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/LICENSE-binary b/LICENSE-binary
index 1031f83686410..df80c6b57ea0e 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -333,7 +333,7 @@ zstd-jni-1.5.5-11 see: licenses/zstd-jni-BSD-2-clause
 ---------------------------------------
 BSD 3-Clause
 
-jline-3.22.0, see: licenses/jline-BSD-3-clause
+jline-3.25.1, see: licenses/jline-BSD-3-clause
 jsr305-3.0.2, see: licenses/jsr305-BSD-3-clause
 paranamer-2.8, see: licenses/paranamer-BSD-3-clause
 protobuf-java-3.23.4, see: licenses/protobuf-java-BSD-3-clause
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 6503b51e192e9..4f31bcd25c1d4 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -108,7 +108,7 @@ versions += [
   javassist: "3.29.2-GA",
   jetty: "9.4.53.v20231009",
   jersey: "2.39.1",
-  jline: "3.22.0",
+  jline: "3.25.1",
   jmh: "1.37",
   hamcrest: "2.2",
   scalaLogging: "3.9.4",

From 0b30dc2e8276bb044ab31cb31e518997071aceee Mon Sep 17 00:00:00 2001
From: Truc Nguyen <119303279+trnguyencflt@users.noreply.github.com>
Date: Wed, 6 Mar 2024 12:33:52 +0000
Subject: [PATCH 118/258] Update jetty to 9.4.54.v20240208 for CVE-2024-22201
 fix (#1064)

Update jetty to version [9.4.54.v20240208](https://github.com/jetty/jetty.project/releases/tag/jetty-9.4.54.v20240208) to fix CVE-2024-22201
---
 gradle/dependencies.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 659931bada299..7fe6bba6874de 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -106,7 +106,7 @@ versions += [
   jackson: "2.16.1",
   jacoco: "0.8.10",
   javassist: "3.29.2-GA",
-  jetty: "9.4.53.v20231009",
+  jetty: "9.4.54.v20240208",
   jersey: "2.39.1",
   jline: "3.25.1",
   jmh: "1.37",

From 6d7b25bb25ee817c03606e666608cece5cbf621d Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Wed, 6 Mar 2024 10:27:12 -0800
Subject: [PATCH 119/258] KAFKA-15797: Fix flaky EOS_v2 upgrade test (#15449)

Originally, we set commit-interval to MAX_VALUE for this test,
to ensure we only commit expliclity. However, we needed to decrease it
later on when adding the tx-timeout verification.

We did see failing test for which commit-interval hit, resulting in
failing test runs. This PR increase the commit-interval close to
test-timeout to avoid commit-interval from triggering.

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 .../kafka/streams/integration/EosV2UpgradeIntegrationTest.java  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
index 011090c152667..bac9ae37eb8e7 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/EosV2UpgradeIntegrationTest.java
@@ -939,7 +939,7 @@ public void close() {}
         final Properties properties = new Properties();
         properties.put(StreamsConfig.CLIENT_ID_CONFIG, appDir);
         properties.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, processingGuarantee);
-        final long commitInterval = Duration.ofMinutes(1L).toMillis();
+        final long commitInterval = Duration.ofMinutes(5L).toMillis();
         properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, commitInterval);
         properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.METADATA_MAX_AGE_CONFIG), Duration.ofSeconds(1L).toMillis());
         properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG), "earliest");

From 62998b72642109ac455564d26c08e37eb0e9f0ea Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Wed, 6 Mar 2024 13:31:33 -0500
Subject: [PATCH 120/258] KAFKA-14683: Migrate WorkerSinkTaskTest to Mockito
 (3/3) (#15316)

Reviewers: Greg Harris <greg.harris@aiven.io>
---
 checkstyle/suppressions.xml                   |    6 +-
 .../kafka/connect/runtime/WorkerSinkTask.java |    5 +
 .../runtime/WorkerSinkTaskMockitoTest.java    |  945 +++++++++++-
 .../connect/runtime/WorkerSinkTaskTest.java   | 1298 -----------------
 4 files changed, 899 insertions(+), 1355 deletions(-)
 delete mode 100644 connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index c65cd675a9ef3..32916fe6602fd 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -96,7 +96,7 @@
               files="(AbstractFetch|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/>
 
     <suppress checks="JavaNCSS"
-              files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest|KafkaRaftClientTest).java"/>
+              files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|WorkerSinkTaskMockitoTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest|KafkaRaftClientTest).java"/>
 
     <suppress checks="NPathComplexity"
               files="(ConsumerCoordinator|BufferPool|MetricName|Node|ConfigDef|RecordBatch|SslFactory|SslTransportLayer|MetadataResponse|KerberosLogin|Selector|Sender|Serdes|TokenInformation|Agent|Values|PluginUtils|MiniTrogdorCluster|TasksRequest|KafkaProducer|AbstractStickyAssignor|KafkaRaftClient|Authorizer|FetchSessionHandler|RecordAccumulator).java"/>
@@ -173,12 +173,12 @@
               files="(DistributedHerder|KafkaBasedLog|WorkerSourceTaskWithTopicCreation|WorkerSourceTask)Test.java"/>
 
     <suppress checks="ClassFanOutComplexity"
-              files="(WorkerSink|WorkerSource|ErrorHandling)Task(|WithTopicCreation)Test.java"/>
+              files="(WorkerSink|WorkerSource|ErrorHandling)Task(|WithTopicCreation|Mockito)Test.java"/>
     <suppress checks="ClassFanOutComplexity"
               files="DistributedHerderTest.java"/>
 
     <suppress checks="MethodLength"
-              files="(RequestResponse|WorkerSinkTask)Test.java"/>
+              files="(RequestResponse|WorkerSinkTask|WorkerSinkTaskMockito)Test.java"/>
 
     <suppress checks="JavaNCSS"
               files="(DistributedHerder|Worker)Test.java"/>
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
index 78a61e755f4ba..25d8b54d4d6d5 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
@@ -354,6 +354,11 @@ Map<TopicPartition, OffsetAndMetadata> lastCommittedOffsets() {
         return Collections.unmodifiableMap(lastCommittedOffsets);
     }
 
+    //VisibleForTesting
+    Map<TopicPartition, OffsetAndMetadata> currentOffsets() {
+        return Collections.unmodifiableMap(currentOffsets);
+    }
+
     private void doCommitSync(Map<TopicPartition, OffsetAndMetadata> offsets, int seqno) {
         log.debug("{} Committing offsets synchronously using sequence number {}: {}", this, seqno, offsets);
         try {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
index f39dce6646070..83bf4fe0c7b0a 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
@@ -16,6 +16,46 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import static java.util.Arrays.asList;
+import static java.util.Collections.singleton;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -36,6 +76,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.WorkerSinkTask.SinkTaskMetricsGroup;
@@ -69,41 +110,6 @@
 import org.mockito.stubbing.Answer;
 import org.mockito.stubbing.OngoingStubbing;
 
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.function.Supplier;
-import java.util.regex.Pattern;
-
-import static java.util.Arrays.asList;
-import static java.util.Collections.singleton;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyList;
-import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.atLeastOnce;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(MockitoJUnitRunner.StrictStubs.class)
 public class WorkerSinkTaskMockitoTest {
     // These are fixed to keep this code simpler. In this example we assume byte[] raw values
@@ -371,9 +377,9 @@ public void testPollRedelivery() {
                 .thenAnswer(expectConsumerPoll(0));
         expectConversionAndTransformation(null, new RecordHeaders());
 
-        doAnswer(invocation -> null)
+        doNothing()
                 .doThrow(new RetriableException("retry"))
-                .doAnswer(invocation -> null)
+                .doNothing()
                 .when(sinkTask).put(anyList());
 
         workerTask.iteration();
@@ -447,6 +453,86 @@ public void testPollRedelivery() {
         assertSinkMetricValue("offset-commit-completion-total", 1.0);
     }
 
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testPollRedeliveryWithConsumerRebalance() {
+        createTask(initialState);
+        expectTaskGetTopic();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        Set<TopicPartition> newAssignment = new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3));
+
+        when(consumer.assignment())
+                .thenReturn(INITIAL_ASSIGNMENT, INITIAL_ASSIGNMENT, INITIAL_ASSIGNMENT)
+                .thenReturn(newAssignment, newAssignment, newAssignment)
+                .thenReturn(Collections.singleton(TOPIC_PARTITION3),
+                        Collections.singleton(TOPIC_PARTITION3),
+                        Collections.singleton(TOPIC_PARTITION3));
+
+        INITIAL_ASSIGNMENT.forEach(tp -> when(consumer.position(tp)).thenReturn(FIRST_OFFSET));
+        when(consumer.position(TOPIC_PARTITION3)).thenReturn(FIRST_OFFSET);
+
+        when(consumer.poll(any(Duration.class)))
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                })
+                .thenAnswer(expectConsumerPoll(1))
+                // Empty consumer poll (all partitions are paused) with rebalance; one new partition is assigned
+                .thenAnswer(invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
+                    return ConsumerRecords.empty();
+                })
+                .thenAnswer(expectConsumerPoll(0))
+                // Non-empty consumer poll; all initially-assigned partitions are revoked in rebalance, and new partitions are allowed to resume
+                .thenAnswer(invocation -> {
+                    ConsumerRecord<byte[], byte[]> newRecord = new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET, RAW_KEY, RAW_VALUE);
+
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.emptyList());
+                    return new ConsumerRecords<>(Collections.singletonMap(TOPIC_PARTITION3, Collections.singletonList(newRecord)));
+                });
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        doNothing()
+                // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
+                .doThrow(new RetriableException("retry"))
+                .doThrow(new RetriableException("retry"))
+                .doThrow(new RetriableException("retry"))
+                .doNothing()
+                .when(sinkTask).put(any(Collection.class));
+
+        workerTask.iteration();
+
+        // Pause
+        workerTask.iteration();
+        verify(consumer).pause(INITIAL_ASSIGNMENT);
+
+        workerTask.iteration();
+        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION3));
+        // All partitions are re-paused in order to pause any newly-assigned partitions so that redelivery efforts can continue
+        verify(consumer).pause(newAssignment);
+
+        workerTask.iteration();
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = INITIAL_ASSIGNMENT.stream()
+                .collect(Collectors.toMap(Function.identity(), tp -> new OffsetAndMetadata(FIRST_OFFSET)));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+        newAssignment = Collections.singleton(TOPIC_PARTITION3);
+
+        workerTask.iteration();
+        verify(sinkTask).close(INITIAL_ASSIGNMENT);
+
+        // All partitions are resumed, as all previously paused-for-redelivery partitions were revoked
+        newAssignment.forEach(tp -> {
+            verify(consumer).resume(Collections.singleton(tp));
+        });
+    }
+
     @Test
     public void testErrorInRebalancePartitionLoss() {
         RuntimeException exception = new RuntimeException("Revocation error");
@@ -601,45 +687,796 @@ public void testPartialRevocationAndAssignment() {
         verify(sinkTask, times(4)).put(Collections.emptyList());
     }
 
+    @Test
     @SuppressWarnings("unchecked")
+    public void testPreCommitFailureAfterPartialRevocationAndAssignment() {
+        createTask(initialState);
+        expectTaskGetTopic();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        when(consumer.assignment())
+                .thenReturn(INITIAL_ASSIGNMENT, INITIAL_ASSIGNMENT)
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3)))
+                .thenReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3)));
+
+        INITIAL_ASSIGNMENT.forEach(tp -> when(consumer.position(tp)).thenReturn(FIRST_OFFSET));
+        when(consumer.position(TOPIC_PARTITION3)).thenReturn(FIRST_OFFSET);
+
+        // First poll; assignment is [TP1, TP2]
+        when(consumer.poll(any(Duration.class)))
+                .thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                })
+                // Second poll; a single record is delivered from TP1
+                .thenAnswer(expectConsumerPoll(1))
+                // Third poll; assignment changes to [TP2]
+                .thenAnswer(invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
+                    return ConsumerRecords.empty();
+                })
+                // Fourth poll; assignment changes to [TP2, TP3]
+                .thenAnswer(invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
+                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
+                    return ConsumerRecords.empty();
+                })
+                // Fifth poll; an offset commit takes place
+                .thenAnswer(expectConsumerPoll(0));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        // First iteration--first call to poll, first consumer assignment
+        workerTask.iteration();
+        // Second iteration--second call to poll, delivery of one record
+        workerTask.iteration();
+        // Third iteration--third call to poll, partial consumer revocation
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+        doNothing().when(consumer).commitSync(offsets);
+
+        workerTask.iteration();
+        verify(sinkTask).close(Collections.singleton(TOPIC_PARTITION));
+        verify(sinkTask, times(2)).put(Collections.emptyList());
+
+        // Fourth iteration--fourth call to poll, partial consumer assignment
+        workerTask.iteration();
+
+        verify(sinkTask).open(Collections.singleton(TOPIC_PARTITION3));
+
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        workerCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenThrow(new ConnectException("Failed to flush"));
+
+        // Fifth iteration--task-requested offset commit with failure in SinkTask::preCommit
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration();
+
+        verify(consumer).seek(TOPIC_PARTITION2, FIRST_OFFSET);
+        verify(consumer).seek(TOPIC_PARTITION3, FIRST_OFFSET);
+    }
+
     @Test
-    public void testTaskCancelPreventsFinalOffsetCommit() {
+    public void testWakeupInCommitSyncCausesRetry() {
         createTask(initialState);
 
         workerTask.initialize(TASK_CONFIG);
+        time.sleep(30000L);
         workerTask.initializeAndStart();
+        time.sleep(30000L);
         verifyInitializeTask();
 
         expectTaskGetTopic();
         expectPollInitialAssignment()
-                // Put one message through the task to get some offsets to commit
                 .thenAnswer(expectConsumerPoll(1))
-                // the second put will return after the task is stopped and cancelled (asynchronously)
-                .thenAnswer(expectConsumerPoll(1));
-
+                .thenAnswer(invocation -> {
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    return ConsumerRecords.empty();
+                });
         expectConversionAndTransformation(null, new RecordHeaders());
 
-        doAnswer(invocation -> null)
-                .doAnswer(invocation -> null)
-                .doAnswer(invocation -> {
+        workerTask.iteration(); // poll for initial assignment
+        time.sleep(30000L);
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        // first one raises wakeup
+        doThrow(new WakeupException())
+                // and succeed the second time
+                .doNothing()
+                .when(consumer).commitSync(eq(offsets));
+
+        workerTask.iteration(); // first record delivered
+
+        workerTask.iteration(); // now rebalance with the wakeup triggered
+        time.sleep(30000L);
+
+        verify(sinkTask).close(INITIAL_ASSIGNMENT);
+        verify(sinkTask, times(2)).open(INITIAL_ASSIGNMENT);
+
+        INITIAL_ASSIGNMENT.forEach(tp -> {
+            verify(consumer).resume(Collections.singleton(tp));
+        });
+
+        verify(statusListener).onResume(taskId);
+
+        assertSinkMetricValue("partition-count", 2);
+        assertSinkMetricValue("sink-record-read-total", 1.0);
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        assertSinkMetricValue("sink-record-active-count", 0.0);
+        assertSinkMetricValue("sink-record-active-count-max", 1.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.33333);
+        assertSinkMetricValue("offset-commit-seq-no", 1.0);
+        assertSinkMetricValue("offset-commit-completion-total", 1.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 1.0);
+        assertTaskMetricValue("batch-size-avg", 1.0);
+        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testWakeupNotThrownDuringShutdown() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(invocation -> {
+                    // stop the task during its second iteration
                     workerTask.stop();
-                    workerTask.cancel();
-                    return null;
-                })
-                .when(sinkTask).put(anyList());
+                    return new ConsumerRecords<>(Collections.emptyMap());
+                });
+        expectConversionAndTransformation(null, new RecordHeaders());
 
-        // task performs normal steps in advance of committing offsets
         final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
         offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
         when(sinkTask.preCommit(offsets)).thenReturn(offsets);
 
+        // fail the first time
+        doThrow(new WakeupException())
+                // and succeed the second time
+                .doNothing()
+                .when(consumer).commitSync(eq(offsets));
+
         workerTask.execute();
 
-        // stop wakes up the consumer
+        assertEquals(0, workerTask.commitFailures());
         verify(consumer).wakeup();
+        verify(sinkTask).close(any(Collection.class));
+    }
 
-        verify(sinkTask).close(any());
+    @Test
+    public void testRequestCommit() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(expectConsumerPoll(0));
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        // Initial assignment
+        time.sleep(30000L);
+        workerTask.iteration();
+        assertSinkMetricValue("partition-count", 2);
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        // First record delivered
+        workerTask.iteration();
+        assertSinkMetricValue("partition-count", 2);
+        assertSinkMetricValue("sink-record-read-total", 1.0);
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        assertSinkMetricValue("sink-record-active-count", 1.0);
+        assertSinkMetricValue("sink-record-active-count-max", 1.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.333333);
+        assertSinkMetricValue("offset-commit-seq-no", 0.0);
+        assertSinkMetricValue("offset-commit-completion-total", 0.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 1.0);
+        assertTaskMetricValue("batch-size-avg", 0.5);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
+
+        // Grab the commit time prior to requesting a commit.
+        // This time should advance slightly after committing.
+        // KAFKA-8229
+        final long previousCommitValue = workerTask.getNextCommit();
+        sinkTaskContext.getValue().requestCommit();
+        assertTrue(sinkTaskContext.getValue().isCommitRequested());
+        assertNotEquals(offsets, workerTask.lastCommittedOffsets());
+
+        ArgumentCaptor<OffsetCommitCallback> callback = ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        time.sleep(10000L);
+        workerTask.iteration(); // triggers the commit
+        verify(consumer).commitAsync(eq(offsets), callback.capture());
+        callback.getValue().onComplete(offsets, null);
+        time.sleep(10000L);
+
+        assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
+        assertEquals(offsets, workerTask.lastCommittedOffsets());
+        assertEquals(0, workerTask.commitFailures());
+
+        // Assert the next commit time advances slightly, the amount it advances
+        // is the normal commit time less the two sleeps since it started each
+        // of those sleeps were 10 seconds.
+        // KAFKA-8229
+        assertEquals("Should have only advanced by 40 seconds",
+                previousCommitValue  +
+                        (WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT - 10000L * 2),
+                workerTask.getNextCommit());
+
+        assertSinkMetricValue("partition-count", 2);
+        assertSinkMetricValue("sink-record-read-total", 1.0);
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        assertSinkMetricValue("sink-record-active-count", 0.0);
+        assertSinkMetricValue("sink-record-active-count-max", 1.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.2);
+        assertSinkMetricValue("offset-commit-seq-no", 1.0);
+        assertSinkMetricValue("offset-commit-completion-total", 1.0);
+        assertSinkMetricValue("offset-commit-skip-total", 0.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 1.0);
+        assertTaskMetricValue("batch-size-avg", 0.33333);
+        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
+    }
+
+    @Test
+    public void testPreCommit() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(2))
+                .thenAnswer(expectConsumerPoll(0));
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+
+        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
+        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
+        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        assertEquals(workerStartingOffsets, workerTask.currentOffsets());
+
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        final Map<TopicPartition, OffsetAndMetadata> taskOffsets = new HashMap<>();
+        taskOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); // act like FIRST_OFFSET+2 has not yet been flushed by the task
+        taskOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 1)); // should be ignored because > current offset
+        taskOffsets.put(new TopicPartition(TOPIC, 3), new OffsetAndMetadata(FIRST_OFFSET)); // should be ignored because this partition is not assigned
+
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(taskOffsets);
+
+        workerTask.iteration(); // iter 2 -- deliver 2 records
+
+        final Map<TopicPartition, OffsetAndMetadata> committableOffsets = new HashMap<>();
+        committableOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        committableOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        assertEquals(workerCurrentOffsets, workerTask.currentOffsets());
+        assertEquals(workerStartingOffsets, workerTask.lastCommittedOffsets());
+
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 3 -- commit
+
+        // Expect extra invalid topic partition to be filtered, which causes the consumer assignment to be logged
+        ArgumentCaptor<OffsetCommitCallback> callback = ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        verify(consumer).commitAsync(eq(committableOffsets), callback.capture());
+        callback.getValue().onComplete(committableOffsets, null);
+
+        assertEquals(committableOffsets, workerTask.lastCommittedOffsets());
+    }
+
+    @Test
+    public void testPreCommitFailure() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                // Put one message through the task to get some offsets to commit
+                .thenAnswer(expectConsumerPoll(2))
+                .thenAnswer(expectConsumerPoll(0));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+
+        workerTask.iteration(); // iter 2 -- deliver 2 records
+
+        // iter 3
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenThrow(new ConnectException("Failed to flush"));
+
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 3 -- commit
+
+        verify(consumer).seek(TOPIC_PARTITION, FIRST_OFFSET);
+        verify(consumer).seek(TOPIC_PARTITION2, FIRST_OFFSET);
+    }
+
+    @Test
+    public void testIgnoredCommit() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        // iter 1
+        expectPollInitialAssignment()
+                // iter 2
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(expectConsumerPoll(0));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+
+        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
+        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
+        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        assertEquals(workerStartingOffsets, workerTask.currentOffsets());
+        assertEquals(workerStartingOffsets, workerTask.lastCommittedOffsets());
+
+        workerTask.iteration(); // iter 2 -- deliver 2 records
+
+    // iter 3
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(workerStartingOffsets);
+
+        sinkTaskContext.getValue().requestCommit();
+        // no actual consumer.commit() triggered
+        workerTask.iteration(); // iter 3 -- commit
+    }
+
+    // Test that the commitTimeoutMs timestamp is correctly computed and checked in WorkerSinkTask.iteration()
+    // when there is a long running commit in process. See KAFKA-4942 for more information.
+    @Test
+    public void testLongRunningCommitWithoutTimeout() throws InterruptedException {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                .thenAnswer(expectConsumerPoll(1))
+                // no actual consumer.commit() triggered
+                .thenAnswer(expectConsumerPoll(0));
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
+        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
+        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        workerTask.iteration(); // iter 1 -- initial assignment
+        assertEquals(workerStartingOffsets, workerTask.currentOffsets());
+        assertEquals(workerStartingOffsets, workerTask.lastCommittedOffsets());
+
+        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
+        workerTask.iteration(); // iter 2 -- deliver 2 records
+
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        // iter 3 - note that we return the current offset to indicate they should be committed
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(workerCurrentOffsets);
+
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 3 -- commit in progress
+
+        // Make sure the "committing" flag didn't immediately get flipped back to false due to an incorrect timeout
+        assertTrue("Expected worker to be in the process of committing offsets", workerTask.isCommitting());
+
+        // Delay the result of trying to commit offsets to Kafka via the consumer.commitAsync method.
+        ArgumentCaptor<OffsetCommitCallback> offsetCommitCallbackArgumentCaptor =
+            ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        verify(consumer).commitAsync(eq(workerCurrentOffsets), offsetCommitCallbackArgumentCaptor.capture());
+
+        final OffsetCommitCallback callback = offsetCommitCallbackArgumentCaptor.getValue();
+        callback.onComplete(workerCurrentOffsets, null);
+
+        assertEquals(workerCurrentOffsets, workerTask.currentOffsets());
+        assertEquals(workerCurrentOffsets, workerTask.lastCommittedOffsets());
+        assertFalse(workerTask.isCommitting());
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testSinkTasksHandleCloseErrors() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                // Put one message through the task to get some offsets to commit
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        doNothing()
+                .doAnswer(invocation -> {
+                    workerTask.stop();
+                    return null;
+                })
+                .when(sinkTask).put(anyList());
+
+        Throwable closeException = new RuntimeException();
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+
+        // Throw another exception while closing the task's assignment
+        doThrow(closeException).when(sinkTask).close(any(Collection.class));
+
+        try {
+            workerTask.execute();
+            fail("workerTask.execute should have thrown an exception");
+        } catch (RuntimeException e) {
+            assertSame("Exception from close should propagate as-is", closeException, e);
+        }
+
+        verify(consumer).wakeup();
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testSuppressCloseErrors() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                // Put one message through the task to get some offsets to commit
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        Throwable putException = new RuntimeException();
+        Throwable closeException = new RuntimeException();
+
+        doNothing()
+                // Throw an exception on the next put to trigger shutdown behavior
+                // This exception is the true "cause" of the failure
+                .doThrow(putException)
+                .when(sinkTask).put(anyList());
+
+        when(sinkTask.preCommit(anyMap())).thenReturn(Collections.emptyMap());
+
+        // Throw another exception while closing the task's assignment
+        doThrow(closeException).when(sinkTask).close(any(Collection.class));
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        try {
+            workerTask.execute();
+            fail("workerTask.execute should have thrown an exception");
+        } catch (ConnectException e) {
+            assertSame("Exception from put should be the cause", putException, e.getCause());
+            assertTrue("Exception from close should be suppressed", e.getSuppressed().length > 0);
+            assertSame(closeException, e.getSuppressed()[0]);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testTaskCancelPreventsFinalOffsetCommit() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        expectTaskGetTopic();
+        expectPollInitialAssignment()
+                // Put one message through the task to get some offsets to commit
+                .thenAnswer(expectConsumerPoll(1))
+                // the second put will return after the task is stopped and cancelled (asynchronously)
+                .thenAnswer(expectConsumerPoll(1));
+
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        doNothing()
+                .doNothing()
+                .doAnswer(invocation -> {
+                    workerTask.stop();
+                    workerTask.cancel();
+                    return null;
+                })
+                .when(sinkTask).put(anyList());
+
+        // task performs normal steps in advance of committing offsets
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        when(sinkTask.preCommit(offsets)).thenReturn(offsets);
+
+        workerTask.execute();
+
+        // stop wakes up the consumer
+        verify(consumer).wakeup();
+
+        verify(sinkTask).close(any());
+    }
+
+    // Verify that when commitAsync is called but the supplied callback is not called by the consumer before a
+    // rebalance occurs, the async callback does not reset the last committed offset from the rebalance.
+    // See KAFKA-5731 for more information.
+    @Test
+    public void testCommitWithOutOfOrderCallback() {
+        createTask(initialState);
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        verifyInitializeTask();
+
+        // iter 1
+        Answer<ConsumerRecords<byte[], byte[]>> consumerPollRebalance = invocation -> {
+            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+            return ConsumerRecords.empty();
+        };
+
+        // iter 2
+        expectTaskGetTopic();
+        expectConversionAndTransformation(null, new RecordHeaders());
+
+        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
+        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
+        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
+        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+
+        final List<TopicPartition> originalPartitions = new ArrayList<>(INITIAL_ASSIGNMENT);
+        final List<TopicPartition> rebalancedPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3);
+        final Map<TopicPartition, OffsetAndMetadata> rebalanceOffsets = new HashMap<>();
+        rebalanceOffsets.put(TOPIC_PARTITION, workerCurrentOffsets.get(TOPIC_PARTITION));
+        rebalanceOffsets.put(TOPIC_PARTITION2, workerCurrentOffsets.get(TOPIC_PARTITION2));
+        rebalanceOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
+
+        final Map<TopicPartition, OffsetAndMetadata> postRebalanceCurrentOffsets = new HashMap<>();
+        postRebalanceCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 3));
+        postRebalanceCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        postRebalanceCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 2));
+
+        // iter 3 - note that we return the current offset to indicate they should be committed
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(workerCurrentOffsets);
+
+        // We need to delay the result of trying to commit offsets to Kafka via the consumer.commitAsync
+        // method. We do this so that we can test that the callback is not called until after the rebalance
+        // changes the lastCommittedOffsets. To fake this for tests we have the commitAsync build a function
+        // that will call the callback with the appropriate parameters, and we'll run that function later.
+        final AtomicReference<Runnable> asyncCallbackRunner = new AtomicReference<>();
+        final AtomicBoolean asyncCallbackRan = new AtomicBoolean();
+
+        doAnswer(invocation -> {
+            final Map<TopicPartition, OffsetAndMetadata> offsets = invocation.getArgument(0);
+            final OffsetCommitCallback callback =  invocation.getArgument(1);
+            asyncCallbackRunner.set(() -> {
+                callback.onComplete(offsets, null);
+                asyncCallbackRan.set(true);
+            });
+
+            return null;
+        }).when(consumer).commitAsync(eq(workerCurrentOffsets), any(OffsetCommitCallback.class));
+
+        // Expect the next poll to discover and perform the rebalance, THEN complete the previous callback handler,
+        // and then return one record for TP1 and one for TP3.
+        final AtomicBoolean rebalanced = new AtomicBoolean();
+        Answer<ConsumerRecords<byte[], byte[]>> consumerPollRebalanced = invocation -> {
+            // Rebalance always begins with revoking current partitions ...
+            rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
+            // Respond to the rebalance
+            Map<TopicPartition, Long> offsets = new HashMap<>();
+            offsets.put(TOPIC_PARTITION, rebalanceOffsets.get(TOPIC_PARTITION).offset());
+            offsets.put(TOPIC_PARTITION2, rebalanceOffsets.get(TOPIC_PARTITION2).offset());
+            offsets.put(TOPIC_PARTITION3, rebalanceOffsets.get(TOPIC_PARTITION3).offset());
+            sinkTaskContext.getValue().offset(offsets);
+            rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions);
+            rebalanced.set(true);
+
+            // Run the previous async commit handler
+            asyncCallbackRunner.get().run();
+
+            // And prep the two records to return
+            long timestamp = RecordBatch.NO_TIMESTAMP;
+            TimestampType timestampType = TimestampType.NO_TIMESTAMP_TYPE;
+            List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
+            records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, timestamp, timestampType,
+                    0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
+            records.add(new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET + recordsReturnedTp3 + 1, timestamp, timestampType,
+                    0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
+            recordsReturnedTp1 += 1;
+            recordsReturnedTp3 += 1;
+            return new ConsumerRecords<>(Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records));
+        };
+
+        // onPartitionsRevoked
+        when(sinkTask.preCommit(workerCurrentOffsets)).thenReturn(workerCurrentOffsets);
+
+        // onPartitionsAssigned - step 1
+        final long offsetTp1 = rebalanceOffsets.get(TOPIC_PARTITION).offset();
+        final long offsetTp2 = rebalanceOffsets.get(TOPIC_PARTITION2).offset();
+        final long offsetTp3 = rebalanceOffsets.get(TOPIC_PARTITION3).offset();
+
+        // iter 4 - note that we return the current offset to indicate they should be committed
+        when(sinkTask.preCommit(postRebalanceCurrentOffsets)).thenReturn(postRebalanceCurrentOffsets);
+
+        // Setup mocks
+        when(consumer.assignment())
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(INITIAL_ASSIGNMENT)
+                .thenReturn(new HashSet<>(rebalancedPartitions))
+                .thenReturn(new HashSet<>(rebalancedPartitions))
+                .thenReturn(new HashSet<>(rebalancedPartitions))
+                .thenReturn(new HashSet<>(rebalancedPartitions))
+                .thenReturn(new HashSet<>(rebalancedPartitions));
+
+        when(consumer.position(TOPIC_PARTITION))
+                .thenReturn(FIRST_OFFSET)
+                .thenReturn(offsetTp1);
+
+        when(consumer.position(TOPIC_PARTITION2))
+                .thenReturn(FIRST_OFFSET)
+                .thenReturn(offsetTp2);
+
+        when(consumer.position(TOPIC_PARTITION3))
+                .thenReturn(offsetTp3);
+
+        when(consumer.poll(any(Duration.class)))
+                .thenAnswer(consumerPollRebalance)
+                .thenAnswer(expectConsumerPoll(1))
+                .thenAnswer(consumerPollRebalanced)
+                .thenAnswer(expectConsumerPoll(1));
+
+        // Run the iterations
+        workerTask.iteration(); // iter 1 -- initial assignment
+
+        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
+        workerTask.iteration(); // iter 2 -- deliver records
+
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 3 -- commit in progress
+
+        assertSinkMetricValue("partition-count", 3);
+        assertSinkMetricValue("sink-record-read-total", 3.0);
+        assertSinkMetricValue("sink-record-send-total", 3.0);
+        assertSinkMetricValue("sink-record-active-count", 4.0);
+        assertSinkMetricValue("sink-record-active-count-max", 4.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.71429);
+        assertSinkMetricValue("offset-commit-seq-no", 2.0);
+        assertSinkMetricValue("offset-commit-completion-total", 1.0);
+        assertSinkMetricValue("offset-commit-skip-total", 1.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 2.0);
+        assertTaskMetricValue("batch-size-avg", 1.0);
+        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
+
+        assertTrue(asyncCallbackRan.get());
+        assertTrue(rebalanced.get());
+
+        // Check that the offsets were not reset by the out-of-order async commit callback
+        assertEquals(postRebalanceCurrentOffsets, workerTask.currentOffsets());
+        assertEquals(rebalanceOffsets, workerTask.lastCommittedOffsets());
+
+        // onPartitionsRevoked
+        verify(sinkTask).close(new ArrayList<>(workerCurrentOffsets.keySet()));
+        verify(consumer).commitSync(anyMap());
+
+        // onPartitionsAssigned - step 2
+        verify(sinkTask).open(rebalancedPartitions);
+
+        // onPartitionsAssigned - step 3 rewind
+        verify(consumer).seek(TOPIC_PARTITION, offsetTp1);
+        verify(consumer).seek(TOPIC_PARTITION2, offsetTp2);
+        verify(consumer).seek(TOPIC_PARTITION3, offsetTp3);
+
+        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
+        sinkTaskContext.getValue().requestCommit();
+        workerTask.iteration(); // iter 4 -- commit in progress
+
+        final ArgumentCaptor<OffsetCommitCallback> callback = ArgumentCaptor.forClass(OffsetCommitCallback.class);
+        verify(consumer).commitAsync(eq(postRebalanceCurrentOffsets), callback.capture());
+        callback.getValue().onComplete(postRebalanceCurrentOffsets, null);
+
+        // Check that the offsets were not reset by the out-of-order async commit callback
+        assertEquals(postRebalanceCurrentOffsets, workerTask.currentOffsets());
+        assertEquals(postRebalanceCurrentOffsets, workerTask.lastCommittedOffsets());
+
+        assertSinkMetricValue("partition-count", 3);
+        assertSinkMetricValue("sink-record-read-total", 4.0);
+        assertSinkMetricValue("sink-record-send-total", 4.0);
+        assertSinkMetricValue("sink-record-active-count", 0.0);
+        assertSinkMetricValue("sink-record-active-count-max", 4.0);
+        assertSinkMetricValue("sink-record-active-count-avg", 0.5555555);
+        assertSinkMetricValue("offset-commit-seq-no", 3.0);
+        assertSinkMetricValue("offset-commit-completion-total", 2.0);
+        assertSinkMetricValue("offset-commit-skip-total", 1.0);
+        assertTaskMetricValue("status", "running");
+        assertTaskMetricValue("running-ratio", 1.0);
+        assertTaskMetricValue("pause-ratio", 0.0);
+        assertTaskMetricValue("batch-size-max", 2.0);
+        assertTaskMetricValue("batch-size-avg", 1.0);
+        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
+        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
+        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
     }
 
     @Test
@@ -1081,4 +1918,4 @@ private void assertTaskMetricValue(String name, String expected) {
         String measured = metrics.currentMetricValueAsString(taskGroup, name);
         assertEquals(expected, measured);
     }
-}
+}
\ No newline at end of file
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
deleted file mode 100644
index e103c30157b9d..0000000000000
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
+++ /dev/null
@@ -1,1298 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.connect.runtime;
-
-import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.consumer.ConsumerRecords;
-import org.apache.kafka.clients.consumer.KafkaConsumer;
-import org.apache.kafka.clients.consumer.OffsetAndMetadata;
-import org.apache.kafka.clients.consumer.OffsetCommitCallback;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.WakeupException;
-import org.apache.kafka.common.header.Header;
-import org.apache.kafka.common.header.Headers;
-import org.apache.kafka.common.header.internals.RecordHeaders;
-import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.record.TimestampType;
-import org.apache.kafka.common.utils.MockTime;
-import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.connect.data.Schema;
-import org.apache.kafka.connect.data.SchemaAndValue;
-import org.apache.kafka.connect.errors.ConnectException;
-import org.apache.kafka.connect.errors.RetriableException;
-import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
-import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
-import org.apache.kafka.connect.runtime.errors.ErrorReporter;
-import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
-import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
-import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
-import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
-import org.apache.kafka.connect.sink.SinkConnector;
-import org.apache.kafka.connect.sink.SinkRecord;
-import org.apache.kafka.connect.sink.SinkTask;
-import org.apache.kafka.connect.storage.ClusterConfigState;
-import org.apache.kafka.connect.storage.Converter;
-import org.apache.kafka.connect.storage.HeaderConverter;
-import org.apache.kafka.connect.storage.StatusBackingStore;
-import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.easymock.Capture;
-import org.easymock.EasyMock;
-import org.easymock.IExpectationSetters;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.powermock.api.easymock.PowerMock;
-import org.powermock.api.easymock.annotation.Mock;
-import org.powermock.core.classloader.annotations.PowerMockIgnore;
-import org.powermock.core.classloader.annotations.PrepareForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
-import org.powermock.reflect.Whitebox;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Function;
-import java.util.function.Supplier;
-import java.util.stream.Collectors;
-
-import static java.util.Arrays.asList;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
-
-@RunWith(PowerMockRunner.class)
-@PrepareForTest(WorkerSinkTask.class)
-@PowerMockIgnore("javax.management.*")
-public class WorkerSinkTaskTest {
-    // These are fixed to keep this code simpler. In this example we assume byte[] raw values
-    // with mix of integer/string in Connect
-    private static final String TOPIC = "test";
-    private static final int PARTITION = 12;
-    private static final int PARTITION2 = 13;
-    private static final int PARTITION3 = 14;
-    private static final long FIRST_OFFSET = 45;
-    private static final Schema KEY_SCHEMA = Schema.INT32_SCHEMA;
-    private static final int KEY = 12;
-    private static final Schema VALUE_SCHEMA = Schema.STRING_SCHEMA;
-    private static final String VALUE = "VALUE";
-    private static final byte[] RAW_KEY = "key".getBytes();
-    private static final byte[] RAW_VALUE = "value".getBytes();
-
-    private static final TopicPartition TOPIC_PARTITION = new TopicPartition(TOPIC, PARTITION);
-    private static final TopicPartition TOPIC_PARTITION2 = new TopicPartition(TOPIC, PARTITION2);
-    private static final TopicPartition TOPIC_PARTITION3 = new TopicPartition(TOPIC, PARTITION3);
-
-    private static final Set<TopicPartition> INITIAL_ASSIGNMENT =
-        new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2));
-
-    private static final Map<String, String> TASK_PROPS = new HashMap<>();
-    static {
-        TASK_PROPS.put(SinkConnector.TOPICS_CONFIG, TOPIC);
-        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, TestSinkTask.class.getName());
-    }
-    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
-
-    private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
-    private TargetState initialState = TargetState.STARTED;
-    private MockTime time;
-    private WorkerSinkTask workerTask;
-    @Mock
-    private SinkTask sinkTask;
-    private Capture<WorkerSinkTaskContext> sinkTaskContext = EasyMock.newCapture();
-    private WorkerConfig workerConfig;
-    private MockConnectMetrics metrics;
-    @Mock
-    private PluginClassLoader pluginLoader;
-    @Mock
-    private Converter keyConverter;
-    @Mock
-    private Converter valueConverter;
-    @Mock
-    private HeaderConverter headerConverter;
-    @Mock
-    private TransformationChain<ConsumerRecord<byte[], byte[]>, SinkRecord> transformationChain;
-    @Mock
-    private TaskStatus.Listener statusListener;
-    @Mock
-    private StatusBackingStore statusBackingStore;
-    @Mock
-    private KafkaConsumer<byte[], byte[]> consumer;
-    @Mock
-    private ErrorHandlingMetrics errorHandlingMetrics;
-    private Capture<ConsumerRebalanceListener> rebalanceListener = EasyMock.newCapture();
-
-    private long recordsReturnedTp1;
-    private long recordsReturnedTp3;
-
-    @Before
-    public void setUp() {
-        time = new MockTime();
-        Map<String, String> workerProps = new HashMap<>();
-        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
-        workerConfig = new StandaloneConfig(workerProps);
-        pluginLoader = PowerMock.createMock(PluginClassLoader.class);
-        metrics = new MockConnectMetrics(time);
-        recordsReturnedTp1 = 0;
-        recordsReturnedTp3 = 0;
-    }
-
-    private void createTask(TargetState initialState) {
-        createTask(initialState, keyConverter, valueConverter, headerConverter);
-    }
-
-    private void createTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter) {
-        createTask(initialState, keyConverter, valueConverter, headerConverter, RetryWithToleranceOperatorTest.noopOperator(), Collections::emptyList);
-    }
-
-    private void createTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter,
-                            RetryWithToleranceOperator<ConsumerRecord<byte[], byte[]>> retryWithToleranceOperator, Supplier<List<ErrorReporter<ConsumerRecord<byte[], byte[]>>>> errorReportersSupplier) {
-        workerTask = new WorkerSinkTask(
-                taskId, sinkTask, statusListener, initialState, workerConfig, ClusterConfigState.EMPTY, metrics,
-                keyConverter, valueConverter, errorHandlingMetrics, headerConverter,
-                transformationChain, consumer, pluginLoader, time,
-                retryWithToleranceOperator, null, statusBackingStore, errorReportersSupplier);
-    }
-
-    @After
-    public void tearDown() {
-        if (metrics != null) metrics.stop();
-    }
-
-    @Test
-    public void testPollRedeliveryWithConsumerRebalance() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall().andThrow(new RetriableException("retry"));
-        // Pause
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT);
-        consumer.pause(INITIAL_ASSIGNMENT);
-        PowerMock.expectLastCall();
-
-        // Empty consumer poll (all partitions are paused) with rebalance; one new partition is assigned
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
-                rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
-                return ConsumerRecords.empty();
-            });
-        Set<TopicPartition> newAssignment = new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3));
-        EasyMock.expect(consumer.assignment()).andReturn(newAssignment).times(3);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET);
-        sinkTask.open(Collections.singleton(TOPIC_PARTITION3));
-        EasyMock.expectLastCall();
-        // All partitions are re-paused in order to pause any newly-assigned partitions so that redelivery efforts can continue
-        consumer.pause(newAssignment);
-        EasyMock.expectLastCall();
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall().andThrow(new RetriableException("retry"));
-
-        // Next delivery attempt fails again
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall().andThrow(new RetriableException("retry"));
-
-        // Non-empty consumer poll; all initially-assigned partitions are revoked in rebalance, and new partitions are allowed to resume
-        ConsumerRecord<byte[], byte[]> newRecord = new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET, RAW_KEY, RAW_VALUE);
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                rebalanceListener.getValue().onPartitionsAssigned(Collections.emptyList());
-                return new ConsumerRecords<>(Collections.singletonMap(TOPIC_PARTITION3, Collections.singletonList(newRecord)));
-            });
-        newAssignment = Collections.singleton(TOPIC_PARTITION3);
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(newAssignment)).times(3);
-        final Map<TopicPartition, OffsetAndMetadata> offsets = INITIAL_ASSIGNMENT.stream()
-                .collect(Collectors.toMap(Function.identity(), tp -> new OffsetAndMetadata(FIRST_OFFSET)));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-        sinkTask.close(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall();
-        // All partitions are resumed, as all previously paused-for-redelivery partitions were revoked
-        newAssignment.forEach(tp -> {
-            consumer.resume(Collections.singleton(tp));
-            EasyMock.expectLastCall();
-        });
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration();
-        workerTask.iteration();
-        workerTask.iteration();
-        workerTask.iteration();
-        workerTask.iteration();
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPreCommitFailureAfterPartialRevocationAndAssignment() throws Exception {
-        createTask(initialState);
-
-        // First poll; assignment is [TP1, TP2]
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        // Second poll; a single record is delivered from TP1
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        // Third poll; assignment changes to [TP2]
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-                () -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION));
-                    rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet());
-                    return ConsumerRecords.empty();
-                });
-        EasyMock.expect(consumer.assignment()).andReturn(Collections.singleton(TOPIC_PARTITION)).times(2);
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-        consumer.commitSync(offsets);
-        EasyMock.expectLastCall();
-        sinkTask.close(Collections.singleton(TOPIC_PARTITION));
-        EasyMock.expectLastCall();
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        // Fourth poll; assignment changes to [TP2, TP3]
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-                () -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet());
-                    rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3));
-                    return ConsumerRecords.empty();
-                });
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET);
-        sinkTask.open(Collections.singleton(TOPIC_PARTITION3));
-        EasyMock.expectLastCall();
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        // Fifth poll; an offset commit takes place
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2);
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        workerCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andThrow(new ConnectException("Failed to flush"));
-
-        consumer.seek(TOPIC_PARTITION2, FIRST_OFFSET);
-        EasyMock.expectLastCall();
-        consumer.seek(TOPIC_PARTITION3, FIRST_OFFSET);
-        EasyMock.expectLastCall();
-
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.eq(Collections.emptyList()));
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        // First iteration--first call to poll, first consumer assignment
-        workerTask.iteration();
-        // Second iteration--second call to poll, delivery of one record
-        workerTask.iteration();
-        // Third iteration--third call to poll, partial consumer revocation
-        workerTask.iteration();
-        // Fourth iteration--fourth call to poll, partial consumer assignment
-        workerTask.iteration();
-        // Fifth iteration--task-requested offset commit with failure in SinkTask::preCommit
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration();
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testWakeupInCommitSyncCausesRetry() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-
-        // first one raises wakeup
-        consumer.commitSync(EasyMock.<Map<TopicPartition, OffsetAndMetadata>>anyObject());
-        EasyMock.expectLastCall().andThrow(new WakeupException());
-
-        // we should retry and complete the commit
-        consumer.commitSync(EasyMock.<Map<TopicPartition, OffsetAndMetadata>>anyObject());
-        EasyMock.expectLastCall();
-
-        sinkTask.close(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall();
-
-        INITIAL_ASSIGNMENT.forEach(tp -> EasyMock.expect(consumer.position(tp)).andReturn(FIRST_OFFSET));
-
-        sinkTask.open(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(5);
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
-                return ConsumerRecords.empty();
-            });
-
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            consumer.resume(Collections.singleton(tp));
-            EasyMock.expectLastCall();
-        });
-
-        statusListener.onResume(taskId);
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        time.sleep(30000L);
-        workerTask.initializeAndStart();
-        time.sleep(30000L);
-
-        workerTask.iteration(); // poll for initial assignment
-        time.sleep(30000L);
-        workerTask.iteration(); // first record delivered
-        workerTask.iteration(); // now rebalance with the wakeup triggered
-        time.sleep(30000L);
-
-        assertSinkMetricValue("partition-count", 2);
-        assertSinkMetricValue("sink-record-read-total", 1.0);
-        assertSinkMetricValue("sink-record-send-total", 1.0);
-        assertSinkMetricValue("sink-record-active-count", 0.0);
-        assertSinkMetricValue("sink-record-active-count-max", 1.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.33333);
-        assertSinkMetricValue("offset-commit-seq-no", 1.0);
-        assertSinkMetricValue("offset-commit-completion-total", 1.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 1.0);
-        assertTaskMetricValue("batch-size-avg", 1.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testWakeupNotThrownDuringShutdown() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(() -> {
-            // stop the task during its second iteration
-            workerTask.stop();
-            return new ConsumerRecords<>(Collections.emptyMap());
-        });
-        consumer.wakeup();
-        EasyMock.expectLastCall();
-
-        sinkTask.put(EasyMock.eq(Collections.emptyList()));
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(1);
-
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-
-        sinkTask.close(EasyMock.anyObject());
-        PowerMock.expectLastCall();
-
-        // fail the first time
-        consumer.commitSync(EasyMock.eq(offsets));
-        EasyMock.expectLastCall().andThrow(new WakeupException());
-
-        // and succeed the second time
-        consumer.commitSync(EasyMock.eq(offsets));
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.execute();
-
-        assertEquals(0, workerTask.commitFailures());
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testRequestCommit() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        expectPollInitialAssignment();
-
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(offsets);
-        EasyMock.expectLastCall().andReturn(offsets);
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-
-        final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
-        consumer.commitAsync(EasyMock.eq(offsets), EasyMock.capture(callback));
-        EasyMock.expectLastCall().andAnswer(() -> {
-            callback.getValue().onComplete(offsets, null);
-            return null;
-        });
-
-        expectConsumerPoll(0);
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-
-        // Initial assignment
-        time.sleep(30000L);
-        workerTask.iteration();
-        assertSinkMetricValue("partition-count", 2);
-
-        // First record delivered
-        workerTask.iteration();
-        assertSinkMetricValue("partition-count", 2);
-        assertSinkMetricValue("sink-record-read-total", 1.0);
-        assertSinkMetricValue("sink-record-send-total", 1.0);
-        assertSinkMetricValue("sink-record-active-count", 1.0);
-        assertSinkMetricValue("sink-record-active-count-max", 1.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.333333);
-        assertSinkMetricValue("offset-commit-seq-no", 0.0);
-        assertSinkMetricValue("offset-commit-completion-total", 0.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 1.0);
-        assertTaskMetricValue("batch-size-avg", 0.5);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 0.0);
-
-        // Grab the commit time prior to requesting a commit.
-        // This time should advance slightly after committing.
-        // KAFKA-8229
-        final long previousCommitValue = workerTask.getNextCommit();
-        sinkTaskContext.getValue().requestCommit();
-        assertTrue(sinkTaskContext.getValue().isCommitRequested());
-        assertNotEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-        time.sleep(10000L);
-        workerTask.iteration(); // triggers the commit
-        time.sleep(10000L);
-        assertFalse(sinkTaskContext.getValue().isCommitRequested()); // should have been cleared
-        assertEquals(offsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-        assertEquals(0, workerTask.commitFailures());
-        // Assert the next commit time advances slightly, the amount it advances
-        // is the normal commit time less the two sleeps since it started each
-        // of those sleeps were 10 seconds.
-        // KAFKA-8229
-        assertEquals("Should have only advanced by 40 seconds",
-                     previousCommitValue  +
-                     (WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT - 10000L * 2),
-                     workerTask.getNextCommit());
-
-        assertSinkMetricValue("partition-count", 2);
-        assertSinkMetricValue("sink-record-read-total", 1.0);
-        assertSinkMetricValue("sink-record-send-total", 1.0);
-        assertSinkMetricValue("sink-record-active-count", 0.0);
-        assertSinkMetricValue("sink-record-active-count-max", 1.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.2);
-        assertSinkMetricValue("offset-commit-seq-no", 1.0);
-        assertSinkMetricValue("offset-commit-completion-total", 1.0);
-        assertSinkMetricValue("offset-commit-skip-total", 0.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 1.0);
-        assertTaskMetricValue("batch-size-avg", 0.33333);
-        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPreCommit() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        // iter 1
-        expectPollInitialAssignment();
-
-        // iter 2
-        expectConsumerPoll(2);
-        expectConversionAndTransformation(2);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
-        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> taskOffsets = new HashMap<>();
-        taskOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); // act like FIRST_OFFSET+2 has not yet been flushed by the task
-        taskOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 1)); // should be ignored because > current offset
-        taskOffsets.put(new TopicPartition(TOPIC, 3), new OffsetAndMetadata(FIRST_OFFSET)); // should be ignored because this partition is not assigned
-
-        final Map<TopicPartition, OffsetAndMetadata> committableOffsets = new HashMap<>();
-        committableOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        committableOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(taskOffsets);
-        // Expect extra invalid topic partition to be filtered, which causes the consumer assignment to be logged
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-        final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
-        consumer.commitAsync(EasyMock.eq(committableOffsets), EasyMock.capture(callback));
-        EasyMock.expectLastCall().andAnswer(() -> {
-            callback.getValue().onComplete(committableOffsets, null);
-            return null;
-        });
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
-        workerTask.iteration(); // iter 2 -- deliver 2 records
-
-        assertEquals(workerCurrentOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 3 -- commit
-        assertEquals(committableOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPreCommitFailure() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-        EasyMock.expect(consumer.assignment()).andStubReturn(INITIAL_ASSIGNMENT);
-
-        // iter 1
-        expectPollInitialAssignment();
-
-        // iter 2
-        expectConsumerPoll(2);
-        expectConversionAndTransformation(2);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        // iter 3
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andThrow(new ConnectException("Failed to flush"));
-
-        consumer.seek(TOPIC_PARTITION, FIRST_OFFSET);
-        EasyMock.expectLastCall();
-        consumer.seek(TOPIC_PARTITION2, FIRST_OFFSET);
-        EasyMock.expectLastCall();
-
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.eq(Collections.emptyList()));
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-        workerTask.iteration(); // iter 2 -- deliver 2 records
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 3 -- commit
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testIgnoredCommit() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        // iter 1
-        expectPollInitialAssignment();
-
-        // iter 2
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
-        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-
-        // iter 3
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(workerStartingOffsets);
-        // no actual consumer.commit() triggered
-        expectConsumerPoll(0);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-
-        workerTask.iteration(); // iter 2 -- deliver 2 records
-
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 3 -- commit
-
-        PowerMock.verifyAll();
-    }
-
-    // Test that the commitTimeoutMs timestamp is correctly computed and checked in WorkerSinkTask.iteration()
-    // when there is a long running commit in process. See KAFKA-4942 for more information.
-    @Test
-    public void testLongRunningCommitWithoutTimeout() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        // iter 1
-        expectPollInitialAssignment();
-
-        // iter 2
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
-        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-
-        // iter 3 - note that we return the current offset to indicate they should be committed
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
-
-        // We need to delay the result of trying to commit offsets to Kafka via the consumer.commitAsync
-        // method. We do this so that we can test that we do not erroneously mark a commit as timed out
-        // while it is still running and under time. To fake this for tests we have the commit run in a
-        // separate thread and wait for a latch which we control back in the main thread.
-        final ExecutorService executor = Executors.newSingleThreadExecutor();
-        final CountDownLatch latch = new CountDownLatch(1);
-
-        consumer.commitAsync(EasyMock.eq(workerCurrentOffsets), EasyMock.anyObject());
-        EasyMock.expectLastCall().andAnswer(() -> {
-            // Grab the arguments passed to the consumer.commitAsync method
-            final Object[] args = EasyMock.getCurrentArguments();
-            @SuppressWarnings("unchecked")
-            final Map<TopicPartition, OffsetAndMetadata> offsets = (Map<TopicPartition, OffsetAndMetadata>) args[0];
-            final OffsetCommitCallback callback = (OffsetCommitCallback) args[1];
-
-            executor.execute(() -> {
-                try {
-                    latch.await();
-                } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                }
-
-                callback.onComplete(offsets, null);
-            });
-
-            return null;
-        });
-
-        // no actual consumer.commit() triggered
-        expectConsumerPoll(0);
-
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
-        assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-
-        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
-        workerTask.iteration(); // iter 2 -- deliver 2 records
-
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 3 -- commit in progress
-
-        // Make sure the "committing" flag didn't immediately get flipped back to false due to an incorrect timeout
-        assertTrue("Expected worker to be in the process of committing offsets", workerTask.isCommitting());
-
-        // Let the async commit finish and wait for it to end
-        latch.countDown();
-        executor.shutdown();
-        executor.awaitTermination(30, TimeUnit.SECONDS);
-
-        assertEquals(workerCurrentOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
-        assertEquals(workerCurrentOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testSinkTasksHandleCloseErrors() throws Exception {
-        createTask(initialState);
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        expectPollInitialAssignment();
-
-        // Put one message through the task to get some offsets to commit
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall().andVoid();
-
-        // Stop the task during the next put
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall().andAnswer(() -> {
-            workerTask.stop();
-            return null;
-        });
-
-        consumer.wakeup();
-        PowerMock.expectLastCall();
-
-        // Throw another exception while closing the task's assignment
-        EasyMock.expect(sinkTask.preCommit(EasyMock.anyObject()))
-            .andStubReturn(Collections.emptyMap());
-        Throwable closeException = new RuntimeException();
-        sinkTask.close(EasyMock.anyObject());
-        PowerMock.expectLastCall().andThrow(closeException);
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        try {
-            workerTask.execute();
-            fail("workerTask.execute should have thrown an exception");
-        } catch (RuntimeException e) {
-            PowerMock.verifyAll();
-            assertSame("Exception from close should propagate as-is", closeException, e);
-        }
-    }
-
-    @Test
-    public void testSuppressCloseErrors() throws Exception {
-        createTask(initialState);
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        expectPollInitialAssignment();
-
-        // Put one message through the task to get some offsets to commit
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall().andVoid();
-
-        // Throw an exception on the next put to trigger shutdown behavior
-        // This exception is the true "cause" of the failure
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(1);
-        Throwable putException = new RuntimeException();
-        sinkTask.put(EasyMock.anyObject());
-        PowerMock.expectLastCall().andThrow(putException);
-
-        // Throw another exception while closing the task's assignment
-        EasyMock.expect(sinkTask.preCommit(EasyMock.anyObject()))
-            .andStubReturn(Collections.emptyMap());
-        Throwable closeException = new RuntimeException();
-        sinkTask.close(EasyMock.anyObject());
-        PowerMock.expectLastCall().andThrow(closeException);
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        try {
-            workerTask.execute();
-            fail("workerTask.execute should have thrown an exception");
-        } catch (ConnectException e) {
-            PowerMock.verifyAll();
-            assertSame("Exception from put should be the cause", putException, e.getCause());
-            assertTrue("Exception from close should be suppressed", e.getSuppressed().length > 0);
-            assertSame(closeException, e.getSuppressed()[0]);
-        }
-    }
-
-    // Verify that when commitAsync is called but the supplied callback is not called by the consumer before a
-    // rebalance occurs, the async callback does not reset the last committed offset from the rebalance.
-    // See KAFKA-5731 for more information.
-    @Test
-    public void testCommitWithOutOfOrderCallback() throws Exception {
-        createTask(initialState);
-
-        expectInitializeTask();
-        expectTaskGetTopic(true);
-
-        // iter 1
-        expectPollInitialAssignment();
-
-        // iter 2
-        expectConsumerPoll(1);
-        expectConversionAndTransformation(4);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
-        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
-        workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
-        workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final List<TopicPartition> originalPartitions = new ArrayList<>(INITIAL_ASSIGNMENT);
-        final List<TopicPartition> rebalancedPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3);
-        final Map<TopicPartition, OffsetAndMetadata> rebalanceOffsets = new HashMap<>();
-        rebalanceOffsets.put(TOPIC_PARTITION, workerCurrentOffsets.get(TOPIC_PARTITION));
-        rebalanceOffsets.put(TOPIC_PARTITION2, workerCurrentOffsets.get(TOPIC_PARTITION2));
-        rebalanceOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
-
-        final Map<TopicPartition, OffsetAndMetadata> postRebalanceCurrentOffsets = new HashMap<>();
-        postRebalanceCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 3));
-        postRebalanceCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-        postRebalanceCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 2));
-
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(originalPartitions)).times(2);
-
-        // iter 3 - note that we return the current offset to indicate they should be committed
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
-
-        // We need to delay the result of trying to commit offsets to Kafka via the consumer.commitAsync
-        // method. We do this so that we can test that the callback is not called until after the rebalance
-        // changes the lastCommittedOffsets. To fake this for tests we have the commitAsync build a function
-        // that will call the callback with the appropriate parameters, and we'll run that function later.
-        final AtomicReference<Runnable> asyncCallbackRunner = new AtomicReference<>();
-        final AtomicBoolean asyncCallbackRan = new AtomicBoolean();
-
-        consumer.commitAsync(EasyMock.eq(workerCurrentOffsets), EasyMock.anyObject());
-        EasyMock.expectLastCall().andAnswer(() -> {
-            // Grab the arguments passed to the consumer.commitAsync method
-            final Object[] args = EasyMock.getCurrentArguments();
-            @SuppressWarnings("unchecked")
-            final Map<TopicPartition, OffsetAndMetadata> offsets = (Map<TopicPartition, OffsetAndMetadata>) args[0];
-            final OffsetCommitCallback callback = (OffsetCommitCallback) args[1];
-            asyncCallbackRunner.set(() -> {
-                callback.onComplete(offsets, null);
-                asyncCallbackRan.set(true);
-            });
-            return null;
-        });
-
-        // Expect the next poll to discover and perform the rebalance, THEN complete the previous callback handler,
-        // and then return one record for TP1 and one for TP3.
-        final AtomicBoolean rebalanced = new AtomicBoolean();
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                // Rebalance always begins with revoking current partitions ...
-                rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
-                // Respond to the rebalance
-                Map<TopicPartition, Long> offsets = new HashMap<>();
-                offsets.put(TOPIC_PARTITION, rebalanceOffsets.get(TOPIC_PARTITION).offset());
-                offsets.put(TOPIC_PARTITION2, rebalanceOffsets.get(TOPIC_PARTITION2).offset());
-                offsets.put(TOPIC_PARTITION3, rebalanceOffsets.get(TOPIC_PARTITION3).offset());
-                sinkTaskContext.getValue().offset(offsets);
-                rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions);
-                rebalanced.set(true);
-
-                // Run the previous async commit handler
-                asyncCallbackRunner.get().run();
-
-                 // And prep the two records to return
-                long timestamp = RecordBatch.NO_TIMESTAMP;
-                TimestampType timestampType = TimestampType.NO_TIMESTAMP_TYPE;
-                List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
-                records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, timestamp, timestampType,
-                    0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
-                records.add(new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET + recordsReturnedTp3 + 1, timestamp, timestampType,
-                    0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
-                recordsReturnedTp1 += 1;
-                recordsReturnedTp3 += 1;
-                return new ConsumerRecords<>(Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records));
-            });
-
-        // onPartitionsRevoked
-        sinkTask.preCommit(workerCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-        sinkTask.close(new ArrayList<>(workerCurrentOffsets.keySet()));
-        EasyMock.expectLastCall();
-        consumer.commitSync(workerCurrentOffsets);
-        EasyMock.expectLastCall();
-
-        // onPartitionsAssigned - step 1
-        final long offsetTp1 = rebalanceOffsets.get(TOPIC_PARTITION).offset();
-        final long offsetTp2 = rebalanceOffsets.get(TOPIC_PARTITION2).offset();
-        final long offsetTp3 = rebalanceOffsets.get(TOPIC_PARTITION3).offset();
-        EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(offsetTp1);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(offsetTp2);
-        EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(offsetTp3);
-        EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(rebalancedPartitions)).times(5);
-
-        // onPartitionsAssigned - step 2
-        sinkTask.open(EasyMock.eq(rebalancedPartitions));
-        EasyMock.expectLastCall();
-
-        // onPartitionsAssigned - step 3 rewind
-        consumer.seek(TOPIC_PARTITION, offsetTp1);
-        EasyMock.expectLastCall();
-        consumer.seek(TOPIC_PARTITION2, offsetTp2);
-        EasyMock.expectLastCall();
-        consumer.seek(TOPIC_PARTITION3, offsetTp3);
-        EasyMock.expectLastCall();
-
-        // iter 4 - note that we return the current offset to indicate they should be committed
-        sinkTask.preCommit(postRebalanceCurrentOffsets);
-        EasyMock.expectLastCall().andReturn(postRebalanceCurrentOffsets);
-
-        final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
-        consumer.commitAsync(EasyMock.eq(postRebalanceCurrentOffsets), EasyMock.capture(callback));
-        EasyMock.expectLastCall().andAnswer(() -> {
-            callback.getValue().onComplete(postRebalanceCurrentOffsets, null);
-            return null;
-        });
-
-        // no actual consumer.commit() triggered
-        expectConsumerPoll(1);
-
-        sinkTask.put(EasyMock.anyObject());
-        EasyMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        workerTask.initialize(TASK_CONFIG);
-        workerTask.initializeAndStart();
-        workerTask.iteration(); // iter 1 -- initial assignment
-
-        assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
-        assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
-
-        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
-        workerTask.iteration(); // iter 2 -- deliver 2 records
-
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 3 -- commit in progress
-
-        assertSinkMetricValue("partition-count", 3);
-        assertSinkMetricValue("sink-record-read-total", 3.0);
-        assertSinkMetricValue("sink-record-send-total", 3.0);
-        assertSinkMetricValue("sink-record-active-count", 4.0);
-        assertSinkMetricValue("sink-record-active-count-max", 4.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.71429);
-        assertSinkMetricValue("offset-commit-seq-no", 2.0);
-        assertSinkMetricValue("offset-commit-completion-total", 1.0);
-        assertSinkMetricValue("offset-commit-skip-total", 1.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 2.0);
-        assertTaskMetricValue("batch-size-avg", 1.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
-
-        assertTrue(asyncCallbackRan.get());
-        assertTrue(rebalanced.get());
-
-        // Check that the offsets were not reset by the out-of-order async commit callback
-        assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
-        assertEquals(rebalanceOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
-
-        time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
-        sinkTaskContext.getValue().requestCommit();
-        workerTask.iteration(); // iter 4 -- commit in progress
-
-        // Check that the offsets were not reset by the out-of-order async commit callback
-        assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
-        assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
-
-        assertSinkMetricValue("partition-count", 3);
-        assertSinkMetricValue("sink-record-read-total", 4.0);
-        assertSinkMetricValue("sink-record-send-total", 4.0);
-        assertSinkMetricValue("sink-record-active-count", 0.0);
-        assertSinkMetricValue("sink-record-active-count-max", 4.0);
-        assertSinkMetricValue("sink-record-active-count-avg", 0.5555555);
-        assertSinkMetricValue("offset-commit-seq-no", 3.0);
-        assertSinkMetricValue("offset-commit-completion-total", 2.0);
-        assertSinkMetricValue("offset-commit-skip-total", 1.0);
-        assertTaskMetricValue("status", "running");
-        assertTaskMetricValue("running-ratio", 1.0);
-        assertTaskMetricValue("pause-ratio", 0.0);
-        assertTaskMetricValue("batch-size-max", 2.0);
-        assertTaskMetricValue("batch-size-avg", 1.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
-        assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
-        assertTaskMetricValue("offset-commit-success-percentage", 1.0);
-
-        PowerMock.verifyAll();
-    }
-
-    private void expectInitializeTask() {
-        consumer.subscribe(EasyMock.eq(asList(TOPIC)), EasyMock.capture(rebalanceListener));
-        PowerMock.expectLastCall();
-
-        sinkTask.initialize(EasyMock.capture(sinkTaskContext));
-        PowerMock.expectLastCall();
-        sinkTask.start(TASK_PROPS);
-        PowerMock.expectLastCall();
-    }
-
-    private void expectPollInitialAssignment() {
-        sinkTask.open(INITIAL_ASSIGNMENT);
-        EasyMock.expectLastCall();
-
-        EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
-
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(() -> {
-            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
-            return ConsumerRecords.empty();
-        });
-        INITIAL_ASSIGNMENT.forEach(tp -> EasyMock.expect(consumer.position(tp)).andReturn(FIRST_OFFSET));
-
-        sinkTask.put(Collections.emptyList());
-        EasyMock.expectLastCall();
-    }
-
-    private void expectConsumerPoll(final int numMessages) {
-        expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE, emptyHeaders());
-    }
-
-    private void expectConsumerPoll(final int numMessages, final long timestamp, final TimestampType timestampType, Headers headers) {
-        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
-            () -> {
-                List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
-                for (int i = 0; i < numMessages; i++)
-                    records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + i, timestamp, timestampType,
-                        0, 0, RAW_KEY, RAW_VALUE, headers, Optional.empty()));
-                recordsReturnedTp1 += numMessages;
-                return new ConsumerRecords<>(
-                        numMessages > 0 ?
-                                Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records) :
-                                Collections.emptyMap()
-                );
-            });
-    }
-
-    private void expectConversionAndTransformation(final int numMessages) {
-        expectConversionAndTransformation(numMessages, null);
-    }
-
-    private void expectConversionAndTransformation(final int numMessages, final String topicPrefix) {
-        expectConversionAndTransformation(numMessages, topicPrefix, emptyHeaders());
-    }
-
-    private void expectConversionAndTransformation(final int numMessages, final String topicPrefix, final Headers headers) {
-        EasyMock.expect(keyConverter.toConnectData(TOPIC, headers, RAW_KEY)).andReturn(new SchemaAndValue(KEY_SCHEMA, KEY)).times(numMessages);
-        EasyMock.expect(valueConverter.toConnectData(TOPIC, headers, RAW_VALUE)).andReturn(new SchemaAndValue(VALUE_SCHEMA, VALUE)).times(numMessages);
-
-        for (Header header : headers) {
-            EasyMock.expect(headerConverter.toConnectHeader(TOPIC, header.key(), header.value())).andReturn(new SchemaAndValue(VALUE_SCHEMA, new String(header.value()))).times(1);
-        }
-
-        expectTransformation(numMessages, topicPrefix);
-    }
-
-    private void expectTransformation(final int numMessages, final String topicPrefix) {
-        final Capture<SinkRecord> recordCapture = EasyMock.newCapture();
-        EasyMock.expect(transformationChain.apply(EasyMock.anyObject(), EasyMock.capture(recordCapture)))
-                .andAnswer(() -> {
-                    SinkRecord origRecord = recordCapture.getValue();
-                    return topicPrefix != null && !topicPrefix.isEmpty()
-                           ? origRecord.newRecord(
-                                   topicPrefix + origRecord.topic(),
-                                   origRecord.kafkaPartition(),
-                                   origRecord.keySchema(),
-                                   origRecord.key(),
-                                   origRecord.valueSchema(),
-                                   origRecord.value(),
-                                   origRecord.timestamp(),
-                                   origRecord.headers()
-                           )
-                           : origRecord;
-                }).times(numMessages);
-    }
-
-    private void expectTaskGetTopic(boolean anyTimes) {
-        final Capture<String> connectorCapture = EasyMock.newCapture();
-        final Capture<String> topicCapture = EasyMock.newCapture();
-        IExpectationSetters<TopicStatus> expect = EasyMock.expect(statusBackingStore.getTopic(
-                EasyMock.capture(connectorCapture),
-                EasyMock.capture(topicCapture)));
-        if (anyTimes) {
-            expect.andStubAnswer(() -> new TopicStatus(
-                    topicCapture.getValue(),
-                    new ConnectorTaskId(connectorCapture.getValue(), 0),
-                    Time.SYSTEM.milliseconds()));
-        } else {
-            expect.andAnswer(() -> new TopicStatus(
-                    topicCapture.getValue(),
-                    new ConnectorTaskId(connectorCapture.getValue(), 0),
-                    Time.SYSTEM.milliseconds()));
-        }
-        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
-            assertEquals("job", connectorCapture.getValue());
-            assertEquals(TOPIC, topicCapture.getValue());
-        }
-    }
-
-    private void assertSinkMetricValue(String name, double expected) {
-        MetricGroup sinkTaskGroup = workerTask.sinkTaskMetricsGroup().metricGroup();
-        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
-        assertEquals(expected, measured, 0.001d);
-    }
-
-    private void assertTaskMetricValue(String name, double expected) {
-        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
-        double measured = metrics.currentMetricValueAsDouble(taskGroup, name);
-        assertEquals(expected, measured, 0.001d);
-    }
-
-    private void assertTaskMetricValue(String name, String expected) {
-        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
-        String measured = metrics.currentMetricValueAsString(taskGroup, name);
-        assertEquals(expected, measured);
-    }
-
-    private RecordHeaders emptyHeaders() {
-        return new RecordHeaders();
-    }
-
-    private abstract static class TestSinkTask extends SinkTask  {
-    }
-}

From ccf4bd5f4621b80f5b1d3700df30e3b444381927 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Wed, 6 Mar 2024 12:02:58 -0800
Subject: [PATCH 121/258] MINOR: Add 3.7 to Kafka Streams system tests (#15443)

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 build.gradle                                  |  15 +-
 settings.gradle                               |   1 +
 .../apache/kafka/streams/StreamsConfig.java   |   8 +-
 .../streams/internals/UpgradeFromValues.java  |   3 +-
 .../assignment/AssignorConfiguration.java     |   2 +
 .../kafka/streams/tests/SmokeTestClient.java  | 299 ++++++++
 .../kafka/streams/tests/SmokeTestDriver.java  | 670 ++++++++++++++++++
 .../kafka/streams/tests/SmokeTestUtil.java    | 131 ++++
 .../kafka/streams/tests/StreamsSmokeTest.java | 100 +++
 .../streams/tests/StreamsUpgradeTest.java     | 120 ++++
 .../streams_application_upgrade_test.py       |   4 +-
 .../streams_broker_compatibility_test.py      |  64 +-
 .../tests/streams/streams_upgrade_test.py     |   6 +-
 13 files changed, 1365 insertions(+), 58 deletions(-)
 create mode 100644 streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestClient.java
 create mode 100644 streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestDriver.java
 create mode 100644 streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestUtil.java
 create mode 100644 streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsSmokeTest.java
 create mode 100644 streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java

diff --git a/build.gradle b/build.gradle
index cba4381aae5e0..786c731dccd8d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2303,6 +2303,7 @@ project(':streams') {
             ':streams:upgrade-system-tests-34:test',
             ':streams:upgrade-system-tests-35:test',
             ':streams:upgrade-system-tests-36:test',
+            ':streams:upgrade-system-tests-37:test',
             ':streams:examples:test'
     ]
   )
@@ -2721,7 +2722,6 @@ project(':streams:upgrade-system-tests-35') {
   }
 }
 
-
 project(':streams:upgrade-system-tests-36') {
   archivesBaseName = "kafka-streams-upgrade-system-tests-36"
 
@@ -2735,6 +2735,19 @@ project(':streams:upgrade-system-tests-36') {
   }
 }
 
+project(':streams:upgrade-system-tests-37') {
+  archivesBaseName = "kafka-streams-upgrade-system-tests-37"
+
+  dependencies {
+    testImplementation libs.kafkaStreams_37
+    testRuntimeOnly libs.junitJupiter
+  }
+
+  systemTestLibs {
+    dependsOn testJar
+  }
+}
+
 project(':jmh-benchmarks') {
 
   apply plugin: 'com.github.johnrengelman.shadow'
diff --git a/settings.gradle b/settings.gradle
index 4a4141f1b760c..c2d66475ad413 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -96,6 +96,7 @@ include 'clients',
     'streams:upgrade-system-tests-34',
     'streams:upgrade-system-tests-35',
     'streams:upgrade-system-tests-36',
+    'streams:upgrade-system-tests-37',
     'tools',
     'tools:tools-api',
     'transaction-coordinator',
diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 2f4841269b5dd..54b30fea4d5d2 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -418,6 +418,12 @@ public class StreamsConfig extends AbstractConfig {
     @SuppressWarnings("WeakerAccess")
     public static final String UPGRADE_FROM_36 = UpgradeFromValues.UPGRADE_FROM_36.toString();
 
+    /**
+     * Config value for parameter {@link #UPGRADE_FROM_CONFIG "upgrade.from"} for upgrading an application from version {@code 3.7.x}.
+     */
+    @SuppressWarnings("WeakerAccess")
+    public static final String UPGRADE_FROM_37 = UpgradeFromValues.UPGRADE_FROM_37.toString();
+
     /**
      * Config value for parameter {@link #PROCESSING_GUARANTEE_CONFIG "processing.guarantee"} for at-least-once processing guarantees.
      */
@@ -774,7 +780,7 @@ public class StreamsConfig extends AbstractConfig {
         UPGRADE_FROM_25 + "\", \"" + UPGRADE_FROM_26 + "\", \"" + UPGRADE_FROM_27 + "\", \"" +
         UPGRADE_FROM_28 + "\", \"" + UPGRADE_FROM_30 + "\", \"" + UPGRADE_FROM_31 + "\", \"" +
         UPGRADE_FROM_32 + "\", \"" + UPGRADE_FROM_33 + "\", \"" + UPGRADE_FROM_34 + "\", \"" +
-        UPGRADE_FROM_35 + "\", \"" + UPGRADE_FROM_36 + "(for upgrading from the corresponding old version).";
+        UPGRADE_FROM_35 + "\", \"" + UPGRADE_FROM_36 + "\", \"" + UPGRADE_FROM_37 + "(for upgrading from the corresponding old version).";
 
     /** {@code windowstore.changelog.additional.retention.ms} */
     @SuppressWarnings("WeakerAccess")
diff --git a/streams/src/main/java/org/apache/kafka/streams/internals/UpgradeFromValues.java b/streams/src/main/java/org/apache/kafka/streams/internals/UpgradeFromValues.java
index abc2c7cb453fd..2bf19da39b618 100644
--- a/streams/src/main/java/org/apache/kafka/streams/internals/UpgradeFromValues.java
+++ b/streams/src/main/java/org/apache/kafka/streams/internals/UpgradeFromValues.java
@@ -38,7 +38,8 @@ public enum UpgradeFromValues {
     UPGRADE_FROM_33("3.3"),
     UPGRADE_FROM_34("3.4"),
     UPGRADE_FROM_35("3.5"),
-    UPGRADE_FROM_36("3.6");
+    UPGRADE_FROM_36("3.6"),
+    UPGRADE_FROM_37("3.7");
 
     private final String value;
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
index 039fd2b258b3a..6d99d93536e5d 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
@@ -127,6 +127,7 @@ public RebalanceProtocol rebalanceProtocol() {
                 case UPGRADE_FROM_34:
                 case UPGRADE_FROM_35:
                 case UPGRADE_FROM_36:
+                case UPGRADE_FROM_37:
                     // we need to add new version when new "upgrade.from" values become available
 
                     // This config is for explicitly sending FK response to a requested partition
@@ -187,6 +188,7 @@ public int configuredMetadataVersion(final int priorVersion) {
                 case UPGRADE_FROM_34:
                 case UPGRADE_FROM_35:
                 case UPGRADE_FROM_36:
+                case UPGRADE_FROM_37:
                     // we need to add new version when new "upgrade.from" values become available
 
                     // This config is for explicitly sending FK response to a requested partition
diff --git a/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestClient.java b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestClient.java
new file mode 100644
index 0000000000000..dc0ad4d5601c8
--- /dev/null
+++ b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestClient.java
@@ -0,0 +1,299 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.tests;
+
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.utils.Bytes;
+import org.apache.kafka.common.utils.KafkaThread;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.errors.StreamsUncaughtExceptionHandler;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.kstream.Grouped;
+import org.apache.kafka.streams.kstream.KGroupedStream;
+import org.apache.kafka.streams.kstream.KStream;
+import org.apache.kafka.streams.kstream.KTable;
+import org.apache.kafka.streams.kstream.Materialized;
+import org.apache.kafka.streams.kstream.Produced;
+import org.apache.kafka.streams.kstream.Suppressed.BufferConfig;
+import org.apache.kafka.streams.kstream.TimeWindows;
+import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.state.Stores;
+import org.apache.kafka.streams.state.WindowStore;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Properties;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.streams.kstream.Suppressed.untilWindowCloses;
+
+public class SmokeTestClient extends SmokeTestUtil {
+
+    private final String name;
+
+    private KafkaStreams streams;
+    private boolean uncaughtException = false;
+    private boolean started;
+    private volatile boolean closed;
+
+    private static void addShutdownHook(final String name, final Runnable runnable) {
+        if (name != null) {
+            Runtime.getRuntime().addShutdownHook(KafkaThread.nonDaemon(name, runnable));
+        } else {
+            Runtime.getRuntime().addShutdownHook(new Thread(runnable));
+        }
+    }
+
+    private static File tempDirectory() {
+        final String prefix = "kafka-";
+        final File file;
+        try {
+            file = Files.createTempDirectory(prefix).toFile();
+        } catch (final IOException ex) {
+            throw new RuntimeException("Failed to create a temp dir", ex);
+        }
+        file.deleteOnExit();
+
+        addShutdownHook("delete-temp-file-shutdown-hook", () -> {
+            try {
+                Utils.delete(file);
+            } catch (final IOException e) {
+                System.out.println("Error deleting " + file.getAbsolutePath());
+                e.printStackTrace(System.out);
+            }
+        });
+
+        return file;
+    }
+
+    public SmokeTestClient(final String name) {
+        this.name = name;
+    }
+
+    public boolean started() {
+        return started;
+    }
+
+    public boolean closed() {
+        return closed;
+    }
+
+    public void start(final Properties streamsProperties) {
+        final Topology build = getTopology();
+        streams = new KafkaStreams(build, getStreamsConfig(streamsProperties));
+
+        final CountDownLatch countDownLatch = new CountDownLatch(1);
+        streams.setStateListener((newState, oldState) -> {
+            System.out.printf("%s %s: %s -> %s%n", name, Instant.now(), oldState, newState);
+            if (oldState == KafkaStreams.State.REBALANCING && newState == KafkaStreams.State.RUNNING) {
+                started = true;
+                countDownLatch.countDown();
+            }
+
+            if (newState == KafkaStreams.State.NOT_RUNNING) {
+                closed = true;
+            }
+        });
+
+        streams.setUncaughtExceptionHandler(e -> {
+            System.out.println(name + ": SMOKE-TEST-CLIENT-EXCEPTION");
+            System.out.println(name + ": FATAL: An unexpected exception is encountered: " + e);
+            e.printStackTrace(System.out);
+            uncaughtException = true;
+            return StreamsUncaughtExceptionHandler.StreamThreadExceptionResponse.SHUTDOWN_CLIENT;
+        });
+
+        addShutdownHook("streams-shutdown-hook", this::close);
+
+        streams.start();
+        try {
+            if (!countDownLatch.await(1, TimeUnit.MINUTES)) {
+                System.out.println(name + ": SMOKE-TEST-CLIENT-EXCEPTION: Didn't start in one minute");
+            }
+        } catch (final InterruptedException e) {
+            System.out.println(name + ": SMOKE-TEST-CLIENT-EXCEPTION: " + e);
+            e.printStackTrace(System.out);
+        }
+        System.out.println(name + ": SMOKE-TEST-CLIENT-STARTED");
+        System.out.println(name + " started at " + Instant.now());
+    }
+
+    public void closeAsync() {
+        streams.close(Duration.ZERO);
+    }
+
+    public void close() {
+        final boolean closed = streams.close(Duration.ofMinutes(1));
+
+        if (closed && !uncaughtException) {
+            System.out.println(name + ": SMOKE-TEST-CLIENT-CLOSED");
+        } else if (closed) {
+            System.out.println(name + ": SMOKE-TEST-CLIENT-EXCEPTION");
+        } else {
+            System.out.println(name + ": SMOKE-TEST-CLIENT-EXCEPTION: Didn't close");
+        }
+    }
+
+    private Properties getStreamsConfig(final Properties props) {
+        final Properties fullProps = new Properties(props);
+        fullProps.put(StreamsConfig.APPLICATION_ID_CONFIG, "SmokeTest");
+        fullProps.put(StreamsConfig.CLIENT_ID_CONFIG, "SmokeTest-" + name);
+        fullProps.put(StreamsConfig.STATE_DIR_CONFIG, tempDirectory().getAbsolutePath());
+        fullProps.putAll(props);
+        return fullProps;
+    }
+
+    public Topology getTopology() {
+        final StreamsBuilder builder = new StreamsBuilder();
+        final Consumed<String, Integer> stringIntConsumed = Consumed.with(stringSerde, intSerde);
+        final KStream<String, Integer> source = builder.stream("data", stringIntConsumed);
+        source.filterNot((k, v) -> k.equals("flush"))
+              .to("echo", Produced.with(stringSerde, intSerde));
+        final KStream<String, Integer> data = source.filter((key, value) -> value == null || value != END);
+        data.process(SmokeTestUtil.printProcessorSupplier("data", name));
+
+        // min
+        final KGroupedStream<String, Integer> groupedData = data.groupByKey(Grouped.with(stringSerde, intSerde));
+
+        final KTable<Windowed<String>, Integer> minAggregation = groupedData
+            .windowedBy(TimeWindows.ofSizeAndGrace(Duration.ofDays(1), Duration.ofMinutes(1)))
+            .aggregate(
+                () -> Integer.MAX_VALUE,
+                (aggKey, value, aggregate) -> (value < aggregate) ? value : aggregate,
+                Materialized
+                    .<String, Integer, WindowStore<Bytes, byte[]>>as("uwin-min")
+                    .withValueSerde(intSerde)
+                    .withRetention(Duration.ofHours(25))
+            );
+
+        streamify(minAggregation, "min-raw");
+
+        streamify(minAggregation.suppress(untilWindowCloses(BufferConfig.unbounded())), "min-suppressed");
+
+        minAggregation
+            .toStream(new Unwindow<>())
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("min", Produced.with(stringSerde, intSerde));
+
+        final KTable<Windowed<String>, Integer> smallWindowSum = groupedData
+            .windowedBy(TimeWindows.ofSizeAndGrace(Duration.ofSeconds(2), Duration.ofSeconds(30)).advanceBy(Duration.ofSeconds(1)))
+            .reduce(Integer::sum);
+
+        streamify(smallWindowSum, "sws-raw");
+        streamify(smallWindowSum.suppress(untilWindowCloses(BufferConfig.unbounded())), "sws-suppressed");
+
+        final KTable<String, Integer> minTable = builder.table(
+            "min",
+            Consumed.with(stringSerde, intSerde),
+            Materialized.as("minStoreName"));
+
+        minTable.toStream().process(SmokeTestUtil.printProcessorSupplier("min", name));
+
+        // max
+        groupedData
+            .windowedBy(TimeWindows.ofSizeWithNoGrace(Duration.ofDays(2)))
+            .aggregate(
+                () -> Integer.MIN_VALUE,
+                (aggKey, value, aggregate) -> (value > aggregate) ? value : aggregate,
+                Materialized.<String, Integer, WindowStore<Bytes, byte[]>>as("uwin-max").withValueSerde(intSerde))
+            .toStream(new Unwindow<>())
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("max", Produced.with(stringSerde, intSerde));
+
+        final KTable<String, Integer> maxTable = builder.table(
+            "max",
+            Consumed.with(stringSerde, intSerde),
+            Materialized.as("maxStoreName"));
+        maxTable.toStream().process(SmokeTestUtil.printProcessorSupplier("max", name));
+
+        // sum
+        groupedData
+            .windowedBy(TimeWindows.ofSizeWithNoGrace(Duration.ofDays(2)))
+            .aggregate(
+                () -> 0L,
+                (aggKey, value, aggregate) -> (long) value + aggregate,
+                Materialized.<String, Long, WindowStore<Bytes, byte[]>>as("win-sum").withValueSerde(longSerde))
+            .toStream(new Unwindow<>())
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("sum", Produced.with(stringSerde, longSerde));
+
+        final Consumed<String, Long> stringLongConsumed = Consumed.with(stringSerde, longSerde);
+        final KTable<String, Long> sumTable = builder.table("sum", stringLongConsumed);
+        sumTable.toStream().process(SmokeTestUtil.printProcessorSupplier("sum", name));
+
+        // cnt
+        groupedData
+            .windowedBy(TimeWindows.ofSizeWithNoGrace(Duration.ofDays(2)))
+            .count(Materialized.as("uwin-cnt"))
+            .toStream(new Unwindow<>())
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("cnt", Produced.with(stringSerde, longSerde));
+
+        final KTable<String, Long> cntTable = builder.table(
+            "cnt",
+            Consumed.with(stringSerde, longSerde),
+            Materialized.as("cntStoreName"));
+        cntTable.toStream().process(SmokeTestUtil.printProcessorSupplier("cnt", name));
+
+        // dif
+        maxTable
+            .join(
+                minTable,
+                (value1, value2) -> value1 - value2)
+            .toStream()
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("dif", Produced.with(stringSerde, intSerde));
+
+        // avg
+        sumTable
+            .join(
+                cntTable,
+                (value1, value2) -> (double) value1 / (double) value2)
+            .toStream()
+            .filterNot((k, v) -> k.equals("flush"))
+            .to("avg", Produced.with(stringSerde, doubleSerde));
+
+        // test repartition
+        final Agg agg = new Agg();
+        cntTable.groupBy(agg.selector(), Grouped.with(stringSerde, longSerde))
+                .aggregate(agg.init(), agg.adder(), agg.remover(),
+                           Materialized.<String, Long>as(Stores.inMemoryKeyValueStore("cntByCnt"))
+                               .withKeySerde(Serdes.String())
+                               .withValueSerde(Serdes.Long()))
+                .toStream()
+                .to("tagg", Produced.with(stringSerde, longSerde));
+
+        return builder.build();
+    }
+
+    private static void streamify(final KTable<Windowed<String>, Integer> windowedTable, final String topic) {
+        windowedTable
+            .toStream()
+            .filterNot((k, v) -> k.key().equals("flush"))
+            .map((key, value) -> new KeyValue<>(key.toString(), value))
+            .to(topic, Produced.with(stringSerde, intSerde));
+    }
+}
diff --git a/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestDriver.java b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestDriver.java
new file mode 100644
index 0000000000000..dbacbb9625b61
--- /dev/null
+++ b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestDriver.java
@@ -0,0 +1,670 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.tests;
+
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.serialization.Serde;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Utils;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Random;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static java.util.Collections.emptyMap;
+import static org.apache.kafka.common.utils.Utils.mkEntry;
+
+public class SmokeTestDriver extends SmokeTestUtil {
+    private static final String[] NUMERIC_VALUE_TOPICS = {
+        "data",
+        "echo",
+        "max",
+        "min", "min-suppressed", "min-raw",
+        "dif",
+        "sum",
+        "sws-raw", "sws-suppressed",
+        "cnt",
+        "avg",
+        "tagg"
+    };
+    private static final String[] STRING_VALUE_TOPICS = {
+        "fk"
+    };
+
+    private static final String[] TOPICS = new String[NUMERIC_VALUE_TOPICS.length + STRING_VALUE_TOPICS.length];
+    static {
+        System.arraycopy(NUMERIC_VALUE_TOPICS, 0, TOPICS, 0, NUMERIC_VALUE_TOPICS.length);
+        System.arraycopy(STRING_VALUE_TOPICS, 0, TOPICS, NUMERIC_VALUE_TOPICS.length, STRING_VALUE_TOPICS.length);
+    }
+
+    private static final int MAX_RECORD_EMPTY_RETRIES = 30;
+
+    private static class ValueList {
+        public final String key;
+        private final int[] values;
+        private int index;
+
+        ValueList(final int min, final int max) {
+            key = min + "-" + max;
+
+            values = new int[max - min + 1];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = min + i;
+            }
+            // We want to randomize the order of data to test not completely predictable processing order
+            // However, values are also use as a timestamp of the record. (TODO: separate data and timestamp)
+            // We keep some correlation of time and order. Thus, the shuffling is done with a sliding window
+            shuffle(values, 10);
+
+            index = 0;
+        }
+
+        int next() {
+            return (index < values.length) ? values[index++] : -1;
+        }
+    }
+
+    public static String[] topics() {
+        return Arrays.copyOf(TOPICS, TOPICS.length);
+    }
+
+    static void generatePerpetually(final String kafka,
+                                    final int numKeys,
+                                    final int maxRecordsPerKey) {
+        final Properties producerProps = generatorProperties(kafka);
+
+        int numRecordsProduced = 0;
+
+        final ValueList[] data = new ValueList[numKeys];
+        for (int i = 0; i < numKeys; i++) {
+            data[i] = new ValueList(i, i + maxRecordsPerKey - 1);
+        }
+
+        final Random rand = new Random();
+
+        try (final KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(producerProps)) {
+            while (true) {
+                final int index = rand.nextInt(numKeys);
+                final String key = data[index].key;
+                final int value = data[index].next();
+
+                final ProducerRecord<byte[], byte[]> record =
+                    new ProducerRecord<>(
+                        "data",
+                        stringSerde.serializer().serialize("", key),
+                        intSerde.serializer().serialize("", value)
+                    );
+                producer.send(record);
+
+                final ProducerRecord<byte[], byte[]> fkRecord =
+                    new ProducerRecord<>(
+                        "fk",
+                        intSerde.serializer().serialize("", value),
+                        stringSerde.serializer().serialize("", key)
+                    );
+                producer.send(fkRecord);
+
+                numRecordsProduced++;
+                if (numRecordsProduced % 100 == 0) {
+                    System.out.println(Instant.now() + " " + numRecordsProduced + " records produced");
+                }
+                Utils.sleep(2);
+            }
+        }
+    }
+
+    public static Map<String, Set<Integer>> generate(final String kafka,
+                                                     final int numKeys,
+                                                     final int maxRecordsPerKey,
+                                                     final Duration timeToSpend) {
+        final Properties producerProps = generatorProperties(kafka);
+
+        int numRecordsProduced = 0;
+
+        final Map<String, Set<Integer>> allData = new HashMap<>();
+        final ValueList[] data = new ValueList[numKeys];
+        for (int i = 0; i < numKeys; i++) {
+            data[i] = new ValueList(i, i + maxRecordsPerKey - 1);
+            allData.put(data[i].key, new HashSet<>());
+        }
+        final Random rand = new Random();
+
+        int remaining = data.length;
+
+        final long recordPauseTime = timeToSpend.toMillis() / numKeys / maxRecordsPerKey;
+
+        final List<ProducerRecord<byte[], byte[]>> dataNeedRetry = new ArrayList<>();
+        final List<ProducerRecord<byte[], byte[]>> fkNeedRetry = new ArrayList<>();
+
+        try (final KafkaProducer<byte[], byte[]> producer = new KafkaProducer<>(producerProps)) {
+            while (remaining > 0) {
+                final int index = rand.nextInt(remaining);
+                final String key = data[index].key;
+                final int value = data[index].next();
+
+                if (value < 0) {
+                    remaining--;
+                    data[index] = data[remaining];
+                } else {
+                    final ProducerRecord<byte[], byte[]> record =
+                        new ProducerRecord<>(
+                            "data",
+                            stringSerde.serializer().serialize("", key),
+                            intSerde.serializer().serialize("", value)
+                        );
+
+                    producer.send(record, new TestCallback(record, dataNeedRetry));
+
+                    final ProducerRecord<byte[], byte[]> fkRecord =
+                        new ProducerRecord<>(
+                            "fk",
+                            intSerde.serializer().serialize("", value),
+                            stringSerde.serializer().serialize("", key)
+                        );
+
+                    producer.send(fkRecord, new TestCallback(fkRecord, fkNeedRetry));
+
+                    numRecordsProduced++;
+                    allData.get(key).add(value);
+                    if (numRecordsProduced % 100 == 0) {
+                        System.out.println(Instant.now() + " " + numRecordsProduced + " records produced");
+                    }
+                    Utils.sleep(Math.max(recordPauseTime, 2));
+                }
+            }
+            producer.flush();
+
+            retry(producer, dataNeedRetry, stringSerde);
+            retry(producer, fkNeedRetry, intSerde);
+
+            flush(producer,
+                "data",
+                stringSerde.serializer().serialize("", "flush"),
+                intSerde.serializer().serialize("", 0)
+            );
+            flush(producer,
+                "fk",
+                intSerde.serializer().serialize("", 0),
+                stringSerde.serializer().serialize("", "flush")
+            );
+        }
+        return Collections.unmodifiableMap(allData);
+    }
+
+    private static void retry(final KafkaProducer<byte[], byte[]> producer,
+        List<ProducerRecord<byte[], byte[]>> needRetry,
+        final Serde<?> keySerde) {
+        int remainingRetries = 5;
+        while (!needRetry.isEmpty()) {
+            final List<ProducerRecord<byte[], byte[]>> needRetry2 = new ArrayList<>();
+            for (final ProducerRecord<byte[], byte[]> record : needRetry) {
+                System.out.println(
+                    "retry producing " + keySerde.deserializer().deserialize("", record.key()));
+                producer.send(record, new TestCallback(record, needRetry2));
+            }
+            producer.flush();
+            needRetry = needRetry2;
+            if (--remainingRetries == 0 && !needRetry.isEmpty()) {
+                System.err.println("Failed to produce all records after multiple retries");
+                Exit.exit(1);
+            }
+        }
+    }
+
+    private static void flush(final KafkaProducer<byte[], byte[]> producer,
+        final String topic,
+        final byte[] keyBytes,
+        final byte[] valBytes) {
+        // now that we've sent everything, we'll send some final records with a timestamp high enough to flush out
+        // all suppressed records.
+        final List<PartitionInfo> partitions = producer.partitionsFor(topic);
+        for (final PartitionInfo partition : partitions) {
+            producer.send(new ProducerRecord<>(
+                partition.topic(),
+                partition.partition(),
+                System.currentTimeMillis() + Duration.ofDays(2).toMillis(),
+                keyBytes,
+                valBytes
+            ));
+        }
+    }
+
+    private static Properties generatorProperties(final String kafka) {
+        final Properties producerProps = new Properties();
+        producerProps.put(ProducerConfig.CLIENT_ID_CONFIG, "SmokeTest");
+        producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka);
+        producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class);
+        producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class);
+        producerProps.put(ProducerConfig.ACKS_CONFIG, "all");
+        return producerProps;
+    }
+
+    private static class TestCallback implements Callback {
+        private final ProducerRecord<byte[], byte[]> originalRecord;
+        private final List<ProducerRecord<byte[], byte[]>> needRetry;
+
+        TestCallback(final ProducerRecord<byte[], byte[]> originalRecord,
+                     final List<ProducerRecord<byte[], byte[]>> needRetry) {
+            this.originalRecord = originalRecord;
+            this.needRetry = needRetry;
+        }
+
+        @Override
+        public void onCompletion(final RecordMetadata metadata, final Exception exception) {
+            if (exception != null) {
+                if (exception instanceof TimeoutException) {
+                    needRetry.add(originalRecord);
+                } else {
+                    exception.printStackTrace();
+                    Exit.exit(1);
+                }
+            }
+        }
+    }
+
+    private static void shuffle(final int[] data, @SuppressWarnings("SameParameterValue") final int windowSize) {
+        final Random rand = new Random();
+        for (int i = 0; i < data.length; i++) {
+            // we shuffle data within windowSize
+            final int j = rand.nextInt(Math.min(data.length - i, windowSize)) + i;
+
+            // swap
+            final int tmp = data[i];
+            data[i] = data[j];
+            data[j] = tmp;
+        }
+    }
+
+    public static class NumberDeserializer implements Deserializer<Number> {
+        @Override
+        public Number deserialize(final String topic, final byte[] data) {
+            final Number value;
+            switch (topic) {
+                case "data":
+                case "echo":
+                case "min":
+                case "min-raw":
+                case "min-suppressed":
+                case "sws-raw":
+                case "sws-suppressed":
+                case "max":
+                case "dif":
+                    value = intSerde.deserializer().deserialize(topic, data);
+                    break;
+                case "sum":
+                case "cnt":
+                case "tagg":
+                    value = longSerde.deserializer().deserialize(topic, data);
+                    break;
+                case "avg":
+                    value = doubleSerde.deserializer().deserialize(topic, data);
+                    break;
+                default:
+                    throw new RuntimeException("unknown topic: " + topic);
+            }
+            return value;
+        }
+    }
+
+    public static VerificationResult verify(final String kafka,
+                                            final Map<String, Set<Integer>> inputs,
+                                            final int maxRecordsPerKey) {
+        final Properties props = new Properties();
+        props.put(ConsumerConfig.CLIENT_ID_CONFIG, "verifier");
+        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka);
+        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, NumberDeserializer.class);
+        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
+
+        final KafkaConsumer<String, Number> consumer = new KafkaConsumer<>(props);
+        final List<TopicPartition> partitions = getAllPartitions(consumer, NUMERIC_VALUE_TOPICS);
+        consumer.assign(partitions);
+        consumer.seekToBeginning(partitions);
+
+        final int recordsGenerated = inputs.size() * maxRecordsPerKey;
+        int recordsProcessed = 0;
+        final Map<String, AtomicInteger> processed =
+            Stream.of(NUMERIC_VALUE_TOPICS)
+                  .collect(Collectors.toMap(t -> t, t -> new AtomicInteger(0)));
+
+        final Map<String, Map<String, LinkedList<ConsumerRecord<String, Number>>>> events = new HashMap<>();
+
+        VerificationResult verificationResult = new VerificationResult(false, "no results yet");
+        int retry = 0;
+        final long start = System.currentTimeMillis();
+        while (System.currentTimeMillis() - start < TimeUnit.MINUTES.toMillis(6)) {
+            final ConsumerRecords<String, Number> records = consumer.poll(Duration.ofSeconds(5));
+            if (records.isEmpty() && recordsProcessed >= recordsGenerated) {
+                verificationResult = verifyAll(inputs, events, false);
+                if (verificationResult.passed()) {
+                    break;
+                } else if (retry++ > MAX_RECORD_EMPTY_RETRIES) {
+                    System.out.println(Instant.now() + " Didn't get any more results, verification hasn't passed, and out of retries.");
+                    break;
+                } else {
+                    System.out.println(Instant.now() + " Didn't get any more results, but verification hasn't passed (yet). Retrying..." + retry);
+                }
+            } else {
+                System.out.println(Instant.now() + " Get some more results from " + records.partitions() + ", resetting retry.");
+
+                retry = 0;
+                for (final ConsumerRecord<String, Number> record : records) {
+                    final String key = record.key();
+
+                    final String topic = record.topic();
+                    processed.get(topic).incrementAndGet();
+
+                    if (topic.equals("echo")) {
+                        recordsProcessed++;
+                        if (recordsProcessed % 100 == 0) {
+                            System.out.println("Echo records processed = " + recordsProcessed);
+                        }
+                    }
+
+                    events.computeIfAbsent(topic, t -> new HashMap<>())
+                          .computeIfAbsent(key, k -> new LinkedList<>())
+                          .add(record);
+                }
+
+                System.out.println(processed);
+            }
+        }
+        consumer.close();
+        final long finished = System.currentTimeMillis() - start;
+        System.out.println("Verification time=" + finished);
+        System.out.println("-------------------");
+        System.out.println("Result Verification");
+        System.out.println("-------------------");
+        System.out.println("recordGenerated=" + recordsGenerated);
+        System.out.println("recordProcessed=" + recordsProcessed);
+
+        if (recordsProcessed > recordsGenerated) {
+            System.out.println("PROCESSED-MORE-THAN-GENERATED");
+        } else if (recordsProcessed < recordsGenerated) {
+            System.out.println("PROCESSED-LESS-THAN-GENERATED");
+        }
+
+        boolean success;
+
+        final Map<String, Set<Number>> received =
+            events.get("echo")
+                  .entrySet()
+                  .stream()
+                  .map(entry -> mkEntry(
+                      entry.getKey(),
+                      entry.getValue().stream().map(ConsumerRecord::value).collect(Collectors.toSet()))
+                  )
+                  .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        success = inputs.equals(received);
+
+        if (success) {
+            System.out.println("ALL-RECORDS-DELIVERED");
+        } else {
+            int missedCount = 0;
+            for (final Map.Entry<String, Set<Integer>> entry : inputs.entrySet()) {
+                missedCount += received.get(entry.getKey()).size();
+            }
+            System.out.println("missedRecords=" + missedCount);
+        }
+
+        // give it one more try if it's not already passing.
+        if (!verificationResult.passed()) {
+            verificationResult = verifyAll(inputs, events, true);
+        }
+        success &= verificationResult.passed();
+
+        System.out.println(verificationResult.result());
+
+        System.out.println(success ? "SUCCESS" : "FAILURE");
+        return verificationResult;
+    }
+
+    public static class VerificationResult {
+        private final boolean passed;
+        private final String result;
+
+        VerificationResult(final boolean passed, final String result) {
+            this.passed = passed;
+            this.result = result;
+        }
+
+        public boolean passed() {
+            return passed;
+        }
+
+        public String result() {
+            return result;
+        }
+    }
+
+    private static VerificationResult verifyAll(final Map<String, Set<Integer>> inputs,
+                                                final Map<String, Map<String, LinkedList<ConsumerRecord<String, Number>>>> events,
+                                                final boolean printResults) {
+        final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        boolean pass;
+        try (final PrintStream resultStream = new PrintStream(byteArrayOutputStream)) {
+            pass = verifyTAgg(resultStream, inputs, events.get("tagg"), printResults);
+            pass &= verifySuppressed(resultStream, "min-suppressed", events, printResults);
+            pass &= verify(resultStream, "min-suppressed", inputs, events, windowedKey -> {
+                final String unwindowedKey = windowedKey.substring(1, windowedKey.length() - 1).replaceAll("@.*", "");
+                return getMin(unwindowedKey);
+            }, printResults);
+            pass &= verifySuppressed(resultStream, "sws-suppressed", events, printResults);
+            pass &= verify(resultStream, "min", inputs, events, SmokeTestDriver::getMin, printResults);
+            pass &= verify(resultStream, "max", inputs, events, SmokeTestDriver::getMax, printResults);
+            pass &= verify(resultStream, "dif", inputs, events, key -> getMax(key).intValue() - getMin(key).intValue(), printResults);
+            pass &= verify(resultStream, "sum", inputs, events, SmokeTestDriver::getSum, printResults);
+            pass &= verify(resultStream, "cnt", inputs, events, key1 -> getMax(key1).intValue() - getMin(key1).intValue() + 1L, printResults);
+            pass &= verify(resultStream, "avg", inputs, events, SmokeTestDriver::getAvg, printResults);
+        }
+        return new VerificationResult(pass, new String(byteArrayOutputStream.toByteArray(), StandardCharsets.UTF_8));
+    }
+
+    private static boolean verify(final PrintStream resultStream,
+                                  final String topic,
+                                  final Map<String, Set<Integer>> inputData,
+                                  final Map<String, Map<String, LinkedList<ConsumerRecord<String, Number>>>> events,
+                                  final Function<String, Number> keyToExpectation,
+                                  final boolean printResults) {
+        final Map<String, LinkedList<ConsumerRecord<String, Number>>> observedInputEvents = events.get("data");
+        final Map<String, LinkedList<ConsumerRecord<String, Number>>> outputEvents = events.getOrDefault(topic, emptyMap());
+        if (outputEvents.isEmpty()) {
+            resultStream.println(topic + " is empty");
+            return false;
+        } else {
+            resultStream.printf("verifying %s with %d keys%n", topic, outputEvents.size());
+
+            if (outputEvents.size() != inputData.size()) {
+                resultStream.printf("fail: resultCount=%d expectedCount=%s%n\tresult=%s%n\texpected=%s%n",
+                                    outputEvents.size(), inputData.size(), outputEvents.keySet(), inputData.keySet());
+                return false;
+            }
+            for (final Map.Entry<String, LinkedList<ConsumerRecord<String, Number>>> entry : outputEvents.entrySet()) {
+                final String key = entry.getKey();
+                final Number expected = keyToExpectation.apply(key);
+                final Number actual = entry.getValue().getLast().value();
+                if (!expected.equals(actual)) {
+                    resultStream.printf("%s fail: key=%s actual=%s expected=%s%n", topic, key, actual, expected);
+
+                    if (printResults) {
+                        resultStream.printf("\t inputEvents=%n%s%n\t" +
+                                "echoEvents=%n%s%n\tmaxEvents=%n%s%n\tminEvents=%n%s%n\tdifEvents=%n%s%n\tcntEvents=%n%s%n\ttaggEvents=%n%s%n",
+                            indent("\t\t", observedInputEvents.get(key)),
+                            indent("\t\t", events.getOrDefault("echo", emptyMap()).getOrDefault(key, new LinkedList<>())),
+                            indent("\t\t", events.getOrDefault("max", emptyMap()).getOrDefault(key, new LinkedList<>())),
+                            indent("\t\t", events.getOrDefault("min", emptyMap()).getOrDefault(key, new LinkedList<>())),
+                            indent("\t\t", events.getOrDefault("dif", emptyMap()).getOrDefault(key, new LinkedList<>())),
+                            indent("\t\t", events.getOrDefault("cnt", emptyMap()).getOrDefault(key, new LinkedList<>())),
+                            indent("\t\t", events.getOrDefault("tagg", emptyMap()).getOrDefault(key, new LinkedList<>())));
+
+                        if (!Utils.mkSet("echo", "max", "min", "dif", "cnt", "tagg").contains(topic))
+                            resultStream.printf("%sEvents=%n%s%n", topic, indent("\t\t", entry.getValue()));
+                    }
+
+                    return false;
+                }
+            }
+            return true;
+        }
+    }
+
+
+    private static boolean verifySuppressed(final PrintStream resultStream,
+                                            @SuppressWarnings("SameParameterValue") final String topic,
+                                            final Map<String, Map<String, LinkedList<ConsumerRecord<String, Number>>>> events,
+                                            final boolean printResults) {
+        resultStream.println("verifying suppressed " + topic);
+        final Map<String, LinkedList<ConsumerRecord<String, Number>>> topicEvents = events.getOrDefault(topic, emptyMap());
+        for (final Map.Entry<String, LinkedList<ConsumerRecord<String, Number>>> entry : topicEvents.entrySet()) {
+            if (entry.getValue().size() != 1) {
+                final String unsuppressedTopic = topic.replace("-suppressed", "-raw");
+                final String key = entry.getKey();
+                final String unwindowedKey = key.substring(1, key.length() - 1).replaceAll("@.*", "");
+                resultStream.printf("fail: key=%s%n\tnon-unique result:%n%s%n",
+                                    key,
+                                    indent("\t\t", entry.getValue()));
+
+                if (printResults)
+                    resultStream.printf("\tresultEvents:%n%s%n\tinputEvents:%n%s%n",
+                        indent("\t\t", events.get(unsuppressedTopic).get(key)),
+                        indent("\t\t", events.get("data").get(unwindowedKey)));
+
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private static String indent(@SuppressWarnings("SameParameterValue") final String prefix,
+                                 final Iterable<ConsumerRecord<String, Number>> list) {
+        final StringBuilder stringBuilder = new StringBuilder();
+        for (final ConsumerRecord<String, Number> record : list) {
+            stringBuilder.append(prefix).append(record).append('\n');
+        }
+        return stringBuilder.toString();
+    }
+
+    private static Long getSum(final String key) {
+        final int min = getMin(key).intValue();
+        final int max = getMax(key).intValue();
+        return ((long) min + max) * (max - min + 1L) / 2L;
+    }
+
+    private static Double getAvg(final String key) {
+        final int min = getMin(key).intValue();
+        final int max = getMax(key).intValue();
+        return ((long) min + max) / 2.0;
+    }
+
+
+    private static boolean verifyTAgg(final PrintStream resultStream,
+                                      final Map<String, Set<Integer>> allData,
+                                      final Map<String, LinkedList<ConsumerRecord<String, Number>>> taggEvents,
+                                      final boolean printResults) {
+        if (taggEvents == null) {
+            resultStream.println("tagg is missing");
+            return false;
+        } else if (taggEvents.isEmpty()) {
+            resultStream.println("tagg is empty");
+            return false;
+        } else {
+            resultStream.println("verifying tagg");
+
+            // generate expected answer
+            final Map<String, Long> expected = new HashMap<>();
+            for (final String key : allData.keySet()) {
+                final int min = getMin(key).intValue();
+                final int max = getMax(key).intValue();
+                final String cnt = Long.toString(max - min + 1L);
+
+                expected.put(cnt, expected.getOrDefault(cnt, 0L) + 1);
+            }
+
+            // check the result
+            for (final Map.Entry<String, LinkedList<ConsumerRecord<String, Number>>> entry : taggEvents.entrySet()) {
+                final String key = entry.getKey();
+                Long expectedCount = expected.remove(key);
+                if (expectedCount == null) {
+                    expectedCount = 0L;
+                }
+
+                if (entry.getValue().getLast().value().longValue() != expectedCount) {
+                    resultStream.println("fail: key=" + key + " tagg=" + entry.getValue() + " expected=" + expectedCount);
+
+                    if (printResults)
+                        resultStream.println("\t taggEvents: " + entry.getValue());
+                    return false;
+                }
+            }
+
+        }
+        return true;
+    }
+
+    private static Number getMin(final String key) {
+        return Integer.parseInt(key.split("-")[0]);
+    }
+
+    private static Number getMax(final String key) {
+        return Integer.parseInt(key.split("-")[1]);
+    }
+
+    private static List<TopicPartition> getAllPartitions(final KafkaConsumer<?, ?> consumer, final String... topics) {
+        final List<TopicPartition> partitions = new ArrayList<>();
+
+        for (final String topic : topics) {
+            for (final PartitionInfo info : consumer.partitionsFor(topic)) {
+                partitions.add(new TopicPartition(info.topic(), info.partition()));
+            }
+        }
+        return partitions;
+    }
+
+}
diff --git a/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestUtil.java b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestUtil.java
new file mode 100644
index 0000000000000..7cb34056e0e37
--- /dev/null
+++ b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/SmokeTestUtil.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.tests;
+
+import org.apache.kafka.common.serialization.Serde;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.kstream.Aggregator;
+import org.apache.kafka.streams.kstream.Initializer;
+import org.apache.kafka.streams.kstream.KeyValueMapper;
+import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.processor.api.ContextualProcessor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.Record;
+
+import java.time.Instant;
+
+public class SmokeTestUtil {
+
+    final static int END = Integer.MAX_VALUE;
+
+    static ProcessorSupplier<Object, Object, Void, Void> printProcessorSupplier(final String topic) {
+        return printProcessorSupplier(topic, "");
+    }
+
+    static ProcessorSupplier<Object, Object, Void, Void> printProcessorSupplier(final String topic, final String name) {
+        return () -> new ContextualProcessor<Object, Object, Void, Void>() {
+            private int numRecordsProcessed = 0;
+            private long smallestOffset = Long.MAX_VALUE;
+            private long largestOffset = Long.MIN_VALUE;
+
+            @Override
+            public void init(final ProcessorContext<Void, Void> context) {
+                super.init(context);
+                System.out.println("[3.7] initializing processor: topic=" + topic + " taskId=" + context.taskId());
+                System.out.flush();
+                numRecordsProcessed = 0;
+                smallestOffset = Long.MAX_VALUE;
+                largestOffset = Long.MIN_VALUE;
+            }
+
+            @Override
+            public void process(final Record<Object, Object> record) {
+                numRecordsProcessed++;
+                if (numRecordsProcessed % 100 == 0) {
+                    System.out.printf("%s: %s%n", name, Instant.now());
+                    System.out.println("processed " + numRecordsProcessed + " records from topic=" + topic);
+                }
+
+                context().recordMetadata().ifPresent(recordMetadata -> {
+                    if (smallestOffset > recordMetadata.offset()) {
+                        smallestOffset = recordMetadata.offset();
+                    }
+                    if (largestOffset < recordMetadata.offset()) {
+                        largestOffset = recordMetadata.offset();
+                    }
+                });
+            }
+
+            @Override
+            public void close() {
+                System.out.printf("Close processor for task %s%n", context().taskId());
+                System.out.println("processed " + numRecordsProcessed + " records");
+                final long processed;
+                if (largestOffset >= smallestOffset) {
+                    processed = 1L + largestOffset - smallestOffset;
+                } else {
+                    processed = 0L;
+                }
+                System.out.println("offset " + smallestOffset + " to " + largestOffset + " -> processed " + processed);
+                System.out.flush();
+            }
+        };
+    }
+
+    public static final class Unwindow<K, V> implements KeyValueMapper<Windowed<K>, V, K> {
+        @Override
+        public K apply(final Windowed<K> winKey, final V value) {
+            return winKey.key();
+        }
+    }
+
+    public static class Agg {
+
+        KeyValueMapper<String, Long, KeyValue<String, Long>> selector() {
+            return (key, value) -> new KeyValue<>(value == null ? null : Long.toString(value), 1L);
+        }
+
+        public Initializer<Long> init() {
+            return () -> 0L;
+        }
+
+        Aggregator<String, Long, Long> adder() {
+            return (aggKey, value, aggregate) -> aggregate + value;
+        }
+
+        Aggregator<String, Long, Long> remover() {
+            return (aggKey, value, aggregate) -> aggregate - value;
+        }
+    }
+
+    public static Serde<String> stringSerde = Serdes.String();
+
+    public static Serde<Integer> intSerde = Serdes.Integer();
+
+    static Serde<Long> longSerde = Serdes.Long();
+
+    static Serde<Double> doubleSerde = Serdes.Double();
+
+    public static void sleep(final long duration) {
+        try {
+            Thread.sleep(duration);
+        } catch (final Exception ignore) { }
+    }
+
+}
diff --git a/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsSmokeTest.java b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsSmokeTest.java
new file mode 100644
index 0000000000000..5803b2fbd0217
--- /dev/null
+++ b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsSmokeTest.java
@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.tests;
+
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.streams.StreamsConfig;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+import java.util.UUID;
+
+import static org.apache.kafka.streams.tests.SmokeTestDriver.generate;
+import static org.apache.kafka.streams.tests.SmokeTestDriver.generatePerpetually;
+
+public class StreamsSmokeTest {
+
+    /**
+     *  args ::= kafka propFileName command disableAutoTerminate
+     *  command := "run" | "process"
+     *
+     * @param args
+     */
+    public static void main(final String[] args) throws IOException {
+        if (args.length < 2) {
+            System.err.println("StreamsSmokeTest are expecting two parameters: propFile, command; but only see " + args.length + " parameter");
+            Exit.exit(1);
+        }
+
+        final String propFileName = args[0];
+        final String command = args[1];
+        final boolean disableAutoTerminate = args.length > 2;
+
+        final Properties streamsProperties = Utils.loadProps(propFileName);
+        final String kafka = streamsProperties.getProperty(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG);
+        final String processingGuarantee = streamsProperties.getProperty(StreamsConfig.PROCESSING_GUARANTEE_CONFIG);
+
+        if (kafka == null) {
+            System.err.println("No bootstrap kafka servers specified in " + StreamsConfig.BOOTSTRAP_SERVERS_CONFIG);
+            Exit.exit(1);
+        }
+
+        if ("process".equals(command)) {
+            if (!StreamsConfig.AT_LEAST_ONCE.equals(processingGuarantee) &&
+                !StreamsConfig.EXACTLY_ONCE_V2.equals(processingGuarantee)) {
+
+                System.err.println("processingGuarantee must be either " + StreamsConfig.AT_LEAST_ONCE + " or " +
+                    StreamsConfig.EXACTLY_ONCE_V2);
+
+                Exit.exit(1);
+            }
+        }
+
+        System.out.println("StreamsTest instance started (StreamsSmokeTest)");
+        System.out.println("command=" + command);
+        System.out.println("props=" + streamsProperties);
+        System.out.println("disableAutoTerminate=" + disableAutoTerminate);
+
+        switch (command) {
+            case "run":
+                // this starts the driver (data generation and result verification)
+                final int numKeys = 10;
+                final int maxRecordsPerKey = 500;
+                if (disableAutoTerminate) {
+                    generatePerpetually(kafka, numKeys, maxRecordsPerKey);
+                } else {
+                    // slow down data production to span 30 seconds so that system tests have time to
+                    // do their bounces, etc.
+                    final Map<String, Set<Integer>> allData =
+                        generate(kafka, numKeys, maxRecordsPerKey, Duration.ofSeconds(30));
+                    SmokeTestDriver.verify(kafka, allData, maxRecordsPerKey);
+                }
+                break;
+            case "process":
+                // this starts the stream processing app
+                new SmokeTestClient(UUID.randomUUID().toString()).start(streamsProperties);
+                break;
+            default:
+                System.out.println("unknown command: " + command);
+        }
+    }
+
+}
diff --git a/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java
new file mode 100644
index 0000000000000..15769bf16c333
--- /dev/null
+++ b/streams/upgrade-system-tests-37/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java
@@ -0,0 +1,120 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.tests;
+
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.kstream.KStream;
+import org.apache.kafka.streams.kstream.KTable;
+import org.apache.kafka.streams.kstream.Produced;
+import org.apache.kafka.streams.processor.api.ContextualProcessor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.Record;
+
+import java.util.Properties;
+
+import static org.apache.kafka.streams.tests.SmokeTestUtil.intSerde;
+import static org.apache.kafka.streams.tests.SmokeTestUtil.stringSerde;
+
+
+public class StreamsUpgradeTest {
+
+    @SuppressWarnings("unchecked")
+    public static void main(final String[] args) throws Exception {
+        if (args.length < 1) {
+            System.err.println("StreamsUpgradeTest requires one argument (properties-file) but provided none");
+        }
+        final String propFileName = args[0];
+
+        final Properties streamsProperties = Utils.loadProps(propFileName);
+
+        System.out.println("StreamsTest instance started (StreamsUpgradeTest v3.7)");
+        System.out.println("props=" + streamsProperties);
+
+        final StreamsBuilder builder = new StreamsBuilder();
+        final KTable<String, Integer> dataTable = builder.table(
+            "data", Consumed.with(stringSerde, intSerde));
+        final KStream<String, Integer> dataStream = dataTable.toStream();
+        dataStream.process(printProcessorSupplier("data"));
+        dataStream.to("echo");
+
+        final boolean runFkJoin = Boolean.parseBoolean(streamsProperties.getProperty(
+            "test.run_fk_join",
+            "false"));
+        if (runFkJoin) {
+            try {
+                final KTable<Integer, String> fkTable = builder.table(
+                    "fk", Consumed.with(intSerde, stringSerde));
+                buildFKTable(dataStream, fkTable);
+            } catch (final Exception e) {
+                System.err.println("Caught " + e.getMessage());
+            }
+        }
+
+        final Properties config = new Properties();
+        config.setProperty(
+            StreamsConfig.APPLICATION_ID_CONFIG,
+            "StreamsUpgradeTest");
+        config.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000);
+        config.putAll(streamsProperties);
+
+        final KafkaStreams streams = new KafkaStreams(builder.build(), config);
+        streams.start();
+
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+            streams.close();
+            System.out.println("UPGRADE-TEST-CLIENT-CLOSED");
+            System.out.flush();
+        }));
+    }
+
+    private static void buildFKTable(final KStream<String, Integer> primaryTable,
+                                     final KTable<Integer, String> otherTable) {
+        final KStream<String, String> kStream = primaryTable.toTable()
+            .join(otherTable, v -> v, (k0, v0) -> v0)
+            .toStream();
+        kStream.process(printProcessorSupplier("fk"));
+        kStream.to("fk-result", Produced.with(stringSerde, stringSerde));
+    }
+
+    private static <KIn, VIn, KOut, VOut> ProcessorSupplier<KIn, VIn, KOut, VOut> printProcessorSupplier(final String topic) {
+        return () -> new ContextualProcessor<KIn, VIn, KOut, VOut>() {
+            private int numRecordsProcessed = 0;
+
+            @Override
+            public void init(final ProcessorContext<KOut, VOut> context) {
+                System.out.println("[3.7] initializing processor: topic=" + topic + "taskId=" + context.taskId());
+                numRecordsProcessed = 0;
+            }
+
+            @Override
+            public void process(final Record<KIn, VIn> record) {
+                numRecordsProcessed++;
+                if (numRecordsProcessed % 100 == 0) {
+                    System.out.println("processed " + numRecordsProcessed + " records from topic=" + topic);
+                }
+            }
+
+            @Override
+            public void close() {}
+        };
+    }
+}
diff --git a/tests/kafkatest/tests/streams/streams_application_upgrade_test.py b/tests/kafkatest/tests/streams/streams_application_upgrade_test.py
index 057f301df6fd2..6719a59178fb9 100644
--- a/tests/kafkatest/tests/streams/streams_application_upgrade_test.py
+++ b/tests/kafkatest/tests/streams/streams_application_upgrade_test.py
@@ -22,13 +22,13 @@
 from kafkatest.services.streams import StreamsSmokeTestDriverService, StreamsSmokeTestJobRunnerService
 from kafkatest.services.zookeeper import ZookeeperService
 from kafkatest.version import LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, LATEST_2_7, LATEST_2_8, \
-  LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, DEV_VERSION, KafkaVersion
+  LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, LATEST_3_7, DEV_VERSION, KafkaVersion
 
 smoke_test_versions = [str(LATEST_2_2), str(LATEST_2_3), str(LATEST_2_4),
                        str(LATEST_2_5), str(LATEST_2_6), str(LATEST_2_7),
                        str(LATEST_2_8), str(LATEST_3_0), str(LATEST_3_1),
                        str(LATEST_3_2), str(LATEST_3_3), str(LATEST_3_4),
-                       str(LATEST_3_5), str(LATEST_3_6)]
+                       str(LATEST_3_5), str(LATEST_3_6), str(LATEST_3_7)]
 dev_version = [str(DEV_VERSION)]
 
 class StreamsUpgradeTest(Test):
diff --git a/tests/kafkatest/tests/streams/streams_broker_compatibility_test.py b/tests/kafkatest/tests/streams/streams_broker_compatibility_test.py
index 82647687f83a6..1262c6d9f418d 100644
--- a/tests/kafkatest/tests/streams/streams_broker_compatibility_test.py
+++ b/tests/kafkatest/tests/streams/streams_broker_compatibility_test.py
@@ -23,7 +23,7 @@
 from kafkatest.services.zookeeper import ZookeeperService
 from kafkatest.version import LATEST_0_11_0, LATEST_0_10_2, LATEST_0_10_1, LATEST_0_10_0, LATEST_1_0, LATEST_1_1, \
     LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, LATEST_2_7, LATEST_2_8, \
-    LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, KafkaVersion
+    LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, LATEST_3_7, KafkaVersion
 
 
 class StreamsBrokerCompatibility(Test):
@@ -64,25 +64,11 @@ def setUp(self):
 
 
     @cluster(num_nodes=4)
-    @parametrize(broker_version=str(LATEST_3_6))
-    @parametrize(broker_version=str(LATEST_3_5))
-    @parametrize(broker_version=str(LATEST_3_4))
-    @parametrize(broker_version=str(LATEST_3_3))
-    @parametrize(broker_version=str(LATEST_3_2))
-    @parametrize(broker_version=str(LATEST_3_1))
-    @parametrize(broker_version=str(LATEST_3_0))
-    @parametrize(broker_version=str(LATEST_2_8))
-    @parametrize(broker_version=str(LATEST_2_7))
-    @parametrize(broker_version=str(LATEST_2_6))
-    @parametrize(broker_version=str(LATEST_2_5))
-    @parametrize(broker_version=str(LATEST_2_4))
-    @parametrize(broker_version=str(LATEST_2_3))
-    @parametrize(broker_version=str(LATEST_2_2))
-    @parametrize(broker_version=str(LATEST_2_1))
-    @parametrize(broker_version=str(LATEST_2_0))
-    @parametrize(broker_version=str(LATEST_1_1))
-    @parametrize(broker_version=str(LATEST_1_0))
-    @parametrize(broker_version=str(LATEST_0_11_0))
+    @matrix(broker_version=[str(LATEST_0_11_0),str(LATEST_1_0),str(LATEST_1_1),str(LATEST_2_0),
+                            str(LATEST_2_1),str(LATEST_2_2),str(LATEST_2_3),str(LATEST_2_4),
+                            str(LATEST_2_5),str(LATEST_2_6),str(LATEST_2_7),str(LATEST_2_8),
+                            str(LATEST_3_0),str(LATEST_3_1),str(LATEST_3_2),str(LATEST_3_3),
+                            str(LATEST_3_4),str(LATEST_3_5),str(LATEST_3_6),str(LATEST_3_7)])
     def test_compatible_brokers_eos_disabled(self, broker_version):
         self.kafka.set_version(KafkaVersion(broker_version))
         self.kafka.start()
@@ -100,25 +86,11 @@ def test_compatible_brokers_eos_disabled(self, broker_version):
         self.kafka.stop()
 
     @cluster(num_nodes=4)
-    @parametrize(broker_version=str(LATEST_3_6))
-    @parametrize(broker_version=str(LATEST_3_5))
-    @parametrize(broker_version=str(LATEST_3_4))
-    @parametrize(broker_version=str(LATEST_3_3))
-    @parametrize(broker_version=str(LATEST_3_2))
-    @parametrize(broker_version=str(LATEST_3_1))
-    @parametrize(broker_version=str(LATEST_3_0))
-    @parametrize(broker_version=str(LATEST_2_8))
-    @parametrize(broker_version=str(LATEST_2_7))
-    @parametrize(broker_version=str(LATEST_2_6))
-    @parametrize(broker_version=str(LATEST_2_5))
-    @parametrize(broker_version=str(LATEST_2_4))
-    @parametrize(broker_version=str(LATEST_2_3))
-    @parametrize(broker_version=str(LATEST_2_2))
-    @parametrize(broker_version=str(LATEST_2_1))
-    @parametrize(broker_version=str(LATEST_2_0))
-    @parametrize(broker_version=str(LATEST_1_1))
-    @parametrize(broker_version=str(LATEST_1_0))
-    @parametrize(broker_version=str(LATEST_0_11_0))
+    @matrix(broker_version=[str(LATEST_0_11_0),str(LATEST_1_0),str(LATEST_1_1),str(LATEST_2_0),
+                            str(LATEST_2_1),str(LATEST_2_2),str(LATEST_2_3),str(LATEST_2_4),
+                            str(LATEST_2_5),str(LATEST_2_6),str(LATEST_2_7),str(LATEST_2_8),
+                            str(LATEST_3_0),str(LATEST_3_1),str(LATEST_3_2),str(LATEST_3_3),
+                            str(LATEST_3_4),str(LATEST_3_5),str(LATEST_3_6),str(LATEST_3_7)])
     def test_compatible_brokers_eos_alpha_enabled(self, broker_version):
         self.kafka.set_version(KafkaVersion(broker_version))
         self.kafka.start()
@@ -136,17 +108,9 @@ def test_compatible_brokers_eos_alpha_enabled(self, broker_version):
         self.kafka.stop()
 
     @cluster(num_nodes=4)
-    @parametrize(broker_version=str(LATEST_3_6))
-    @parametrize(broker_version=str(LATEST_3_5))
-    @parametrize(broker_version=str(LATEST_3_4))
-    @parametrize(broker_version=str(LATEST_3_3))
-    @parametrize(broker_version=str(LATEST_3_2))
-    @parametrize(broker_version=str(LATEST_3_1))
-    @parametrize(broker_version=str(LATEST_3_0))
-    @parametrize(broker_version=str(LATEST_2_8))
-    @parametrize(broker_version=str(LATEST_2_7))
-    @parametrize(broker_version=str(LATEST_2_6))
-    @parametrize(broker_version=str(LATEST_2_5))
+    @matrix(broker_version=[str(LATEST_2_5),str(LATEST_2_6),str(LATEST_2_7),str(LATEST_2_8),
+                            str(LATEST_3_0),str(LATEST_3_1),str(LATEST_3_2),str(LATEST_3_3),
+                            str(LATEST_3_4),str(LATEST_3_5),str(LATEST_3_6),str(LATEST_3_7)])
     def test_compatible_brokers_eos_v2_enabled(self, broker_version):
         self.kafka.set_version(KafkaVersion(broker_version))
         self.kafka.start()
diff --git a/tests/kafkatest/tests/streams/streams_upgrade_test.py b/tests/kafkatest/tests/streams/streams_upgrade_test.py
index bd04a0bf0788c..371feaf13ca3f 100644
--- a/tests/kafkatest/tests/streams/streams_upgrade_test.py
+++ b/tests/kafkatest/tests/streams/streams_upgrade_test.py
@@ -26,7 +26,7 @@
 from kafkatest.tests.streams.utils import extract_generation_from_logs, extract_generation_id
 from kafkatest.version import LATEST_0_10_0, LATEST_0_10_1, LATEST_0_10_2, LATEST_0_11_0, LATEST_1_0, LATEST_1_1, \
     LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, LATEST_2_7, LATEST_2_8, \
-    LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, DEV_BRANCH, DEV_VERSION, \
+    LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, LATEST_3_7, DEV_BRANCH, DEV_VERSION, \
     KafkaVersion
 
 # broker 0.10.0 is not compatible with newer Kafka Streams versions
@@ -36,7 +36,7 @@
                            str(LATEST_2_4), str(LATEST_2_5), str(LATEST_2_6), str(LATEST_2_7),
                            str(LATEST_2_8), str(LATEST_3_0), str(LATEST_3_1), str(LATEST_3_2),
                            str(LATEST_3_3), str(LATEST_3_4), str(LATEST_3_5), str(LATEST_3_6),
-                           str(DEV_BRANCH)]
+                           str(LATEST_3_7), str(DEV_BRANCH)]
 
 metadata_1_versions = [str(LATEST_0_10_0)]
 metadata_2_versions = [str(LATEST_0_10_1), str(LATEST_0_10_2), str(LATEST_0_11_0), str(LATEST_1_0), str(LATEST_1_1),
@@ -46,7 +46,7 @@
 # -> https://issues.apache.org/jira/browse/KAFKA-14646
 # thus, we cannot test two bounce rolling upgrade because we know it's broken
 # instead we add version 2.4...3.3 to the `metadata_2_versions` upgrade list
-fk_join_versions = [str(LATEST_3_4), str(LATEST_3_5), str(LATEST_3_6)]
+fk_join_versions = [str(LATEST_3_4), str(LATEST_3_5), str(LATEST_3_6), str(LATEST_3_7)]
 
 
 """

From 5f4806fd1c0eb0ef67885a5a7f12de282f494933 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Thu, 7 Mar 2024 02:44:17 +0300
Subject: [PATCH 122/258] KAFKA-14589 [2/4] Tests of ConsoleGroupCommand
 rewritten in java (#15363)

This PR is part of #14471
It contains some of ConsoleGroupCommand tests rewritten in java.
Intention of separate PR is to reduce changes and simplify review.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../admin/DescribeConsumerGroupTest.scala     | 747 ----------------
 .../group/DescribeConsumerGroupTest.java      | 830 ++++++++++++++++++
 2 files changed, 830 insertions(+), 747 deletions(-)
 delete mode 100644 core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java

diff --git a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
deleted file mode 100644
index e98404f496f90..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
+++ /dev/null
@@ -1,747 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.admin
-
-import java.util.Properties
-import kafka.utils.{Exit, TestInfoUtils, TestUtils}
-import org.apache.kafka.clients.consumer.{ConsumerConfig, RoundRobinAssignor}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.TimeoutException
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.{MethodSource, ValueSource}
-
-import scala.concurrent.ExecutionException
-import scala.util.Random
-
-class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
-  private val describeTypeOffsets = Array(Array(""), Array("--offsets"))
-  private val describeTypeMembers = Array(Array("--members"), Array("--members", "--verbose"))
-  private val describeTypeState = Array(Array("--state"))
-  private val describeTypes = describeTypeOffsets ++ describeTypeMembers ++ describeTypeState
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeNonExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val missingGroup = "missing.group"
-
-    for (describeType <- describeTypes) {
-      // note the group to be queried is a different (non-existing) group
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", missingGroup) ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      val output = TestUtils.grabConsoleOutput(service.describeGroups())
-      assertTrue(output.contains(s"Consumer group '$missingGroup' does not exist."),
-        s"Expected error was not detected for describe option '${describeType.mkString(" ")}'")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDescribeWithMultipleSubActions(quorum: String): Unit = {
-    var exitStatus: Option[Int] = None
-    var exitMessage: Option[String] = None
-    Exit.setExitProcedure { (status, err) =>
-      exitStatus = Some(status)
-      exitMessage = err
-      throw new RuntimeException
-    }
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group, "--members", "--state")
-    try {
-      ConsumerGroupCommand.main(cgcArgs)
-    } catch {
-      case e: RuntimeException => //expected
-    } finally {
-      Exit.resetExitProcedure()
-    }
-    assertEquals(Some(1), exitStatus)
-    assertTrue(exitMessage.get.contains("Option [describe] takes at most one of these options"))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDescribeWithStateValue(quorum: String): Unit = {
-    var exitStatus: Option[Int] = None
-    var exitMessage: Option[String] = None
-    Exit.setExitProcedure { (status, err) =>
-      exitStatus = Some(status)
-      exitMessage = err
-      throw new RuntimeException
-    }
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--all-groups", "--state", "Stable")
-    try {
-      ConsumerGroupCommand.main(cgcArgs)
-    } catch {
-      case e: RuntimeException => //expected
-    } finally {
-      Exit.resetExitProcedure()
-    }
-    assertEquals(Some(1), exitStatus)
-    assertTrue(exitMessage.get.contains("Option [describe] does not take a value for [state]"))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeOffsetsOfNonExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    val group = "missing.group"
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-    // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    val (state, assignments) = service.collectGroupOffsets(group)
-    assertTrue(state.contains("Dead") && assignments.contains(List()),
-      s"Expected the state to be 'Dead', with no members in the group '$group'.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeMembersOfNonExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    val group = "missing.group"
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-    // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    val (state, assignments) = service.collectGroupMembers(group, false)
-    assertTrue(state.contains("Dead") && assignments.contains(List()),
-      s"Expected the state to be 'Dead', with no members in the group '$group'.")
-
-    val (state2, assignments2) = service.collectGroupMembers(group, true)
-    assertTrue(state2.contains("Dead") && assignments2.contains(List()),
-      s"Expected the state to be 'Dead', with no members in the group '$group' (verbose option).")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateOfNonExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    val group = "missing.group"
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-    // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    val state = service.collectGroupState(group)
-    assertTrue(state.state == "Dead" && state.numMembers == 0 &&
-      state.coordinator != null && brokers.map(_.config.brokerId).toList.contains(state.coordinator.id),
-      s"Expected the state to be 'Dead', with no members in the group '$group'."
-    )
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    for (describeType <- describeTypes) {
-      val group = this.group + describeType.mkString("")
-      // run one consumer in the group consuming from a single-partition topic
-      addConsumerGroupExecutor(numConsumers = 1, group = group, groupProtocol = groupProtocol)
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group) ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        output.trim.split("\n").length == 2 && error.isEmpty
-      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeExistingGroups(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // Create N single-threaded consumer groups from a single-partition topic
-    val groups = (for (describeType <- describeTypes) yield {
-      val group = this.group + describeType.mkString("")
-      addConsumerGroupExecutor(numConsumers = 1, group = group, groupProtocol = groupProtocol)
-      Array("--group", group)
-    }).flatten
-
-    val expectedNumLines = describeTypes.length * 2
-
-    for (describeType <- describeTypes) {
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe") ++ groups ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        val numLines = output.trim.split("\n").count(line => line.nonEmpty)
-        (numLines == expectedNumLines) && error.isEmpty
-      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeAllExistingGroups(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // Create N single-threaded consumer groups from a single-partition topic
-    for (describeType <- describeTypes) {
-      val group = this.group + describeType.mkString("")
-      addConsumerGroupExecutor(numConsumers = 1, group = group, groupProtocol = groupProtocol)
-    }
-
-    val expectedNumLines = describeTypes.length * 2
-
-    for (describeType <- describeTypes) {
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--all-groups") ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        val numLines = output.trim.split("\n").count(line => line.nonEmpty)
-        (numLines == expectedNumLines) && error.isEmpty
-      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeOffsetsOfExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 1 &&
-        assignments.get.filter(_.group == group).head.consumerId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-        assignments.get.filter(_.group == group).head.clientId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-        assignments.get.filter(_.group == group).head.host.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE)
-    }, s"Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for group $group.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeMembersOfExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(group, false)
-      state.contains("Stable") &&
-        (assignments match {
-          case Some(memberAssignments) =>
-            memberAssignments.count(_.group == group) == 1 &&
-              memberAssignments.filter(_.group == group).head.consumerId != ConsumerGroupCommand.MISSING_COLUMN_VALUE &&
-              memberAssignments.filter(_.group == group).head.clientId != ConsumerGroupCommand.MISSING_COLUMN_VALUE &&
-              memberAssignments.filter(_.group == group).head.host != ConsumerGroupCommand.MISSING_COLUMN_VALUE
-          case None =>
-            false
-        })
-    }, s"Expected a 'Stable' group status, rows and valid member information for group $group.")
-
-    val (_, assignments) = service.collectGroupMembers(group, true)
-    assignments match {
-      case None =>
-        fail(s"Expected partition assignments for members of group $group")
-      case Some(memberAssignments) =>
-        assertTrue(memberAssignments.size == 1 && memberAssignments.head.assignment.size == 1,
-          s"Expected a topic partition assigned to the single group member for group $group")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateOfExistingGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(
-      numConsumers = 1,
-      groupProtocol = groupProtocol,
-      // This is only effective when new protocol is used.
-      remoteAssignor = Some("range")
-    )
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Stable" &&
-        state.numMembers == 1 &&
-        state.assignmentStrategy == "range" &&
-        state.coordinator != null &&
-        brokers.map(_.config.brokerId).toList.contains(state.coordinator.id)
-    }, s"Expected a 'Stable' group status, with one member and round robin assignment strategy for group $group.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateOfExistingGroupWithNonDefaultAssignor(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    val expectedName = if (groupProtocol == "consumer") {
-      addConsumerGroupExecutor(numConsumers = 1, remoteAssignor = Some("range"), groupProtocol = groupProtocol)
-      "range"
-    } else {
-      addConsumerGroupExecutor(numConsumers = 1, strategy = classOf[RoundRobinAssignor].getName, groupProtocol = groupProtocol)
-      "roundrobin"
-    }
-
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Stable" &&
-        state.numMembers == 1 &&
-        state.assignmentStrategy == expectedName &&
-        state.coordinator != null &&
-        brokers.map(_.config.brokerId).toList.contains(state.coordinator.id)
-    }, s"Expected a 'Stable' group status, with one member and $expectedName assignment strategy for group $group.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeExistingGroupWithNoMembers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    for (describeType <- describeTypes) {
-      val group = this.group + describeType.mkString("")
-      // run one consumer in the group consuming from a single-partition topic
-      val executor = addConsumerGroupExecutor(numConsumers = 1, group = group, groupProtocol = groupProtocol)
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group) ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        output.trim.split("\n").length == 2 && error.isEmpty
-      }, s"Expected describe group results with one data row for describe type '${describeType.mkString(" ")}'")
-
-      // stop the consumer so the group has no active member anymore
-      executor.shutdown()
-      TestUtils.waitUntilTrue(() => {
-        TestUtils.grabConsoleError(service.describeGroups()).contains(s"Consumer group '$group' has no active members.")
-      }, s"Expected no active member in describe group results with describe type ${describeType.mkString(" ")}")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeOffsetsOfExistingGroupWithNoMembers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    val executor = addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol, syncCommit = true)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Stable") && assignments.exists(_.exists(assignment => assignment.group == group && assignment.offset.isDefined))
-    }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.")
-
-    // stop the consumer so the group has no active member anymore
-    executor.shutdown()
-
-    val (result, succeeded) = TestUtils.computeUntilTrue(service.collectGroupOffsets(group)) {
-      case (state, assignments) =>
-        val testGroupAssignments = assignments.toSeq.flatMap(_.filter(_.group == group))
-        def assignment = testGroupAssignments.head
-        state.contains("Empty") &&
-          testGroupAssignments.size == 1 &&
-          assignment.consumerId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) && // the member should be gone
-          assignment.clientId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-          assignment.host.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE)
-    }
-    val (state, assignments) = result
-    assertTrue(succeeded, s"Expected no active member in describe group results, state: $state, assignments: $assignments")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeMembersOfExistingGroupWithNoMembers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    val executor = addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(group, false)
-      state.contains("Stable") && assignments.exists(_.exists(_.group == group))
-    }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.")
-
-    // stop the consumer so the group has no active member anymore
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(group, false)
-      state.contains("Empty") && assignments.isDefined && assignments.get.isEmpty
-    }, s"Expected no member in describe group members results for group '$group'")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateOfExistingGroupWithNoMembers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run one consumer in the group consuming from a single-partition topic
-    val executor = addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Stable" &&
-        state.numMembers == 1 &&
-        state.coordinator != null &&
-        brokers.map(_.config.brokerId).toList.contains(state.coordinator.id)
-    }, s"Expected the group '$group' to initially become stable, and have a single member.")
-
-    // stop the consumer so the group has no active member anymore
-    executor.shutdown()
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Empty" && state.numMembers == 0
-    }, s"Expected the group '$group' to become empty after the only member leaving.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeWithConsumersWithoutAssignedPartitions(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    for (describeType <- describeTypes) {
-      val group = this.group + describeType.mkString("")
-      // run two consumers in the group consuming from a single-partition topic
-      addConsumerGroupExecutor(numConsumers = 2, group = group, groupProtocol = groupProtocol)
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group) ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        val expectedNumRows = if (describeTypeMembers.contains(describeType)) 3 else 2
-        error.isEmpty && output.trim.split("\n").size == expectedNumRows
-      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeOffsetsWithConsumersWithoutAssignedPartitions(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run two consumers in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 1 &&
-        assignments.get.count { x => x.group == group && x.partition.isDefined } == 1
-    }, "Expected rows for consumers with no assigned partitions in describe group results")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeMembersWithConsumersWithoutAssignedPartitions(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run two consumers in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(group, false)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 2 &&
-        assignments.get.count { x => x.group == group && x.numPartitions == 1 } == 1 &&
-        assignments.get.count { x => x.group == group && x.numPartitions == 0 } == 1 &&
-        !assignments.get.exists(_.assignment.nonEmpty)
-    }, "Expected rows for consumers with no assigned partitions in describe group results")
-
-    val (state, assignments) = service.collectGroupMembers(group, true)
-    assertTrue(state.contains("Stable") && assignments.get.count(_.assignment.nonEmpty) > 0,
-      "Expected additional columns in verbose version of describe members")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateWithConsumersWithoutAssignedPartitions(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    // run two consumers in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Stable" && state.numMembers == 2
-    }, "Expected two consumers in describe group results")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeWithMultiPartitionTopicAndMultipleConsumers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val topic2 = "foo2"
-    createTopic(topic2, 2, 1)
-
-    for (describeType <- describeTypes) {
-      val group = this.group + describeType.mkString("")
-      // run two consumers in the group consuming from a two-partition topic
-      addConsumerGroupExecutor(2, topic2, group = group, groupProtocol = groupProtocol)
-      val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group) ++ describeType
-      val service = getConsumerGroupService(cgcArgs)
-
-      TestUtils.waitUntilTrue(() => {
-        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
-        val expectedNumRows = if (describeTypeState.contains(describeType)) 2 else 3
-        error.isEmpty && output.trim.split("\n").size == expectedNumRows
-      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val topic2 = "foo2"
-    createTopic(topic2, 2, 1)
-
-    // run two consumers in the group consuming from a two-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, topic2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 2 &&
-        assignments.get.count { x => x.group == group && x.partition.isDefined } == 2 &&
-        assignments.get.count { x => x.group == group && x.partition.isEmpty } == 0
-    }, "Expected two rows (one row per consumer) in describe group results.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val topic2 = "foo2"
-    createTopic(topic2, 2, 1)
-
-    // run two consumers in the group consuming from a two-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, topic2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupMembers(group, false)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 2 &&
-        assignments.get.count { x => x.group == group && x.numPartitions == 1 } == 2 &&
-        assignments.get.count { x => x.group == group && x.numPartitions == 0 } == 0
-    }, "Expected two rows (one row per consumer) in describe group members results.")
-
-    val (state, assignments) = service.collectGroupMembers(group, true)
-    assertTrue(state.contains("Stable") && assignments.get.count(_.assignment.isEmpty) == 0,
-      "Expected additional columns in verbose version of describe members")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeStateWithMultiPartitionTopicAndMultipleConsumers(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val topic2 = "foo2"
-    createTopic(topic2, 2, 1)
-
-    // run two consumers in the group consuming from a two-partition topic
-    addConsumerGroupExecutor(numConsumers = 2, topic2, groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val state = service.collectGroupState(group)
-      state.state == "Stable" && state.group == group && state.numMembers == 2
-    }, "Expected a stable group with two members in describe group state result.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft", "kraft+kip848"))
-  def testDescribeSimpleConsumerGroup(quorum: String): Unit = {
-    // Ensure that the offsets of consumers which don't use group management are still displayed
-
-    createOffsetsTopic()
-    val topic2 = "foo2"
-    createTopic(topic2, 2, 1)
-    addSimpleGroupExecutor(Seq(new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)))
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Empty") && assignments.isDefined && assignments.get.count(_.group == group) == 2
-    }, "Expected a stable group with two members in describe group state result.")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeGroupWithShortInitializationTimeout(quorum: String, groupProtocol: String): Unit = {
-    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
-    // complete before the timeout expires
-
-    val describeType = describeTypes(Random.nextInt(describeTypes.length))
-    val group = this.group + describeType.mkString("")
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-    // set the group initialization timeout too low for the group to stabilize
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--timeout", "1", "--group", group) ++ describeType
-    val service = getConsumerGroupService(cgcArgs)
-
-    val e = assertThrows(classOf[ExecutionException], () => TestUtils.grabConsoleOutputAndError(service.describeGroups()))
-    assertEquals(classOf[TimeoutException], e.getCause.getClass)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeGroupOffsetsWithShortInitializationTimeout(quorum: String, groupProtocol: String): Unit = {
-    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
-    // complete before the timeout expires
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    // set the group initialization timeout too low for the group to stabilize
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group, "--timeout", "1")
-    val service = getConsumerGroupService(cgcArgs)
-
-    val e = assertThrows(classOf[ExecutionException], () => service.collectGroupOffsets(group))
-    assertEquals(classOf[TimeoutException], e.getCause.getClass)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeGroupMembersWithShortInitializationTimeout(quorum: String, groupProtocol: String): Unit = {
-    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
-    // complete before the timeout expires
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    // set the group initialization timeout too low for the group to stabilize
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group, "--timeout", "1")
-    val service = getConsumerGroupService(cgcArgs)
-
-    var e = assertThrows(classOf[ExecutionException], () => service.collectGroupMembers(group, false))
-    assertEquals(classOf[TimeoutException], e.getCause.getClass)
-    e = assertThrows(classOf[ExecutionException], () => service.collectGroupMembers(group, true))
-    assertEquals(classOf[TimeoutException], e.getCause.getClass)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeGroupStateWithShortInitializationTimeout(quorum: String, groupProtocol: String): Unit = {
-    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
-    // complete before the timeout expires
-
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, groupProtocol = groupProtocol)
-
-    // set the group initialization timeout too low for the group to stabilize
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group, "--timeout", "1")
-    val service = getConsumerGroupService(cgcArgs)
-
-    val e = assertThrows(classOf[ExecutionException], () => service.collectGroupState(group))
-    assertEquals(classOf[TimeoutException], e.getCause.getClass)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
-  @ValueSource(strings = Array("zk", "kraft"))
-  def testDescribeWithUnrecognizedNewConsumerOption(quorum: String): Unit = {
-    val cgcArgs = Array("--new-consumer", "--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    assertThrows(classOf[joptsimple.OptionException], () => getConsumerGroupService(cgcArgs))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testDescribeNonOffsetCommitGroup(quorum: String, groupProtocol: String): Unit = {
-    createOffsetsTopic()
-
-    val customProps = new Properties
-    // create a consumer group that never commits offsets
-    customProps.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-    // run one consumer in the group consuming from a single-partition topic
-    addConsumerGroupExecutor(numConsumers = 1, customPropsOpt = Some(customProps), groupProtocol = groupProtocol)
-
-    val cgcArgs = Array("--bootstrap-server", bootstrapServers(), "--describe", "--group", group)
-    val service = getConsumerGroupService(cgcArgs)
-
-    TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = service.collectGroupOffsets(group)
-      state.contains("Stable") &&
-        assignments.isDefined &&
-        assignments.get.count(_.group == group) == 1 &&
-        assignments.get.filter(_.group == group).head.consumerId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-        assignments.get.filter(_.group == group).head.clientId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-        assignments.get.filter(_.group == group).head.host.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE)
-    }, s"Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for non-offset-committing group $group.")
-  }
-
-}
-
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
new file mode 100644
index 0000000000000..f0277d18cd639
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
@@ -0,0 +1,830 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import kafka.admin.ConsumerGroupCommand;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.RangeAssignor;
+import org.apache.kafka.clients.consumer.RoundRobinAssignor;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import scala.Function0;
+import scala.Function1;
+import scala.Option;
+import scala.collection.Seq;
+import scala.runtime.BoxedUnit;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.apache.kafka.test.TestUtils.RANDOM;
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES;
+import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
+    private static final List<List<String>> DESCRIBE_TYPE_OFFSETS = Arrays.asList(Collections.singletonList(""), Collections.singletonList("--offsets"));
+    private static final List<List<String>> DESCRIBE_TYPE_MEMBERS = Arrays.asList(Collections.singletonList("--members"), Arrays.asList("--members", "--verbose"));
+    private static final List<List<String>> DESCRIBE_TYPE_STATE = Collections.singletonList(Collections.singletonList("--state"));
+    private static final List<List<String>> DESCRIBE_TYPES;
+
+    static {
+        List<List<String>> describeTypes = new ArrayList<>();
+
+        describeTypes.addAll(DESCRIBE_TYPE_OFFSETS);
+        describeTypes.addAll(DESCRIBE_TYPE_MEMBERS);
+        describeTypes.addAll(DESCRIBE_TYPE_STATE);
+
+        DESCRIBE_TYPES = describeTypes;
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeNonExistingGroup(String quorum, String groupProtocol) {
+        createOffsetsTopic(listenerName(), new Properties());
+        String missingGroup = "missing.group";
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            // note the group to be queried is a different (non-existing) group
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", missingGroup));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            String output = kafka.utils.TestUtils.grabConsoleOutput(describeGroups(service));
+            assertTrue(output.contains("Consumer group '" + missingGroup + "' does not exist."),
+                "Expected error was not detected for describe option '" + String.join(" ", describeType) + "'");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDescribeWithMultipleSubActions(String quorum) {
+        AtomicInteger exitStatus = new AtomicInteger(0);
+        AtomicReference<String> exitMessage = new AtomicReference<>("");
+        Exit.setExitProcedure((status, err) -> {
+            exitStatus.set(status);
+            exitMessage.set(err);
+            throw new RuntimeException();
+        });
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP, "--members", "--state"};
+        try {
+            assertThrows(RuntimeException.class, () -> ConsumerGroupCommand.main(cgcArgs));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+        assertEquals(1, exitStatus.get());
+        assertTrue(exitMessage.get().contains("Option [describe] takes at most one of these options"));
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDescribeWithStateValue(String quorum) {
+        AtomicInteger exitStatus = new AtomicInteger(0);
+        AtomicReference<String> exitMessage = new AtomicReference<>("");
+        Exit.setExitProcedure((status, err) -> {
+            exitStatus.set(status);
+            exitMessage.set(err);
+            throw new RuntimeException();
+        });
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--all-groups", "--state", "Stable"};
+        try {
+            assertThrows(RuntimeException.class, () -> ConsumerGroupCommand.main(cgcArgs));
+        } finally {
+            Exit.resetExitProcedure();
+        }
+        assertEquals(1, exitStatus.get());
+        assertTrue(exitMessage.get().contains("Option [describe] does not take a value for [state]"));
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupProtocol) {
+        String group = "missing.group";
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+        // note the group to be queried is a different (non-existing) group
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(group);
+        assertTrue(res._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res._2.map(Seq::isEmpty).getOrElse(() -> false),
+            "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeMembersOfNonExistingGroup(String quorum, String groupProtocol) {
+        String group = "missing.group";
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+        // note the group to be queried is a different (non-existing) group
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(group, false);
+        assertTrue(res._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res._2.map(Seq::isEmpty).getOrElse(() -> false),
+            "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res2 = service.collectGroupMembers(group, true);
+        assertTrue(res2._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res2._2.map(Seq::isEmpty).getOrElse(() -> false),
+            "Expected the state to be 'Dead', with no members in the group '" + group + "' (verbose option).");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateOfNonExistingGroup(String quorum, String groupProtocol) {
+        String group = "missing.group";
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+        // note the group to be queried is a different (non-existing) group
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        ConsumerGroupCommand.GroupState state = service.collectGroupState(group);
+        assertTrue(Objects.equals(state.state(), "Dead") && state.numMembers() == 0 &&
+                state.coordinator() != null && !brokers().filter(s -> s.config().brokerId() == state.coordinator().id()).isEmpty(),
+            "Expected the state to be 'Dead', with no members in the group '" + group + "'."
+        );
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeExistingGroup(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            // run one consumer in the group consuming from a single-partition topic
+            addConsumerGroupExecutor(1, TOPIC, group, groupProtocol);
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res._1.trim().split("\n").length == 2 && res._2.isEmpty();
+            }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeExistingGroups(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // Create N single-threaded consumer groups from a single-partition topic
+        List<String> groups = new ArrayList<>();
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            addConsumerGroupExecutor(1, TOPIC, group, groupProtocol);
+            groups.addAll(Arrays.asList("--group", group));
+        }
+
+        int expectedNumLines = DESCRIBE_TYPES.size() * 2;
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe"));
+            cgcArgs.addAll(groups);
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res._1.trim().split("\n")).filter(line -> !line.isEmpty()).count();
+                return (numLines == expectedNumLines) && res._2.isEmpty();
+            }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeAllExistingGroups(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // Create N single-threaded consumer groups from a single-partition topic
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            addConsumerGroupExecutor(1, TOPIC, group, groupProtocol);
+        }
+
+        int expectedNumLines = DESCRIBE_TYPES.size() * 2;
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--all-groups"));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res._1.trim().split("\n")).filter(s -> !s.isEmpty()).count();
+                return (numLines == expectedNumLines) && res._2.isEmpty();
+            }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeOffsetsOfExistingGroup(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+            Option<String> state = groupOffsets._1;
+            Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>> assignments = groupOffsets._2;
+
+            Function1<ConsumerGroupCommand.PartitionAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+
+            boolean res = state.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                assignments.isDefined() &&
+                assignments.get().count(isGrp) == 1;
+
+            if (!res)
+                return false;
+
+            @SuppressWarnings("cast")
+            ConsumerGroupCommand.PartitionAssignmentState partitionState =
+                (ConsumerGroupCommand.PartitionAssignmentState) assignments.get().filter(isGrp).head();
+
+            return !partitionState.consumerId().map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
+                !partitionState.clientId().map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
+                !partitionState.host().map(h -> h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+        }, "Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for group " + GROUP + ".");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeMembersOfExistingGroup(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> groupMembers = service.collectGroupMembers(GROUP, false);
+            Option<String> state = groupMembers._1;
+            Option<Seq<ConsumerGroupCommand.MemberAssignmentState>> assignments = groupMembers._2;
+
+            Function1<ConsumerGroupCommand.MemberAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+
+            boolean res = state.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                assignments.isDefined() &&
+                assignments.get().count(s -> Objects.equals(s.group(), GROUP)) == 1;
+
+            if (!res)
+                return false;
+
+            @SuppressWarnings("cast")
+            ConsumerGroupCommand.MemberAssignmentState assignmentState =
+                (ConsumerGroupCommand.MemberAssignmentState) assignments.get().filter(isGrp).head();
+
+            return !Objects.equals(assignmentState.consumerId(), ConsumerGroupCommand.MISSING_COLUMN_VALUE()) &&
+                !Objects.equals(assignmentState.clientId(), ConsumerGroupCommand.MISSING_COLUMN_VALUE()) &&
+                !Objects.equals(assignmentState.host(), ConsumerGroupCommand.MISSING_COLUMN_VALUE());
+        }, "Expected a 'Stable' group status, rows and valid member information for group " + GROUP + ".");
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+
+        if (res._2.isDefined()) {
+            assertTrue(res._2.get().size() == 1 && res._2.get().iterator().next().assignment().size() == 1,
+                "Expected a topic partition assigned to the single group member for group " + GROUP);
+        } else {
+            fail("Expected partition assignments for members of group " + GROUP);
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateOfExistingGroup(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(
+            1,
+            groupProtocol,
+            // This is only effective when new protocol is used.
+            Optional.of("range")
+        );
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Stable") &&
+                state.numMembers() == 1 &&
+                Objects.equals(state.assignmentStrategy(), "range") &&
+                state.coordinator() != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+        }, "Expected a 'Stable' group status, with one member and round robin assignment strategy for group " + GROUP + ".");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateOfExistingGroupWithNonDefaultAssignor(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        String expectedName;
+        if (groupProtocol.equals("consumer")) {
+            addConsumerGroupExecutor(1, groupProtocol, Optional.of("range"));
+            expectedName = "range";
+        } else {
+            addConsumerGroupExecutor(1, TOPIC, GROUP, RoundRobinAssignor.class.getName(), Optional.empty(), Optional.empty(), false, groupProtocol);
+            expectedName = "roundrobin";
+        }
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Stable") &&
+                state.numMembers() == 1 &&
+                Objects.equals(state.assignmentStrategy(), expectedName) &&
+                state.coordinator() != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+        }, "Expected a 'Stable' group status, with one member and " + expectedName + " assignment strategy for group " + GROUP + ".");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeExistingGroupWithNoMembers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            // run one consumer in the group consuming from a single-partition topic
+            ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, group, groupProtocol);
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res._1.trim().split("\n").length == 2 && res._2.isEmpty();
+            }, "Expected describe group results with one data row for describe type '" + String.join(" ", describeType) + "'");
+
+            // stop the consumer so the group has no active member anymore
+            executor.shutdown();
+            TestUtils.waitForCondition(
+                () -> kafka.utils.TestUtils.grabConsoleError(describeGroups(service)).contains("Consumer group '" + group + "' has no active members."),
+                "Expected no active member in describe group results with describe type " + String.join(" ", describeType));
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeOffsetsOfExistingGroupWithNoMembers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, GROUP, RangeAssignor.class.getName(), Optional.empty(), Optional.empty(), true, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
+                && res._2.map(c -> c.exists(assignment -> Objects.equals(assignment.group(), GROUP) && assignment.offset().isDefined())).getOrElse(() -> false);
+        }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
+
+        // stop the consumer so the group has no active member anymore
+        executor.shutdown();
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> offsets = service.collectGroupOffsets(GROUP);
+            Option<String> state = offsets._1;
+            Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>> assignments = offsets._2;
+            @SuppressWarnings("unchecked")
+            Seq<ConsumerGroupCommand.PartitionAssignmentState> testGroupAssignments = assignments.get().filter(a -> Objects.equals(a.group(), GROUP)).toSeq();
+            ConsumerGroupCommand.PartitionAssignmentState assignment = testGroupAssignments.head();
+            return state.map(s -> s.contains("Empty")).getOrElse(() -> false) &&
+                testGroupAssignments.size() == 1 &&
+                assignment.consumerId().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) && // the member should be gone
+                assignment.clientId().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
+                assignment.host().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+        }, "failed to collect group offsets");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeMembersOfExistingGroupWithNoMembers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
+                && res._2.map(c -> c.exists(m -> Objects.equals(m.group(), GROUP))).getOrElse(() -> false);
+        }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
+
+        // stop the consumer so the group has no active member anymore
+        executor.shutdown();
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res._1.map(s -> s.contains("Empty")).getOrElse(() -> false) && res._2.isDefined() && res._2.get().isEmpty();
+        }, "Expected no member in describe group members results for group '" + GROUP + "'");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateOfExistingGroupWithNoMembers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run one consumer in the group consuming from a single-partition topic
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Stable") &&
+                state.numMembers() == 1 &&
+                state.coordinator() != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+        }, "Expected the group '" + GROUP + "' to initially become stable, and have a single member.");
+
+        // stop the consumer so the group has no active member anymore
+        executor.shutdown();
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Empty") && state.numMembers() == 0;
+        }, "Expected the group '" + GROUP + "' to become empty after the only member leaving.");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeWithConsumersWithoutAssignedPartitions(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            // run two consumers in the group consuming from a single-partition topic
+            addConsumerGroupExecutor(2, TOPIC, group, groupProtocol);
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                int expectedNumRows = DESCRIBE_TYPE_MEMBERS.contains(describeType) ? 3 : 2;
+                return res._2.isEmpty() && res._1.trim().split("\n").length == expectedNumRows;
+            }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeOffsetsWithConsumersWithoutAssignedPartitions(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(2, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                res._2.isDefined() &&
+                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 1 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.partition().isDefined()) == 1;
+        }, "Expected rows for consumers with no assigned partitions in describe group results");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeMembersWithConsumersWithoutAssignedPartitions(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(2, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                res._2.isDefined() &&
+                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 1) == 1 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 0) == 1 &&
+                res._2.get().forall(s -> s.assignment().isEmpty());
+        }, "Expected rows for consumers with no assigned partitions in describe group results");
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
+                && res._2.map(c -> c.exists(s -> !s.assignment().isEmpty())).getOrElse(() -> false),
+            "Expected additional columns in verbose version of describe members");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateWithConsumersWithoutAssignedPartitions(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(2, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Stable") && state.numMembers() == 2;
+        }, "Expected two consumers in describe group results");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeWithMultiPartitionTopicAndMultipleConsumers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String topic2 = "foo2";
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        for (List<String> describeType : DESCRIBE_TYPES) {
+            String group = GROUP + String.join("", describeType);
+            // run two consumers in the group consuming from a two-partition topic
+            addConsumerGroupExecutor(2, topic2, group, groupProtocol);
+            List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group));
+            cgcArgs.addAll(describeType);
+            ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+            TestUtils.waitForCondition(() -> {
+                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                int expectedNumRows = DESCRIBE_TYPE_STATE.contains(describeType) ? 2 : 3;
+                return res._2.isEmpty() && res._1.trim().split("\n").length == expectedNumRows;
+            }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
+        }
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String topic2 = "foo2";
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a two-partition topic
+        addConsumerGroupExecutor(2, topic2, GROUP, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                res._2.isDefined() &&
+                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.partition().isDefined()) == 2 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && !x.partition().isDefined()) == 0;
+        }, "Expected two rows (one row per consumer) in describe group results.");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String topic2 = "foo2";
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a two-partition topic
+        addConsumerGroupExecutor(2, topic2, GROUP, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                res._2.isDefined() &&
+                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 1) == 2 &&
+                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 0) == 0;
+        }, "Expected two rows (one row per consumer) in describe group members results.");
+
+        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) && res._2.map(s -> s.count(x -> x.assignment().isEmpty())).getOrElse(() -> 0) == 0,
+            "Expected additional columns in verbose version of describe members");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeStateWithMultiPartitionTopicAndMultipleConsumers(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+        String topic2 = "foo2";
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        // run two consumers in the group consuming from a two-partition topic
+        addConsumerGroupExecutor(2, topic2, GROUP, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state(), "Stable") && Objects.equals(state.group(), GROUP) && state.numMembers() == 2;
+        }, "Expected a stable group with two members in describe group state result.");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft", "kraft+kip848"})
+    public void testDescribeSimpleConsumerGroup(String quorum) throws Exception {
+        // Ensure that the offsets of consumers which don't use group management are still displayed
+
+        createOffsetsTopic(listenerName(), new Properties());
+        String topic2 = "foo2";
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+        addSimpleGroupExecutor(Arrays.asList(new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)), GROUP);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res._1.map(s -> s.contains("Empty")).getOrElse(() -> false)
+                && res._2.isDefined() && res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2;
+        }, "Expected a stable group with two members in describe group state result.");
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeGroupWithShortInitializationTimeout(String quorum, String groupProtocol) {
+        // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+        // complete before the timeout expires
+
+        List<String> describeType = DESCRIBE_TYPES.get(RANDOM.nextInt(DESCRIBE_TYPES.size()));
+        String group = GROUP + String.join("", describeType);
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+        // set the group initialization timeout too low for the group to stabilize
+        List<String> cgcArgs = new ArrayList<>(Arrays.asList("--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--timeout", "1", "--group", group));
+        cgcArgs.addAll(describeType);
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
+
+        ExecutionException e = assertThrows(ExecutionException.class, service::describeGroups);
+        assertInstanceOf(TimeoutException.class, e.getCause());
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeGroupOffsetsWithShortInitializationTimeout(String quorum, String groupProtocol) {
+        // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+        // complete before the timeout expires
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+
+        // set the group initialization timeout too low for the group to stabilize
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP, "--timeout", "1"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        Throwable e = assertThrows(ExecutionException.class, () -> service.collectGroupOffsets(GROUP));
+        assertEquals(TimeoutException.class, e.getCause().getClass());
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeGroupMembersWithShortInitializationTimeout(String quorum, String groupProtocol) {
+        // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+        // complete before the timeout expires
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+
+        // set the group initialization timeout too low for the group to stabilize
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP, "--timeout", "1"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        Throwable e = assertThrows(ExecutionException.class, () -> service.collectGroupMembers(GROUP, false));
+        assertEquals(TimeoutException.class, e.getCause().getClass());
+        e = assertThrows(ExecutionException.class, () -> service.collectGroupMembers(GROUP, true));
+        assertEquals(TimeoutException.class, e.getCause().getClass());
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeGroupStateWithShortInitializationTimeout(String quorum, String groupProtocol) {
+        // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+        // complete before the timeout expires
+
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, groupProtocol);
+
+        // set the group initialization timeout too low for the group to stabilize
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP, "--timeout", "1"};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        Throwable e = assertThrows(ExecutionException.class, () -> service.collectGroupState(GROUP));
+        assertEquals(TimeoutException.class, e.getCause().getClass());
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testDescribeWithUnrecognizedNewConsumerOption(String quorum) {
+        String[] cgcArgs = new String[]{"--new-consumer", "--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        assertThrows(joptsimple.OptionException.class, () -> getConsumerGroupService(cgcArgs));
+    }
+
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
+    @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
+    public void testDescribeNonOffsetCommitGroup(String quorum, String groupProtocol) throws Exception {
+        createOffsetsTopic(listenerName(), new Properties());
+
+        Properties customProps = new Properties();
+        // create a consumer group that never commits offsets
+        customProps.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
+        // run one consumer in the group consuming from a single-partition topic
+        addConsumerGroupExecutor(1, TOPIC, GROUP, RangeAssignor.class.getName(), Optional.empty(), Optional.of(customProps), false, groupProtocol);
+
+        String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", GROUP};
+        ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
+
+        TestUtils.waitForCondition(() -> {
+            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+
+            Function1<ConsumerGroupCommand.PartitionAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+
+            boolean res = groupOffsets._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
+                groupOffsets._2.isDefined() &&
+                groupOffsets._2.get().count(isGrp) == 1;
+
+            if (!res)
+                return false;
+
+            @SuppressWarnings("cast")
+            ConsumerGroupCommand.PartitionAssignmentState assignmentState =
+                (ConsumerGroupCommand.PartitionAssignmentState) groupOffsets._2.get().filter(isGrp).head();
+
+            return assignmentState.consumerId().map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
+                assignmentState.clientId().map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
+                assignmentState.host().map(h -> !h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+        }, "Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for non-offset-committing group " + GROUP + ".");
+    }
+
+    private Function0<BoxedUnit> describeGroups(ConsumerGroupCommand.ConsumerGroupService service) {
+        return () -> {
+            try {
+                service.describeGroups();
+                return null;
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        };
+    }
+}

From ba0db81e5307cf090dc5876f3c61ddbe5fef2284 Mon Sep 17 00:00:00 2001
From: Dmitry Werner <grimekillah@gmail.com>
Date: Thu, 7 Mar 2024 13:39:16 +0500
Subject: [PATCH 123/258] KAFKA-16246: Cleanups in ConsoleConsumer (#15457)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
---
 .../kafka/tools/consumer/ConsoleConsumer.java |  51 +++-----
 .../consumer/ConsoleConsumerOptions.java      |  31 +++--
 .../consumer/ConsoleConsumerOptionsTest.java  |  52 ++++++---
 .../tools/consumer/ConsoleConsumerTest.java   | 110 +++++++++++-------
 4 files changed, 142 insertions(+), 102 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
index f84fb88c23f5f..bb5ab1443ed49 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
@@ -22,8 +22,6 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
 import java.util.concurrent.CountDownLatch;
 import java.util.regex.Pattern;
 import java.util.Collections;
@@ -68,11 +66,8 @@ public static void main(String[] args) throws Exception {
 
     public static void run(ConsoleConsumerOptions opts) {
         messageCount = 0;
-        long timeoutMs = opts.timeoutMs() >= 0 ? opts.timeoutMs() : Long.MAX_VALUE;
         Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(opts.consumerProps(), new ByteArrayDeserializer(), new ByteArrayDeserializer());
-        ConsumerWrapper consumerWrapper = opts.partitionArg().isPresent()
-            ? new ConsumerWrapper(Optional.of(opts.topicArg()), opts.partitionArg(), OptionalLong.of(opts.offsetArg()), Optional.empty(), consumer, timeoutMs)
-            : new ConsumerWrapper(Optional.of(opts.topicArg()), OptionalInt.empty(), OptionalLong.empty(), Optional.ofNullable(opts.includedTopicsArg()), consumer, timeoutMs);
+        ConsumerWrapper consumerWrapper = new ConsumerWrapper(opts, consumer);
 
         addShutdownHook(consumerWrapper, opts);
 
@@ -148,43 +143,25 @@ static boolean checkErr(PrintStream output) {
     }
 
     public static class ConsumerWrapper {
-        final Optional<String> topic;
-        final OptionalInt partitionId;
-        final OptionalLong offset;
-        final Optional<String> includedTopics;
-        final Consumer<byte[], byte[]> consumer;
-        final long timeoutMs;
         final Time time = Time.SYSTEM;
+        final long timeoutMs;
+        final Consumer<byte[], byte[]> consumer;
 
         Iterator<ConsumerRecord<byte[], byte[]>> recordIter = Collections.emptyIterator();
 
-        public ConsumerWrapper(Optional<String> topic,
-                               OptionalInt partitionId,
-                               OptionalLong offset,
-                               Optional<String> includedTopics,
-                               Consumer<byte[], byte[]> consumer,
-                               long timeoutMs) {
-            this.topic = topic;
-            this.partitionId = partitionId;
-            this.offset = offset;
-            this.includedTopics = includedTopics;
+        public ConsumerWrapper(ConsoleConsumerOptions opts, Consumer<byte[], byte[]> consumer) {
             this.consumer = consumer;
-            this.timeoutMs = timeoutMs;
-
-            if (topic.isPresent() && partitionId.isPresent() && offset.isPresent() && !includedTopics.isPresent()) {
-                seek(topic.get(), partitionId.getAsInt(), offset.getAsLong());
-            } else if (topic.isPresent() && partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
-                // default to latest if no offset is provided
-                seek(topic.get(), partitionId.getAsInt(), ListOffsetsRequest.LATEST_TIMESTAMP);
-            } else if (topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
-                consumer.subscribe(Collections.singletonList(topic.get()));
-            } else if (!topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && includedTopics.isPresent()) {
-                consumer.subscribe(Pattern.compile(includedTopics.get()));
+            timeoutMs = opts.timeoutMs();
+            Optional<String> topic = opts.topicArg();
+
+            if (topic.isPresent()) {
+                if (opts.partitionArg().isPresent()) {
+                    seek(topic.get(), opts.partitionArg().getAsInt(), opts.offsetArg());
+                } else {
+                    consumer.subscribe(Collections.singletonList(topic.get()));
+                }
             } else {
-                throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
-                        "Exactly one of 'topic' or 'include' must be provided. " +
-                        "If 'topic' is provided, an optional 'partition' may also be provided. " +
-                        "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.");
+                opts.includedTopicsArg().ifPresent(topics -> consumer.subscribe(Pattern.compile(topics)));
             }
         }
 
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
index a713afb2bf22c..aa37919515450 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
@@ -34,7 +34,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Objects;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Properties;
 import java.util.Random;
@@ -55,7 +55,7 @@ public final class ConsoleConsumerOptions extends CommandDefaultOptions {
     private final OptionSpec<String> messageFormatterConfigOpt;
     private final OptionSpec<?> resetBeginningOpt;
     private final OptionSpec<Integer> maxMessagesOpt;
-    private final OptionSpec<Integer> timeoutMsOpt;
+    private final OptionSpec<Long> timeoutMsOpt;
     private final OptionSpec<?> skipMessageOnErrorOpt;
     private final OptionSpec<String> bootstrapServerOpt;
     private final OptionSpec<String> keyDeserializerOpt;
@@ -66,6 +66,7 @@ public final class ConsoleConsumerOptions extends CommandDefaultOptions {
 
     private final Properties consumerProps;
     private final long offset;
+    private final long timeoutMs;
     private final MessageFormatter formatter;
 
     public ConsoleConsumerOptions(String[] args) throws IOException {
@@ -139,7 +140,7 @@ public ConsoleConsumerOptions(String[] args) throws IOException {
         timeoutMsOpt = parser.accepts("timeout-ms", "If specified, exit if no message is available for consumption for the specified interval.")
                 .withRequiredArg()
                 .describedAs("timeout_ms")
-                .ofType(Integer.class);
+                .ofType(Long.class);
         skipMessageOnErrorOpt = parser.accepts("skip-message-on-error", "If there is an error when processing a message, " +
                 "skip it instead of halt.");
         bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: The server(s) to connect to.")
@@ -184,12 +185,13 @@ public ConsoleConsumerOptions(String[] args) throws IOException {
         Set<String> groupIdsProvided = checkConsumerGroup(consumerPropsFromFile, extraConsumerProps);
         consumerProps = buildConsumerProps(consumerPropsFromFile, extraConsumerProps, groupIdsProvided);
         offset = parseOffset();
+        timeoutMs = parseTimeoutMs();
         formatter = buildFormatter();
     }
 
     private void checkRequiredArgs() {
-        List<String> topicOrFilterArgs = new ArrayList<>(Arrays.asList(topicArg(), includedTopicsArg()));
-        topicOrFilterArgs.removeIf(Objects::isNull);
+        List<Optional<String>> topicOrFilterArgs = new ArrayList<>(Arrays.asList(topicArg(), includedTopicsArg()));
+        topicOrFilterArgs.removeIf(arg -> !arg.isPresent());
         // user need to specify value for either --topic or one of the include filters options (--include or --whitelist)
         if (topicOrFilterArgs.size() != 1) {
             CommandLineUtils.printUsageAndExit(parser, "Exactly one of --include/--topic is required. " +
@@ -322,6 +324,11 @@ private void invalidOffset(String offset) {
                 "'earliest', 'latest', or a non-negative long.");
     }
 
+    private long parseTimeoutMs() {
+        long timeout = options.has(timeoutMsOpt) ? options.valueOf(timeoutMsOpt) : -1;
+        return timeout >= 0 ? timeout : Long.MAX_VALUE;
+    }
+
     private MessageFormatter buildFormatter() {
         MessageFormatter formatter = null;
         try {
@@ -365,16 +372,16 @@ OptionalInt partitionArg() {
         return OptionalInt.empty();
     }
 
-    String topicArg() {
-        return options.valueOf(topicOpt);
+    Optional<String> topicArg() {
+        return options.has(topicOpt) ? Optional.of(options.valueOf(topicOpt)) : Optional.empty();
     }
 
     int maxMessages() {
         return options.has(maxMessagesOpt) ? options.valueOf(maxMessagesOpt) : -1;
     }
 
-    int timeoutMs() {
-        return options.has(timeoutMsOpt) ? options.valueOf(timeoutMsOpt) : -1;
+    long timeoutMs() {
+        return timeoutMs;
     }
 
     boolean enableSystestEventsLogging() {
@@ -385,10 +392,10 @@ String bootstrapServer() {
         return options.valueOf(bootstrapServerOpt);
     }
 
-    String includedTopicsArg() {
+    Optional<String> includedTopicsArg() {
         return options.has(includeOpt)
-                ? options.valueOf(includeOpt)
-                : options.valueOf(whitelistOpt);
+                ? Optional.of(options.valueOf(includeOpt))
+                : Optional.ofNullable(options.valueOf(whitelistOpt));
     }
 
     Properties formatterArgs() throws IOException {
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
index 523122c4cdd98..3242b642cdb82 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
@@ -48,12 +48,12 @@ public void shouldParseValidConsumerValidConfig() throws IOException {
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertTrue(config.fromBeginning());
         assertFalse(config.enableSystestEventsLogging());
         assertFalse(config.skipMessageOnError());
         assertEquals(-1, config.maxMessages());
-        assertEquals(-1, config.timeoutMs());
+        assertEquals(Long.MAX_VALUE, config.timeoutMs());
     }
 
     @Test
@@ -67,7 +67,7 @@ public void shouldParseIncludeArgument() throws IOException {
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("includeTest*", config.includedTopicsArg());
+        assertEquals("includeTest*", config.includedTopicsArg().orElse(""));
         assertTrue(config.fromBeginning());
     }
 
@@ -82,7 +82,7 @@ public void shouldParseWhitelistArgument() throws IOException {
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("whitelistTest*", config.includedTopicsArg());
+        assertEquals("whitelistTest*", config.includedTopicsArg().orElse(""));
         assertTrue(config.fromBeginning());
     }
 
@@ -96,7 +96,7 @@ public void shouldIgnoreWhitelistArgumentIfIncludeSpecified() throws IOException
         };
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("includeTest*", config.includedTopicsArg());
+        assertEquals("includeTest*", config.includedTopicsArg().orElse(""));
         assertTrue(config.fromBeginning());
     }
 
@@ -112,7 +112,7 @@ public void shouldParseValidSimpleConsumerValidConfigWithNumericOffset() throws
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertTrue(config.partitionArg().isPresent());
         assertEquals(0, config.partitionArg().getAsInt());
         assertEquals(3, config.offsetArg());
@@ -191,7 +191,7 @@ public void shouldParseValidSimpleConsumerValidConfigWithStringOffset() throws E
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertTrue(config.partitionArg().isPresent());
         assertEquals(0, config.partitionArg().getAsInt());
         assertEquals(-1, config.offsetArg());
@@ -211,7 +211,7 @@ public void shouldParseValidConsumerConfigWithAutoOffsetResetLatest() throws IOE
         Properties consumerProperties = config.consumerProps();
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertFalse(config.fromBeginning());
         assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
     }
@@ -228,7 +228,7 @@ public void shouldParseValidConsumerConfigWithAutoOffsetResetEarliest() throws I
         Properties consumerProperties = config.consumerProps();
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertFalse(config.fromBeginning());
         assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
     }
@@ -246,7 +246,7 @@ public void shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBegi
         Properties consumerProperties = config.consumerProps();
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertTrue(config.fromBeginning());
         assertEquals("earliest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
     }
@@ -262,7 +262,7 @@ public void shouldParseValidConsumerConfigWithNoOffsetReset() throws IOException
         Properties consumerProperties = config.consumerProps();
 
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertFalse(config.fromBeginning());
         assertEquals("latest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG));
     }
@@ -442,7 +442,7 @@ public void shouldParseGroupIdFromBeginningGivenTogether() throws IOException {
 
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertEquals(-2, config.offsetArg());
         assertTrue(config.fromBeginning());
 
@@ -455,7 +455,7 @@ public void shouldParseGroupIdFromBeginningGivenTogether() throws IOException {
 
         config = new ConsoleConsumerOptions(args);
         assertEquals("localhost:9092", config.bootstrapServer());
-        assertEquals("test", config.topicArg());
+        assertEquals("test", config.topicArg().orElse(""));
         assertEquals(-1, config.offsetArg());
         assertFalse(config.fromBeginning());
     }
@@ -618,4 +618,30 @@ public void testParseOffset() throws Exception {
             Exit.resetExitProcedure();
         }
     }
+
+    @Test
+    public void testParseTimeoutMs() throws Exception {
+        String[] withoutTimeoutMs = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0"
+        };
+        assertEquals(Long.MAX_VALUE, new ConsoleConsumerOptions(withoutTimeoutMs).timeoutMs());
+
+        String[] negativeTimeoutMs = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0",
+            "--timeout-ms", "-100"
+        };
+        assertEquals(Long.MAX_VALUE, new ConsoleConsumerOptions(negativeTimeoutMs).timeoutMs());
+
+        String[] validTimeoutMs = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--partition", "0",
+            "--timeout-ms", "100"
+        };
+        assertEquals(100, new ConsoleConsumerOptions(validTimeoutMs).timeoutMs());
+    }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
index 008893f9c505a..f4fa6ac3be221 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
@@ -24,21 +24,19 @@
 import org.apache.kafka.common.MessageFormatter;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.requests.ListOffsetsRequest;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.util.MockTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.io.IOException;
 import java.io.PrintStream;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
+import java.util.regex.Pattern;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -58,8 +56,7 @@ public void setup() {
     }
 
     @Test
-    public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() {
-        String topic = "test";
+    public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() throws IOException {
         final Time time = new MockTime();
         final int timeoutMs = 1000;
 
@@ -71,20 +68,22 @@ public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() {
             return ConsumerRecords.EMPTY;
         });
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--timeout-ms", String.valueOf(timeoutMs)
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(topic),
-                OptionalInt.empty(),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                timeoutMs
+            new ConsoleConsumerOptions(args),
+            mockConsumer
         );
 
         assertThrows(TimeoutException.class, consumer::receive);
     }
 
     @Test
-    public void shouldResetUnConsumedOffsetsBeforeExit() {
+    public void shouldResetUnConsumedOffsetsBeforeExit() throws IOException {
         String topic = "test";
         int maxMessages = 123;
         int totalMessages = 700;
@@ -94,13 +93,16 @@ public void shouldResetUnConsumedOffsetsBeforeExit() {
         TopicPartition tp1 = new TopicPartition(topic, 0);
         TopicPartition tp2 = new TopicPartition(topic, 1);
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", topic,
+            "--timeout-ms", "1000"
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(topic),
-                OptionalInt.empty(),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
 
         mockConsumer.rebalance(Arrays.asList(tp1, tp2));
         Map<TopicPartition, Long> offsets = new HashMap<>();
@@ -165,47 +167,75 @@ public void shouldStopWhenOutputCheckErrorFails() {
 
     @Test
     @SuppressWarnings("unchecked")
-    public void shouldSeekWhenOffsetIsSet() {
+    public void shouldSeekWhenOffsetIsSet() throws IOException {
         Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
         TopicPartition tp0 = new TopicPartition("test", 0);
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--timeout-ms", "1000"
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seekToEnd(eq(Collections.singletonList(tp0)));
         consumer.cleanup();
         reset(mockConsumer);
 
-        consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.of(123L),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+        args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--offset", "123",
+            "--timeout-ms", "1000"
+        };
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(new ConsoleConsumerOptions(args), mockConsumer);
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seek(eq(tp0), eq(123L));
         consumer.cleanup();
         reset(mockConsumer);
 
-        consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.of(ListOffsetsRequest.EARLIEST_TIMESTAMP),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+        args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--offset", "earliest",
+            "--timeout-ms", "1000"
+        };
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(new ConsoleConsumerOptions(args), mockConsumer);
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seekToBeginning(eq(Collections.singletonList(tp0)));
         consumer.cleanup();
         reset(mockConsumer);
     }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void shouldWorkWithoutTopicOption() throws IOException {
+        Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--include", "includeTest*",
+            "--from-beginning"
+        };
+
+        ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
+
+        verify(mockConsumer).subscribe(any(Pattern.class));
+        consumer.cleanup();
+    }
 }

From a33c47ea4ddc810d66b6ed17ab74e40c5b7668fb Mon Sep 17 00:00:00 2001
From: Christo Lolov <lolovc@amazon.com>
Date: Thu, 7 Mar 2024 09:33:31 +0000
Subject: [PATCH 124/258] KAFKA-14133: Move consumer mock in TaskManagerTest to
 Mockito - part 2 (#15261)

The previous pull request in this series was #15112.

This pull request continues the migration of the consumer mock in TaskManagerTest test by test for easier reviews.

I envision there will be at least 1 more pull request to clean things up. For example, all calls to taskManager.setMainConsumer should be removed.

Reviewer: Bruno Cadonna <cadonna@apache.org>
---
 .../processor/internals/TaskManagerTest.java  | 379 ++++++++----------
 1 file changed, 177 insertions(+), 202 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index ba1c91e7f7197..681e69d300487 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -58,7 +58,6 @@
 import java.time.Duration;
 import java.util.ArrayList;
 
-import org.easymock.EasyMock;
 import org.easymock.EasyMockRunner;
 import org.easymock.Mock;
 import org.easymock.MockType;
@@ -112,7 +111,6 @@
 import static org.easymock.EasyMock.expect;
 import static org.easymock.EasyMock.expectLastCall;
 import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.reset;
 import static org.easymock.EasyMock.verify;
 import static org.hamcrest.CoreMatchers.hasItem;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -133,6 +131,7 @@
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.when;
@@ -185,6 +184,7 @@ public class TaskManagerTest {
     private final TaskId taskId10 = new TaskId(1, 0);
     private final TopicPartition t2p0 = new TopicPartition(topic2, 0);
     private final Set<TopicPartition> taskId10Partitions = mkSet(t2p0);
+    private final Set<TopicPartition> assignment = singleton(new TopicPartition("assignment", 0));
 
     final java.util.function.Consumer<Set<TopicPartition>> noOpResetter = partitions -> { };
 
@@ -2016,13 +2016,12 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
 
         handleAssignment(taskId00Assignment, taskId01Assignment, emptyMap());
-        reset(consumer);
-        expectConsumerAssignmentPaused(consumer);
-        replay(consumer);
 
         taskManager.handleRebalanceComplete();
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01)));
         verify(stateDirectory);
+
+        Mockito.verify(mockitoConsumer).pause(assignment);
     }
 
     @Test
@@ -2332,19 +2331,11 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         // first `handleAssignment`
-        expectRestoreToBeCompleted(consumer);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        expectLastCall();
-
-        // `handleRevocation`
-        consumer.commitSync(offsets);
-        expectLastCall();
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
-        // second `handleAssignment`
-        consumer.commitSync(offsets);
-        expectLastCall();
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2369,12 +2360,9 @@ public void closeClean() {
             }
         };
 
-        // first `handleAssignment`
-        expectRestoreToBeCompleted(consumer);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         taskManager.handleRevocation(taskId00Partitions);
@@ -2399,7 +2387,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
         // `handleAssignment`
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
@@ -2412,11 +2400,11 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
         expectLockObtainedFor();
         replay(stateDirectory);
 
+        taskManager.setMainConsumer(mockitoConsumer);
+
         taskManager.handleRebalanceStart(emptySet());
         assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01)));
 
-        replay(consumer);
-
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -2455,15 +2443,13 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         // `handleAssignment`
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         // `handleAssignment`
-        consumer.commitSync(offsets);
-        expectLastCall();
         doThrow(new RuntimeException("KABOOM!")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2529,10 +2515,12 @@ public void postCommit(final boolean enforceCheckpoint) {
         };
 
         // `handleAssignment`
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(taskId00Partitions);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        expect(consumer.assignment()).andReturn(taskId00Partitions);
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
@@ -2548,7 +2536,6 @@ public void postCommit(final boolean enforceCheckpoint) {
         assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
 
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2564,10 +2551,12 @@ public void suspend() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(taskId00Partitions);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        expect(consumer.assignment()).andReturn(taskId00Partitions);
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
@@ -2581,7 +2570,6 @@ public void suspend() {
         assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
 
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2592,20 +2580,20 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final StateMachineTask nonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
 
-        final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>(taskId00Assignment);
-        assignment.putAll(taskId01Assignment);
+        final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>(taskId00Assignment);
+        firstAssignment.putAll(taskId01Assignment);
 
         // `handleAssignment`
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedTask, nonCorruptedTask));
-        expectRestoreToBeCompleted(consumer);
-        expect(consumer.assignment()).andReturn(taskId00Partitions);
-        // check that we should not commit empty map either
-        consumer.commitSync(eq(emptyMap()));
-        expectLastCall().andStubThrow(new AssertionError("should not invoke commitSync when offset map is empty"));
-        replay(consumer);
 
-        taskManager.handleAssignment(assignment, emptyMap());
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(taskId00Partitions);
+
+        taskManager.setMainConsumer(mockitoConsumer);
+
+        taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
 
         assertThat(nonCorruptedTask.state(), is(Task.State.RUNNING));
@@ -2618,7 +2606,8 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         assertThat(nonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet()));
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
-        verify(consumer);
+        // check that we should not commit empty map either
+        Mockito.verify(mockitoConsumer, never()).commitSync(emptyMap());
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2637,8 +2626,9 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         // `handleAssignment`
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
             .thenReturn(asList(corruptedTask, nonRunningNonCorruptedTask));
-        expect(consumer.assignment()).andReturn(taskId00Partitions);
-        replay(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(taskId00Partitions);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignment, emptyMap());
 
@@ -2650,7 +2640,6 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
         assertFalse(nonRunningNonCorruptedTask.commitPrepared);
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2732,9 +2721,9 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             .thenReturn(singleton(runningNonCorruptedActive));
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singleton(corruptedStandby));
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId01Assignment, taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2751,7 +2740,6 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
 
         assertThat(corruptedStandby.commitPrepared, is(true));
         assertThat(corruptedStandby.state(), is(Task.State.CREATED));
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2768,21 +2756,23 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
         uncorruptedActive.setCommitNeeded();
 
         // handleAssignment
-        final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
-        assignment.putAll(taskId00Assignment);
-        assignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        final Map<TaskId, Set<TopicPartition>> firstAssignement = new HashMap<>();
+        firstAssignement.putAll(taskId00Assignment);
+        firstAssignement.putAll(taskId01Assignment);
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignement)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
-        expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
-
-        replay(consumer, stateDirectory);
+        replay(stateDirectory);
 
         uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);
 
-        taskManager.handleAssignment(assignment, emptyMap());
+        taskManager.setMainConsumer(mockitoConsumer);
+
+        taskManager.handleAssignment(firstAssignement, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
         assertThat(uncorruptedActive.state(), is(Task.State.RUNNING));
@@ -2800,7 +2790,6 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
         assertThat(uncorruptedActive.commitCompleted, is(false));
 
         assertThat(uncorruptedActive.state(), is(State.RUNNING));
-        verify(consumer);
     }
 
     @Test
@@ -2818,22 +2807,21 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);
 
         // handleAssignment
-        final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
-        assignment.putAll(taskId00Assignment);
-        assignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>();
+        firstAssignment.putAll(taskId00Assignment);
+        firstAssignment.putAll(taskId01Assignment);
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
-        consumer.commitSync(offsets);
-        expectLastCall().andThrow(new TimeoutException());
+        doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(offsets);
 
-        expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
-
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
-        taskManager.handleAssignment(assignment, emptyMap());
+        taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
         assertThat(uncorruptedActive.state(), is(Task.State.RUNNING));
@@ -2861,7 +2849,6 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
 
         assertThat(corruptedActive.state(), is(Task.State.CREATED));
         assertThat(uncorruptedActive.state(), is(Task.State.CREATED));
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2893,24 +2880,24 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         uncorruptedActiveTask.setCommittableOffsetsAndMetadata(offsets);
 
         // handleAssignment
-        final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
-        assignment.putAll(taskId00Assignment);
-        assignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>();
+        firstAssignment.putAll(taskId00Assignment);
+        firstAssignment.putAll(taskId01Assignment);
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask));
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        expect(consumer.groupMetadata()).andReturn(groupMetadata);
+        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
 
         doThrow(new TimeoutException()).when(producer).commitTransaction(offsets, groupMetadata);
 
-        expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
-
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
-        taskManager.handleAssignment(assignment, emptyMap());
+        taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
         assertThat(uncorruptedActiveTask.state(), is(Task.State.RUNNING));
@@ -2944,7 +2931,6 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         assertThat(uncorruptedActiveTask.state(), is(Task.State.CREATED));
         assertThat(corruptedTaskChangelogMarkedAsCorrupted.get(), is(true));
         assertThat(uncorruptedTaskChangelogMarkedAsCorrupted.get(), is(true));
-        verify(consumer);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions);
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions);
     }
@@ -2978,16 +2964,16 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             mkEntry(taskId02, taskId02Partitions)
         );
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded));
-        expectLastCall();
-        consumer.commitSync(expectedCommittedOffsets);
-        expectLastCall().andThrow(new TimeoutException());
-        expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
-        replay(consumer);
+        doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(expectedCommittedOffsets);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3038,19 +3024,19 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             mkEntry(taskId02, taskId02Partitions)
             );
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment())
+            .thenReturn(assignment)
+            .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        expect(consumer.groupMetadata()).andReturn(groupMetadata);
+        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
 
         doThrow(new TimeoutException()).when(producer).commitTransaction(expectedCommittedOffsets, groupMetadata);
 
-        expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
-
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3077,11 +3063,10 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
     public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00));
-        consumer.commitSync(Collections.emptyMap());
-        expectLastCall();
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(emptyMap(), taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3098,12 +3083,11 @@ public void shouldAddNonResumedSuspendedTasks() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
-        // expect these calls twice (because we're going to tryToCompleteRestoration twice)
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3115,18 +3099,20 @@ public void shouldAddNonResumedSuspendedTasks() {
         assertThat(task00.state(), is(Task.State.RUNNING));
         assertThat(task01.state(), is(Task.State.RUNNING));
 
+        // expect these calls twice (because we're going to tryToCompleteRestoration twice)
         Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
+        Mockito.verify(mockitoConsumer, times(2)).assignment();
+        Mockito.verify(mockitoConsumer, times(2)).resume(assignment);
     }
 
     @Test
     public void shouldUpdateInputPartitionsAfterRebalance() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
-        // expect these calls twice (because we're going to tryToCompleteRestoration twice)
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3138,7 +3124,9 @@ public void shouldUpdateInputPartitionsAfterRebalance() {
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
         assertEquals(newPartitionsSet, task00.inputPartitions());
-        verify(consumer);
+        // expect these calls twice (because we're going to tryToCompleteRestoration twice)
+        Mockito.verify(mockitoConsumer, times(2)).resume(assignment);
+        Mockito.verify(mockitoConsumer, times(2)).assignment();
         Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
     }
 
@@ -3244,12 +3232,10 @@ public void shouldSuspendActiveTasksDuringRevocation() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task00.setCommittableOffsetsAndMetadata(offsets);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        consumer.commitSync(offsets);
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3293,7 +3279,7 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = mkMap(
             mkEntry(taskId10, taskId10Partitions)
         );
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
@@ -3303,20 +3289,14 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
             .thenReturn(singletonList(task10));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        expect(consumer.groupMetadata()).andReturn(groupMetadata);
-        producer.commitTransaction(expectedCommittedOffsets, groupMetadata);
-        expectLastCall();
+        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
 
         task00.committedOffsets();
-        EasyMock.expectLastCall();
         task01.committedOffsets();
-        EasyMock.expectLastCall();
         task02.committedOffsets();
-        EasyMock.expectLastCall();
         task10.committedOffsets();
-        EasyMock.expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3331,6 +3311,8 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
         assertThat(task01.commitNeeded, is(false));
         assertThat(task02.commitPrepared, is(false));
         assertThat(task10.commitPrepared, is(false));
+
+        Mockito.verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata);
     }
 
     @Test
@@ -3364,16 +3346,14 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = mkMap(
             mkEntry(taskId10, taskId10Partitions)
         );
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task10));
-        consumer.commitSync(expectedCommittedOffsets);
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3390,6 +3370,8 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         assertThat(task01.commitPrepared, is(true));
         assertThat(task02.commitPrepared, is(false));
         assertThat(task10.commitPrepared, is(false));
+
+        Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -3404,12 +3386,12 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() {
         final Map<TaskId, Set<TopicPartition>> assignmentActive = singletonMap(taskId00, taskId00Partitions);
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = singletonMap(taskId10, taskId10Partitions);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3434,12 +3416,12 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() {
         final Map<TaskId, Set<TopicPartition>> assignmentActive = singletonMap(taskId00, taskId00Partitions);
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = singletonMap(taskId10, taskId10Partitions);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3478,17 +3460,15 @@ public void suspend() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
 
         assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions));
         assertThat(task00.state(), is(Task.State.SUSPENDED));
-
-        verify(consumer);
     }
 
     @Test
@@ -3938,10 +3918,12 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() {
     @Test
     public void shouldInitializeNewActiveTasks() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
+
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3950,17 +3932,17 @@ public void shouldInitializeNewActiveTasks() {
         assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         // verifies that we actually resume the assignment at the end of restoration.
-        verify(consumer);
+        Mockito.verify(mockitoConsumer).resume(assignment);
     }
 
     @Test
     public void shouldInitialiseNewStandbyTasks() {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(emptyMap(), taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3972,7 +3954,6 @@ public void shouldInitialiseNewStandbyTasks() {
 
     @Test
     public void shouldHandleRebalanceEvents() {
-        final Set<TopicPartition> assignment = singleton(new TopicPartition("assignment", 0));
         taskManager.setMainConsumer(mockitoConsumer);
         when(mockitoConsumer.assignment()).thenReturn(assignment);
         expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>());
@@ -3992,15 +3973,13 @@ public void shouldCommitActiveAndStandbyTasks() {
         task00.setCommittableOffsetsAndMetadata(offsets);
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
             .thenReturn(singletonList(task01));
-        consumer.commitSync(offsets);
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4014,6 +3993,8 @@ public void shouldCommitActiveAndStandbyTasks() {
         assertThat(taskManager.commitAll(), equalTo(2));
         assertThat(task00.commitNeeded, is(false));
         assertThat(task01.commitNeeded, is(false));
+
+        Mockito.verify(mockitoConsumer).commitSync(offsets);
     }
 
     @Test
@@ -4036,15 +4017,13 @@ public void shouldCommitProvidedTasksIfNeeded() {
             mkEntry(taskId05, taskId05Partitions)
         );
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(Arrays.asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(Arrays.asList(task03, task04, task05));
 
-        consumer.commitSync(eq(emptyMap()));
-
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4070,10 +4049,10 @@ public void shouldCommitProvidedTasksIfNeeded() {
     public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4094,13 +4073,15 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
         makeTaskFolders(taskId00.toString(), taskId01.toString());
         expectDirectoryNotEmpty(taskId00, taskId01);
         expectLockObtainedFor(taskId00, taskId01);
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
             .thenReturn(singletonList(task01));
 
-        replay(stateDirectory, consumer);
+        replay(stateDirectory);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4200,9 +4181,10 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4225,10 +4207,10 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(emptyMap(), taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4259,10 +4241,10 @@ public Map<TopicPartition, Long> purgeableOffsets() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4294,9 +4276,10 @@ public Map<TopicPartition, Long> purgeableOffsets() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4317,14 +4300,14 @@ public Map<TopicPartition, Long> purgeableOffsets() {
     public void shouldIgnorePurgeDataErrors() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
 
         final KafkaFutureImpl<DeletedRecords> futureDeletedRecords = new KafkaFutureImpl<>();
         final DeleteRecordsResult deleteRecordsResult = new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords));
         futureDeletedRecords.completeExceptionally(new Exception("KABOOM!"));
         when(adminClient.deleteRecords(any())).thenReturn(deleteRecordsResult);
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.addTask(task00);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4366,15 +4349,13 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
             mkEntry(taskId10, taskId10Partitions)
         );
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02, task03));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task04));
-        consumer.commitSync(expectedCommittedOffsets);
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4399,6 +4380,8 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
         task04.setCommitRequested();
 
         assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3));
+
+        Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -4406,16 +4389,17 @@ public void shouldProcessActiveTasks() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
 
-        final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>();
-        assignment.put(taskId00, taskId00Partitions);
-        assignment.put(taskId01, taskId01Partitions);
+        final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>();
+        firstAssignment.put(taskId00, taskId00Partitions);
+        firstAssignment.put(taskId01, taskId01Partitions);
 
-        expectRestoreToBeCompleted(consumer);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(Arrays.asList(task00, task01));
-        replay(consumer);
 
-        taskManager.handleAssignment(assignment, emptyMap());
+        taskManager.setMainConsumer(mockitoConsumer);
+
+        taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -4523,9 +4507,10 @@ public boolean process(final long wallClockTime) {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4547,10 +4532,11 @@ public boolean process(final long wallClockTime) {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4575,9 +4561,10 @@ public boolean maybePunctuateStreamTime() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4596,9 +4583,10 @@ public boolean maybePunctuateStreamTime() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4622,9 +4610,10 @@ public boolean maybePunctuateSystemTime() {
             }
         };
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        replay(consumer);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4660,12 +4649,10 @@ public void shouldHaveRemainingPartitionsUncleared() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task00.setCommittableOffsetsAndMetadata(offsets);
 
-        expectRestoreToBeCompleted(consumer);
+        when(mockitoConsumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        consumer.commitSync(offsets);
-        expectLastCall();
 
-        replay(consumer);
+        taskManager.setMainConsumer(mockitoConsumer);
 
         try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) {
             appender.setClassLoggerToDebug(TaskManager.class);
@@ -4821,8 +4808,9 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
         when(standbyTaskCreator.createTasks(standbyAssignment)).thenReturn(standbyTasks);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(allActiveTasksAssignment))).thenReturn(allActiveTasks);
 
-        expectRestoreToBeCompleted(consumer);
-        replay(consumer);
+        lenient().when(mockitoConsumer.assignment()).thenReturn(assignment);
+
+        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.handleAssignment(allActiveTasksAssignment, standbyAssignment);
         taskManager.tryToCompleteRestoration(time.milliseconds(), null);
@@ -4870,12 +4858,6 @@ private void expectDirectoryNotEmpty(final TaskId... tasks) {
         }
     }
 
-    private static void expectConsumerAssignmentPaused(final Consumer<byte[], byte[]> consumer) {
-        final Set<TopicPartition> assignment = singleton(new TopicPartition("assignment", 0));
-        expect(consumer.assignment()).andReturn(assignment);
-        consumer.pause(assignment);
-    }
-
     @Test
     public void shouldThrowTaskMigratedExceptionOnCommitFailed() {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
@@ -5135,13 +5117,6 @@ public void shouldListNotPausedTasks() {
         assertEquals(taskManager.notPausedTasks().size(), 0);
     }
 
-    private static void expectRestoreToBeCompleted(final Consumer<byte[], byte[]> consumer) {
-        final Set<TopicPartition> assignment = singleton(new TopicPartition("assignment", 0));
-        expect(consumer.assignment()).andReturn(assignment);
-        consumer.resume(assignment);
-        expectLastCall();
-    }
-
     private static KafkaFutureImpl<DeletedRecords> completedFuture() {
         final KafkaFutureImpl<DeletedRecords> futureDeletedRecords = new KafkaFutureImpl<>();
         futureDeletedRecords.complete(null);

From 86e7885a81c7132e522d2c59dc6fcf81026cc60d Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Thu, 7 Mar 2024 06:00:21 -0800
Subject: [PATCH 125/258] KAFKA-16100: Add timeout to all the
 CompletableApplicationEvents (#15455)

This is part of the larger task of enforcing the timeouts for application events, per KAFKA-15974.

This takes a first step by adding a Timer to all of the CompletableApplicationEvent subclasses. For the few classes that already included a timeout, this refactors them to use the Timer mechanism instead.

Reviewers: Andrew Schofield <aschofield@confluent.io>, Bruno Cadonna <cadonna@apache.org>
---
 .../internals/AsyncKafkaConsumer.java         | 41 +++++++++++--------
 .../events/AbstractTopicMetadataEvent.java    | 17 ++------
 .../events/AllTopicsMetadataEvent.java        |  6 ++-
 .../events/ApplicationEventProcessor.java     | 25 ++---------
 .../internals/events/AsyncCommitEvent.java    |  2 +-
 .../internals/events/CommitEvent.java         | 20 +++++++--
 .../events/CompletableApplicationEvent.java   | 20 ++++++++-
 .../events/FetchCommittedOffsetsEvent.java    | 17 ++------
 .../internals/events/LeaveOnCloseEvent.java   |  6 ++-
 .../internals/events/ListOffsetsEvent.java    |  5 ++-
 .../internals/events/ResetPositionsEvent.java |  6 ++-
 .../internals/events/SyncCommitEvent.java     | 21 ++--------
 .../internals/events/TopicMetadataEvent.java  |  6 ++-
 .../internals/events/UnsubscribeEvent.java    |  6 ++-
 .../events/ValidatePositionsEvent.java        |  6 ++-
 .../internals/ConsumerNetworkThreadTest.java  | 19 ++++++---
 .../events/ApplicationEventProcessorTest.java | 17 +++-----
 17 files changed, 120 insertions(+), 120 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index 5354503c01619..fcd57469c2ad9 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -938,13 +938,14 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
                 return Collections.emptyMap();
             }
 
+            final Timer timer = time.timer(timeout);
             final FetchCommittedOffsetsEvent event = new FetchCommittedOffsetsEvent(
                 partitions,
-                timeout.toMillis());
+                timer);
             wakeupTrigger.setActiveTask(event.future());
             try {
                 final Map<TopicPartition, OffsetAndMetadata> committedOffsets = applicationEventHandler.addAndGet(event,
-                    time.timer(timeout));
+                    timer);
                 committedOffsets.forEach(this::updateLastSeenEpochIfNewer);
                 return committedOffsets;
             } catch (TimeoutException e) {
@@ -990,11 +991,12 @@ public List<PartitionInfo> partitionsFor(String topic, Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final TopicMetadataEvent topicMetadataEvent = new TopicMetadataEvent(topic, timeout.toMillis());
+            final Timer timer = time.timer(timeout);
+            final TopicMetadataEvent topicMetadataEvent = new TopicMetadataEvent(topic, timer);
             wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
                 Map<String, List<PartitionInfo>> topicMetadata =
-                        applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
+                        applicationEventHandler.addAndGet(topicMetadataEvent, timer);
 
                 return topicMetadata.getOrDefault(topic, Collections.emptyList());
             } finally {
@@ -1018,10 +1020,11 @@ public Map<String, List<PartitionInfo>> listTopics(Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final AllTopicsMetadataEvent topicMetadataEvent = new AllTopicsMetadataEvent(timeout.toMillis());
+            final Timer timer = time.timer(timeout);
+            final AllTopicsMetadataEvent topicMetadataEvent = new AllTopicsMetadataEvent(timer);
             wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
-                return applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
+                return applicationEventHandler.addAndGet(topicMetadataEvent, timer);
             } finally {
                 wakeupTrigger.clearTask();
             }
@@ -1089,16 +1092,18 @@ public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(Map<TopicPartitio
             if (timestampsToSearch.isEmpty()) {
                 return Collections.emptyMap();
             }
+            final Timer timer = time.timer(timeout);
             final ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampsToSearch,
-                true);
+                true,
+                timer);
 
             // If timeout is set to zero return empty immediately; otherwise try to get the results
             // and throw timeout exception if it cannot complete in time.
             if (timeout.toMillis() == 0L)
                 return listOffsetsEvent.emptyResult();
 
-            return applicationEventHandler.addAndGet(listOffsetsEvent, time.timer(timeout));
+            return applicationEventHandler.addAndGet(listOffsetsEvent, timer);
         } finally {
             release();
         }
@@ -1139,12 +1144,14 @@ private Map<TopicPartition, Long> beginningOrEndOffset(Collection<TopicPartition
             Map<TopicPartition, Long> timestampToSearch = partitions
                 .stream()
                 .collect(Collectors.toMap(Function.identity(), tp -> timestamp));
+            Timer timer = time.timer(timeout);
             ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampToSearch,
-                false);
+                false,
+                timer);
             Map<TopicPartition, OffsetAndTimestamp> offsetAndTimestampMap = applicationEventHandler.addAndGet(
                 listOffsetsEvent,
-                time.timer(timeout));
+                timer);
             return offsetAndTimestampMap
                 .entrySet()
                 .stream()
@@ -1274,7 +1281,7 @@ void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstEx
         completeQuietly(
             () -> {
                 maybeRevokePartitions();
-                applicationEventHandler.addAndGet(new LeaveOnCloseEvent(), timer);
+                applicationEventHandler.addAndGet(new LeaveOnCloseEvent(timer), timer);
             },
             "Failed to send leaveGroup heartbeat with a timeout(ms)=" + timer.timeoutMs(), firstException);
     }
@@ -1351,7 +1358,7 @@ public void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration
         long commitStart = time.nanoseconds();
         try {
             Timer requestTimer = time.timer(timeout.toMillis());
-            SyncCommitEvent syncCommitEvent = new SyncCommitEvent(offsets, timeout.toMillis());
+            SyncCommitEvent syncCommitEvent = new SyncCommitEvent(offsets, requestTimer);
             CompletableFuture<Void> commitFuture = commit(syncCommitEvent);
             wakeupTrigger.setActiveTask(commitFuture);
             ConsumerUtils.getResult(commitFuture, requestTimer);
@@ -1465,10 +1472,10 @@ public void unsubscribe() {
         try {
             fetchBuffer.retainAll(Collections.emptySet());
             if (groupMetadata.get().isPresent()) {
-                UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent();
+                Timer timer = time.timer(Long.MAX_VALUE);
+                UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent(timer);
                 applicationEventHandler.add(unsubscribeEvent);
                 log.info("Unsubscribing all topics or patterns and assigned partitions");
-                Timer timer = time.timer(Long.MAX_VALUE);
 
                 try {
                     processBackgroundEvents(backgroundEventProcessor, unsubscribeEvent.future(), timer);
@@ -1579,7 +1586,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // Validate positions using the partition leader end offsets, to detect if any partition
             // has been truncated due to a leader change. This will trigger an OffsetForLeaderEpoch
             // request, retrieve the partition end offsets, and validate the current position against it.
-            applicationEventHandler.addAndGet(new ValidatePositionsEvent(), timer);
+            applicationEventHandler.addAndGet(new ValidatePositionsEvent(timer), timer);
 
             cachedSubscriptionHasAllFetchPositions = subscriptions.hasAllFetchPositions();
             if (cachedSubscriptionHasAllFetchPositions) return true;
@@ -1602,7 +1609,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // which are awaiting reset. This will trigger a ListOffset request, retrieve the
             // partition offsets according to the strategy (ex. earliest, latest), and update the
             // positions.
-            applicationEventHandler.addAndGet(new ResetPositionsEvent(), timer);
+            applicationEventHandler.addAndGet(new ResetPositionsEvent(timer), timer);
             return true;
         } catch (TimeoutException e) {
             return false;
@@ -1635,7 +1642,7 @@ private boolean initWithCommittedOffsetsIfNeeded(Timer timer) {
             final FetchCommittedOffsetsEvent event =
                 new FetchCommittedOffsetsEvent(
                     initializingPartitions,
-                    timer.remainingMs());
+                    timer);
             final Map<TopicPartition, OffsetAndMetadata> offsets = applicationEventHandler.addAndGet(event, timer);
             refreshCommittedOffsets(offsets, metadata, subscriptions);
             return true;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
index 31c21817d85a3..3347002cc6fea 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java
@@ -17,25 +17,14 @@
 package org.apache.kafka.clients.consumer.internals.events;
 
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.List;
 import java.util.Map;
 
 public abstract class AbstractTopicMetadataEvent extends CompletableApplicationEvent<Map<String, List<PartitionInfo>>> {
 
-    private final long timeoutMs;
-
-    protected AbstractTopicMetadataEvent(final Type type, final long timeoutMs) {
-        super(type);
-        this.timeoutMs = timeoutMs;
-    }
-
-    public long timeoutMs() {
-        return timeoutMs;
-    }
-
-    @Override
-    public String toStringBase() {
-        return super.toStringBase() + ", timeoutMs=" + timeoutMs;
+    protected AbstractTopicMetadataEvent(final Type type, final Timer timer) {
+        super(type, timer);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
index 154703aaee15b..bda18e642105b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java
@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 public class AllTopicsMetadataEvent extends AbstractTopicMetadataEvent {
 
-    public AllTopicsMetadataEvent(final long timeoutMs) {
-        super(Type.ALL_TOPICS_METADATA, timeoutMs);
+    public AllTopicsMetadataEvent(final Timer timer) {
+        super(Type.ALL_TOPICS_METADATA, timer);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
index c86aa8815f250..3382530746532 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
@@ -165,8 +165,7 @@ private void process(final SyncCommitEvent event) {
         }
 
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        long expirationTimeoutMs = getExpirationTimeForTimeout(event.retryTimeoutMs());
-        CompletableFuture<Void> future = manager.commitSync(event.offsets(), expirationTimeoutMs);
+        CompletableFuture<Void> future = manager.commitSync(event.offsets(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -177,8 +176,7 @@ private void process(final FetchCommittedOffsetsEvent event) {
             return;
         }
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        long expirationTimeMs = getExpirationTimeForTimeout(event.timeout());
-        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future = manager.fetchOffsets(event.partitions(), expirationTimeMs);
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future = manager.fetchOffsets(event.partitions(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -250,16 +248,14 @@ private void process(final ValidatePositionsEvent event) {
     }
 
     private void process(final TopicMetadataEvent event) {
-        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
         final CompletableFuture<Map<String, List<PartitionInfo>>> future =
-                requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), expirationTimeMs);
+                requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
     private void process(final AllTopicsMetadataEvent event) {
-        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
         final CompletableFuture<Map<String, List<PartitionInfo>>> future =
-                requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(expirationTimeMs);
+                requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -296,19 +292,6 @@ private void process(final LeaveOnCloseEvent event) {
         future.whenComplete(complete(event.future()));
     }
 
-    /**
-     * @return Expiration time in milliseconds calculated with the current time plus the given
-     * timeout. Returns Long.MAX_VALUE if the expiration overflows it.
-     * Visible for testing.
-     */
-    long getExpirationTimeForTimeout(final long timeoutMs) {
-        long expiration = System.currentTimeMillis() + timeoutMs;
-        if (expiration < 0) {
-            return Long.MAX_VALUE;
-        }
-        return expiration;
-    }
-
     private <T> BiConsumer<? super T, ? super Throwable> complete(final CompletableFuture<T> b) {
         return (value, exception) -> {
             if (exception != null)
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
index 2f03fdfb1e543..c36f0534b3671 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java
@@ -27,6 +27,6 @@
 public class AsyncCommitEvent extends CommitEvent {
 
     public AsyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets) {
-        super(Type.COMMIT_ASYNC, offsets);
+        super(Type.COMMIT_ASYNC, offsets, Long.MAX_VALUE);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
index 253d27e2573b8..1da7b84039ab8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Collections;
 import java.util.Map;
@@ -29,15 +30,28 @@ public abstract class CommitEvent extends CompletableApplicationEvent<Void> {
      */
     private final Map<TopicPartition, OffsetAndMetadata> offsets;
 
-    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets) {
-        super(type);
-        this.offsets = Collections.unmodifiableMap(offsets);
+    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets, final Timer timer) {
+        super(type, timer);
+        this.offsets = validate(offsets);
+    }
+
+    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets, final long deadlineMs) {
+        super(type, deadlineMs);
+        this.offsets = validate(offsets);
+    }
 
+    /**
+     * Validates the offsets are not negative and then returns the given offset map as
+     * {@link Collections#unmodifiableMap(Map) as unmodifiable}.
+     */
+    private static Map<TopicPartition, OffsetAndMetadata> validate(final Map<TopicPartition, OffsetAndMetadata> offsets) {
         for (OffsetAndMetadata offsetAndMetadata : offsets.values()) {
             if (offsetAndMetadata.offset() < 0) {
                 throw new IllegalArgumentException("Invalid offset: " + offsetAndMetadata.offset());
             }
         }
+
+        return Collections.unmodifiableMap(offsets);
     }
 
     public Map<TopicPartition, OffsetAndMetadata> offsets() {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
index a62c3aaa4c43b..dae9e9f1017ba 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java
@@ -16,6 +16,9 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -27,10 +30,19 @@
 public abstract class CompletableApplicationEvent<T> extends ApplicationEvent implements CompletableEvent<T> {
 
     private final CompletableFuture<T> future;
+    private final long deadlineMs;
 
-    protected CompletableApplicationEvent(final Type type) {
+    protected CompletableApplicationEvent(final Type type, final Timer timer) {
         super(type);
         this.future = new CompletableFuture<>();
+        Objects.requireNonNull(timer);
+        this.deadlineMs = timer.remainingMs() + timer.currentTimeMs();
+    }
+
+    protected CompletableApplicationEvent(final Type type, final long deadlineMs) {
+        super(type);
+        this.future = new CompletableFuture<>();
+        this.deadlineMs = deadlineMs;
     }
 
     @Override
@@ -38,8 +50,12 @@ public CompletableFuture<T> future() {
         return future;
     }
 
+    public long deadlineMs() {
+        return deadlineMs;
+    }
+
     @Override
     protected String toStringBase() {
-        return super.toStringBase() + ", future=" + future;
+        return super.toStringBase() + ", future=" + future + ", deadlineMs=" + deadlineMs;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
index 7cf56b990b0d9..980a8f1104261 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Collections;
 import java.util.Map;
@@ -30,27 +31,17 @@ public class FetchCommittedOffsetsEvent extends CompletableApplicationEvent<Map<
      */
     private final Set<TopicPartition> partitions;
 
-    /**
-     * Time until which the request will be retried if it fails with a retriable error.
-     */
-    private final long timeoutMs;
-
-    public FetchCommittedOffsetsEvent(final Set<TopicPartition> partitions, final long timeoutMs) {
-        super(Type.FETCH_COMMITTED_OFFSETS);
+    public FetchCommittedOffsetsEvent(final Set<TopicPartition> partitions, final Timer timer) {
+        super(Type.FETCH_COMMITTED_OFFSETS, timer);
         this.partitions = Collections.unmodifiableSet(partitions);
-        this.timeoutMs = timeoutMs;
     }
 
     public Set<TopicPartition> partitions() {
         return partitions;
     }
 
-    public long timeout() {
-        return timeoutMs;
-    }
-
     @Override
     public String toStringBase() {
-        return super.toStringBase() + ", partitions=" + partitions + ", partitions=" + partitions;
+        return super.toStringBase() + ", partitions=" + partitions;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
index 5ee19a7cc02da..e77b4dfb2893c 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java
@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 public class LeaveOnCloseEvent extends CompletableApplicationEvent<Void> {
 
-    public LeaveOnCloseEvent() {
-        super(Type.LEAVE_ON_CLOSE);
+    public LeaveOnCloseEvent(final Timer timer) {
+        super(Type.LEAVE_ON_CLOSE, timer);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
index fd3b321173f11..e218705846e19 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ListOffsetsEvent.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Collections;
 import java.util.HashMap;
@@ -36,8 +37,8 @@ public class ListOffsetsEvent extends CompletableApplicationEvent<Map<TopicParti
     private final Map<TopicPartition, Long> timestampsToSearch;
     private final boolean requireTimestamps;
 
-    public ListOffsetsEvent(final Map<TopicPartition, Long> timestampToSearch, final boolean requireTimestamps) {
-        super(Type.LIST_OFFSETS);
+    public ListOffsetsEvent(final Map<TopicPartition, Long> timestampToSearch, final boolean requireTimestamps, final Timer timer) {
+        super(Type.LIST_OFFSETS, timer);
         this.timestampsToSearch = Collections.unmodifiableMap(timestampToSearch);
         this.requireTimestamps = requireTimestamps;
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
index 06f6ebbb68a32..65893b62ecaa5 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ResetPositionsEvent.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 /**
  * Event for resetting offsets for all assigned partitions that require it. This is an
  * asynchronous event that generates ListOffsets requests, and completes by updating in-memory
@@ -24,7 +26,7 @@
  */
 public class ResetPositionsEvent extends CompletableApplicationEvent<Void> {
 
-    public ResetPositionsEvent() {
-        super(Type.RESET_POSITIONS);
+    public ResetPositionsEvent(final Timer timer) {
+        super(Type.RESET_POSITIONS, timer);
     }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
index 7e00e0da59683..87945616ea71b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/SyncCommitEvent.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Map;
 
@@ -27,23 +28,7 @@
  */
 public class SyncCommitEvent extends CommitEvent {
 
-    /**
-     * Time to wait for a response, retrying on retriable errors.
-     */
-    private final long retryTimeoutMs;
-
-    public SyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets,
-                           final long retryTimeoutMs) {
-        super(Type.COMMIT_SYNC, offsets);
-        this.retryTimeoutMs = retryTimeoutMs;
-    }
-
-    public Long retryTimeoutMs() {
-        return retryTimeoutMs;
-    }
-
-    @Override
-    public String toStringBase() {
-        return super.toStringBase() + ", offsets=" + offsets() + ", retryTimeoutMs=" + retryTimeoutMs;
+    public SyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets, final Timer timer) {
+        super(Type.COMMIT_SYNC, offsets, timer);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
index ebbb2a6c46892..33e1270ce6040 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/TopicMetadataEvent.java
@@ -16,14 +16,16 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 import java.util.Objects;
 
 public class TopicMetadataEvent extends AbstractTopicMetadataEvent {
 
     private final String topic;
 
-    public TopicMetadataEvent(final String topic, final long timeoutMs) {
-        super(Type.TOPIC_METADATA, timeoutMs);
+    public TopicMetadataEvent(final String topic, final Timer timer) {
+        super(Type.TOPIC_METADATA, timer);
         this.topic = Objects.requireNonNull(topic);
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
index 07af36e5feb85..0b988370014a5 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeEvent.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 /**
  * Application event triggered when a user calls the unsubscribe API. This will make the consumer
  * release all its assignments and send a heartbeat request to leave the consumer group.
@@ -26,8 +28,8 @@
  */
 public class UnsubscribeEvent extends CompletableApplicationEvent<Void> {
 
-    public UnsubscribeEvent() {
-        super(Type.UNSUBSCRIBE);
+    public UnsubscribeEvent(final Timer timer) {
+        super(Type.UNSUBSCRIBE, timer);
     }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
index efa358b4c7882..21e7f3cf6eba1 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ValidatePositionsEvent.java
@@ -17,6 +17,8 @@
 
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 /**
  * Event for validating offsets for all assigned partitions for which a leader change has been
  * detected. This is an asynchronous event that generates OffsetForLeaderEpoch requests, and
@@ -24,7 +26,7 @@
  */
 public class ValidatePositionsEvent extends CompletableApplicationEvent<Void> {
 
-    public ValidatePositionsEvent() {
-        super(Type.VALIDATE_POSITIONS);
+    public ValidatePositionsEvent(final Timer timer) {
+        super(Type.VALIDATE_POSITIONS, timer);
     }
 }
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
index cbd56d8b5e2bd..e4d492fb581a1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkThreadTest.java
@@ -41,6 +41,7 @@
 import org.apache.kafka.common.requests.OffsetCommitResponse;
 import org.apache.kafka.common.requests.RequestTestUtils;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
 import org.apache.kafka.test.TestCondition;
 import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.AfterEach;
@@ -161,7 +162,8 @@ public void testAsyncCommitEvent() {
 
     @Test
     public void testSyncCommitEvent() {
-        ApplicationEvent e = new SyncCommitEvent(new HashMap<>(), 100L);
+        Timer timer = time.timer(100);
+        ApplicationEvent e = new SyncCommitEvent(new HashMap<>(), timer);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor).process(any(SyncCommitEvent.class));
@@ -170,7 +172,8 @@ public void testSyncCommitEvent() {
     @Test
     public void testListOffsetsEventIsProcessed() {
         Map<TopicPartition, Long> timestamps = Collections.singletonMap(new TopicPartition("topic1", 1), 5L);
-        ApplicationEvent e = new ListOffsetsEvent(timestamps, true);
+        Timer timer = time.timer(100);
+        ApplicationEvent e = new ListOffsetsEvent(timestamps, true, timer);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor).process(any(ListOffsetsEvent.class));
@@ -179,7 +182,8 @@ public void testListOffsetsEventIsProcessed() {
 
     @Test
     public void testResetPositionsEventIsProcessed() {
-        ResetPositionsEvent e = new ResetPositionsEvent();
+        Timer timer = time.timer(100);
+        ResetPositionsEvent e = new ResetPositionsEvent(timer);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor).process(any(ResetPositionsEvent.class));
@@ -190,7 +194,8 @@ public void testResetPositionsEventIsProcessed() {
     public void testResetPositionsProcessFailureIsIgnored() {
         doThrow(new NullPointerException()).when(offsetsRequestManager).resetPositionsIfNeeded();
 
-        ResetPositionsEvent event = new ResetPositionsEvent();
+        Timer timer = time.timer(100);
+        ResetPositionsEvent event = new ResetPositionsEvent(timer);
         applicationEventsQueue.add(event);
         assertDoesNotThrow(() -> consumerNetworkThread.runOnce());
 
@@ -199,7 +204,8 @@ public void testResetPositionsProcessFailureIsIgnored() {
 
     @Test
     public void testValidatePositionsEventIsProcessed() {
-        ValidatePositionsEvent e = new ValidatePositionsEvent();
+        Timer timer = time.timer(100);
+        ValidatePositionsEvent e = new ValidatePositionsEvent(timer);
         applicationEventsQueue.add(e);
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor).process(any(ValidatePositionsEvent.class));
@@ -224,7 +230,8 @@ public void testAssignmentChangeEvent() {
 
     @Test
     void testFetchTopicMetadata() {
-        applicationEventsQueue.add(new TopicMetadataEvent("topic", Long.MAX_VALUE));
+        Timer timer = time.timer(Long.MAX_VALUE);
+        applicationEventsQueue.add(new TopicMetadataEvent("topic", timer));
         consumerNetworkThread.runOnce();
         verify(applicationEventProcessor).process(any(TopicMetadataEvent.class));
     }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
index f3e2557ae9417..b23660e5469c8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessorTest.java
@@ -27,6 +27,9 @@
 import org.apache.kafka.clients.consumer.internals.RequestManagers;
 import org.apache.kafka.clients.consumer.internals.TopicMetadataRequestManager;
 import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -36,7 +39,6 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -44,6 +46,7 @@
 import static org.mockito.Mockito.when;
 
 public class ApplicationEventProcessorTest {
+    private final Time time = new MockTime(1);
     private ApplicationEventProcessor processor;
     private BlockingQueue applicationEventQueue = mock(BlockingQueue.class);
     private RequestManagers requestManagers;
@@ -96,18 +99,10 @@ public void testPrepClosingCommitEvents() {
         verify(commitRequestManager).signalClose();
     }
 
-    @Test
-    public void testExpirationCalculation() {
-        assertEquals(Long.MAX_VALUE, processor.getExpirationTimeForTimeout(Long.MAX_VALUE));
-        assertEquals(Long.MAX_VALUE, processor.getExpirationTimeForTimeout(Long.MAX_VALUE - 1));
-        long timeout = processor.getExpirationTimeForTimeout(1000);
-        assertTrue(timeout > 0);
-        assertTrue(timeout < Long.MAX_VALUE);
-    }
-
     @Test
     public void testPrepClosingLeaveGroupEvent() {
-        LeaveOnCloseEvent event = new LeaveOnCloseEvent();
+        Timer timer = time.timer(100);
+        LeaveOnCloseEvent event = new LeaveOnCloseEvent(timer);
         when(heartbeatRequestManager.membershipManager()).thenReturn(membershipManager);
         when(membershipManager.leaveGroup()).thenReturn(CompletableFuture.completedFuture(null));
         processor.process(event);

From 652537f28e5102bfe197528551cd244ce0345319 Mon Sep 17 00:00:00 2001
From: Stanislav Kozlovski <stanislav_kozlovski@outlook.com>
Date: Thu, 7 Mar 2024 15:20:54 +0100
Subject: [PATCH 126/258] MINOR: Add 3.7.0 to core and client's upgrade
 compatibility tests (#15452)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../tests/client/client_compatibility_features_test.py       | 5 +++--
 .../client/client_compatibility_produce_consume_test.py      | 4 +++-
 .../tests/core/compatibility_test_new_broker_test.py         | 3 ++-
 tests/kafkatest/tests/core/downgrade_test.py                 | 5 ++++-
 tests/kafkatest/tests/core/upgrade_test.py                   | 5 ++++-
 5 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/tests/kafkatest/tests/client/client_compatibility_features_test.py b/tests/kafkatest/tests/client/client_compatibility_features_test.py
index b82122b3cdb4b..f29f1df0b4f0d 100644
--- a/tests/kafkatest/tests/client/client_compatibility_features_test.py
+++ b/tests/kafkatest/tests/client/client_compatibility_features_test.py
@@ -28,8 +28,8 @@
 from ducktape.tests.test import Test
 from kafkatest.version import DEV_BRANCH, LATEST_0_10_0, LATEST_0_10_1, LATEST_0_10_2, LATEST_0_11_0, LATEST_1_0, \
     LATEST_1_1, LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, LATEST_2_7, \
-    LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, V_0_11_0_0, \
-    V_0_10_1_0, KafkaVersion
+    LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, LATEST_3_7, \
+    V_0_11_0_0, V_0_10_1_0, KafkaVersion
 
 def get_broker_features(broker_version):
     features = {}
@@ -140,6 +140,7 @@ def invoke_compatibility_program(self, features):
     @parametrize(broker_version=str(LATEST_3_4))
     @parametrize(broker_version=str(LATEST_3_5))
     @parametrize(broker_version=str(LATEST_3_6))
+    @parametrize(broker_version=str(LATEST_3_7))
     def run_compatibility_test(self, broker_version, metadata_quorum=quorum.zk):
         if self.zk:
             self.zk.start()
diff --git a/tests/kafkatest/tests/client/client_compatibility_produce_consume_test.py b/tests/kafkatest/tests/client/client_compatibility_produce_consume_test.py
index 33df4d2fdceaa..afa69eb660569 100644
--- a/tests/kafkatest/tests/client/client_compatibility_produce_consume_test.py
+++ b/tests/kafkatest/tests/client/client_compatibility_produce_consume_test.py
@@ -25,7 +25,8 @@
 from kafkatest.utils import is_int_with_prefix
 from kafkatest.version import DEV_BRANCH, LATEST_0_10_0, LATEST_0_10_1, LATEST_0_10_2, LATEST_0_11_0, LATEST_1_0, \
     LATEST_1_1, LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, LATEST_2_7, \
-    LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, KafkaVersion
+    LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, LATEST_3_7, \
+    KafkaVersion
 
 class ClientCompatibilityProduceConsumeTest(ProduceConsumeValidateTest):
     """
@@ -80,6 +81,7 @@ def min_cluster_size(self):
     @parametrize(broker_version=str(LATEST_3_4))
     @parametrize(broker_version=str(LATEST_3_5))
     @parametrize(broker_version=str(LATEST_3_6))
+    @parametrize(broker_version=str(LATEST_3_7))
     def test_produce_consume(self, broker_version, metadata_quorum=quorum.zk):
         print("running producer_consumer_compat with broker_version = %s" % broker_version, flush=True)
         self.kafka.set_version(KafkaVersion(broker_version))
diff --git a/tests/kafkatest/tests/core/compatibility_test_new_broker_test.py b/tests/kafkatest/tests/core/compatibility_test_new_broker_test.py
index 6d25b0b8e073c..d25f9992c6a40 100644
--- a/tests/kafkatest/tests/core/compatibility_test_new_broker_test.py
+++ b/tests/kafkatest/tests/core/compatibility_test_new_broker_test.py
@@ -24,7 +24,7 @@
 from kafkatest.version import LATEST_0_8_2, LATEST_0_9, LATEST_0_10_0, LATEST_0_10_1, LATEST_0_10_2, LATEST_0_11_0, \
     LATEST_1_0, LATEST_1_1, LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, LATEST_2_6, \
     LATEST_2_7, LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_6, \
-    DEV_BRANCH, KafkaVersion
+    LATEST_3_7, DEV_BRANCH, KafkaVersion
 
 # Compatibility tests for moving to a new broker (e.g., 0.10.x) and using a mix of old and new clients (e.g., 0.9.x)
 class ClientCompatibilityTestNewBroker(ProduceConsumeValidateTest):
@@ -64,6 +64,7 @@ def setUp(self):
     @matrix(producer_version=[str(LATEST_3_4)], consumer_version=[str(LATEST_3_4)], compression_types=[["none"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
     @matrix(producer_version=[str(LATEST_3_5)], consumer_version=[str(LATEST_3_5)], compression_types=[["none"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
     @matrix(producer_version=[str(LATEST_3_6)], consumer_version=[str(LATEST_3_6)], compression_types=[["none"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
+    @matrix(producer_version=[str(LATEST_3_7)], consumer_version=[str(LATEST_3_7)], compression_types=[["none"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
     @matrix(producer_version=[str(LATEST_2_1)], consumer_version=[str(LATEST_2_1)], compression_types=[["zstd"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
     @matrix(producer_version=[str(LATEST_2_0)], consumer_version=[str(LATEST_2_0)], compression_types=[["snappy"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
     @matrix(producer_version=[str(LATEST_1_1)], consumer_version=[str(LATEST_1_1)], compression_types=[["lz4"]], timestamp_type=[str("CreateTime")], metadata_quorum=quorum.all_non_upgrade)
diff --git a/tests/kafkatest/tests/core/downgrade_test.py b/tests/kafkatest/tests/core/downgrade_test.py
index afdb009c5ef3d..a2ada3d868e08 100644
--- a/tests/kafkatest/tests/core/downgrade_test.py
+++ b/tests/kafkatest/tests/core/downgrade_test.py
@@ -21,7 +21,7 @@
 from kafkatest.tests.end_to_end import EndToEndTest
 from kafkatest.version import LATEST_1_1, LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, \
     LATEST_2_6, LATEST_2_7, LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, \
-    LATEST_3_6, DEV_BRANCH, KafkaVersion
+    LATEST_3_6, LATEST_3_7, DEV_BRANCH, KafkaVersion
 
 class TestDowngrade(EndToEndTest):
     PARTITIONS = 3
@@ -81,6 +81,9 @@ def wait_until_rejoin(self):
                     timeout_sec=60, backoff_sec=1, err_msg="Replicas did not rejoin the ISR in a reasonable amount of time")
 
     @cluster(num_nodes=7)
+    @parametrize(version=str(LATEST_3_7), compression_types=["snappy"])
+    @parametrize(version=str(LATEST_3_7), compression_types=["zstd"], security_protocol="SASL_SSL")
+    @matrix(version=[str(LATEST_3_7)], compression_types=[["none"]], static_membership=[False, True])
     @parametrize(version=str(LATEST_3_6), compression_types=["snappy"])
     @parametrize(version=str(LATEST_3_6), compression_types=["zstd"], security_protocol="SASL_SSL")
     @matrix(version=[str(LATEST_3_6)], compression_types=[["none"]], static_membership=[False, True])
diff --git a/tests/kafkatest/tests/core/upgrade_test.py b/tests/kafkatest/tests/core/upgrade_test.py
index 62855b78c87e9..8ba1d0a02de94 100644
--- a/tests/kafkatest/tests/core/upgrade_test.py
+++ b/tests/kafkatest/tests/core/upgrade_test.py
@@ -27,7 +27,7 @@
 from kafkatest.version import LATEST_0_8_2, LATEST_0_9, LATEST_0_10, LATEST_0_10_0, LATEST_0_10_1, LATEST_0_10_2, \
     LATEST_0_11_0, LATEST_1_0, LATEST_1_1, LATEST_2_0, LATEST_2_1, LATEST_2_2, LATEST_2_3, LATEST_2_4, LATEST_2_5, \
     LATEST_2_6, LATEST_2_7, LATEST_2_8, LATEST_3_0, LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, \
-    LATEST_3_6, V_0_11_0_0, V_2_8_0, V_3_0_0, DEV_BRANCH, KafkaVersion
+    LATEST_3_6, LATEST_3_7, V_0_11_0_0, V_2_8_0, V_3_0_0, DEV_BRANCH, KafkaVersion
 from kafkatest.services.kafka.util import new_jdk_not_supported
 
 class TestUpgrade(ProduceConsumeValidateTest):
@@ -94,6 +94,9 @@ def perform_upgrade(self, from_kafka_version, to_message_format_version=None):
             self.wait_until_rejoin()
 
     @cluster(num_nodes=6)
+    @parametrize(from_kafka_version=str(LATEST_3_7), to_message_format_version=None, compression_types=["none"])
+    @parametrize(from_kafka_version=str(LATEST_3_7), to_message_format_version=None, compression_types=["lz4"])
+    @parametrize(from_kafka_version=str(LATEST_3_7), to_message_format_version=None, compression_types=["snappy"])
     @parametrize(from_kafka_version=str(LATEST_3_6), to_message_format_version=None, compression_types=["none"])
     @parametrize(from_kafka_version=str(LATEST_3_6), to_message_format_version=None, compression_types=["lz4"])
     @parametrize(from_kafka_version=str(LATEST_3_6), to_message_format_version=None, compression_types=["snappy"])

From 1d50cbeda85bd3409cff17fd5114e73f25cf865c Mon Sep 17 00:00:00 2001
From: Andrew Schofield <aschofield@confluent.io>
Date: Thu, 7 Mar 2024 15:24:11 +0000
Subject: [PATCH 127/258] KAFKA-16319: Divide DeleteTopics requests by leader
 node (#15479)

Reviewers: Reviewers: Mickael Maison <mickael.maison@gmail.com>, Kirk True <kirk@kirktrue.pro>, Daniel Gospodinow <dgospodinov@confluent.io>
---
 .../admin/internals/DeleteRecordsHandler.java |  6 +-
 .../internals/DeleteRecordsHandlerTest.java   | 73 +++++++++++++++++--
 2 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
index 2daad226034d2..9f40d19f00b48 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
@@ -79,15 +79,15 @@ public static SimpleAdminApiFuture<TopicPartition, DeletedRecords> newFuture(
     @Override
     public DeleteRecordsRequest.Builder buildBatchedRequest(int brokerId, Set<TopicPartition> keys) {
         Map<String, DeleteRecordsRequestData.DeleteRecordsTopic> deletionsForTopic = new HashMap<>();
-        for (Map.Entry<TopicPartition, RecordsToDelete> entry: recordsToDelete.entrySet()) {
-            TopicPartition topicPartition = entry.getKey();
+        for (TopicPartition topicPartition : keys) {
+            RecordsToDelete toDelete = recordsToDelete.get(topicPartition);
             DeleteRecordsRequestData.DeleteRecordsTopic deleteRecords = deletionsForTopic.computeIfAbsent(
                     topicPartition.topic(),
                     key -> new DeleteRecordsRequestData.DeleteRecordsTopic().setName(topicPartition.topic())
             );
             deleteRecords.partitions().add(new DeleteRecordsRequestData.DeleteRecordsPartition()
                     .setPartitionIndex(topicPartition.partition())
-                    .setOffset(entry.getValue().beforeOffset()));
+                    .setOffset(toDelete.beforeOffset()));
         }
 
         DeleteRecordsRequestData data = new DeleteRecordsRequestData()
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandlerTest.java
index c39747f1fbafe..58492696c4da8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandlerTest.java
@@ -22,15 +22,24 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
+
 import org.apache.kafka.clients.admin.DeletedRecords;
 import org.apache.kafka.clients.admin.RecordsToDelete;
+import org.apache.kafka.clients.admin.internals.AdminApiLookupStrategy.LookupResult;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.message.DeleteRecordsRequestData;
 import org.apache.kafka.common.message.DeleteRecordsResponseData;
+import org.apache.kafka.common.message.MetadataResponseData;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponsePartition;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.DeleteRecordsRequest;
 import org.apache.kafka.common.requests.DeleteRecordsResponse;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.utils.LogContext;
 import org.junit.jupiter.api.Test;
 
@@ -41,6 +50,7 @@
 
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class DeleteRecordsHandlerTest {
@@ -50,7 +60,8 @@ public class DeleteRecordsHandlerTest {
     private final TopicPartition t0p1 = new TopicPartition("t0", 1);
     private final TopicPartition t0p2 = new TopicPartition("t0", 2);
     private final TopicPartition t0p3 = new TopicPartition("t0", 3);
-    private final Node node = new Node(1, "host", 1234);
+    private final Node node1 = new Node(1, "host", 1234);
+    private final Node node2 = new Node(2, "host", 1235);
     private final Map<TopicPartition, RecordsToDelete> recordsToDelete = new HashMap<TopicPartition, RecordsToDelete>() {
         {
             put(t0p0, RecordsToDelete.beforeOffset(10L));
@@ -63,11 +74,11 @@ public class DeleteRecordsHandlerTest {
     @Test
     public void testBuildRequestSimple() {
         DeleteRecordsHandler handler = new DeleteRecordsHandler(recordsToDelete, logContext, timeout);
-        DeleteRecordsRequest request = handler.buildBatchedRequest(node.id(), mkSet(t0p0, t0p1)).build();
-        List<DeleteRecordsRequestData.DeleteRecordsTopic> topicPartitions = request.data().topics();
-        assertEquals(1, topicPartitions.size());
-        DeleteRecordsRequestData.DeleteRecordsTopic topic = topicPartitions.get(0);
-        assertEquals(4, topic.partitions().size());
+        DeleteRecordsRequest request = handler.buildBatchedRequest(node1.id(), mkSet(t0p0, t0p1)).build();
+        List<DeleteRecordsRequestData.DeleteRecordsTopic> topics = request.data().topics();
+        assertEquals(1, topics.size());
+        DeleteRecordsRequestData.DeleteRecordsTopic topic = topics.get(0);
+        assertEquals(2, topic.partitions().size());
     }
 
     @Test
@@ -199,6 +210,54 @@ public void testHandleResponseSanityCheck() {
         assertTrue(result.unmappedKeys.isEmpty());
     }
 
+    // This is a more complicated test which ensures that DeleteRecords requests for multiple
+    // leader nodes are correctly divided up among the nodes based on leadership.
+    // node1 leads t0p0 and t0p2, while node2 leads t0p1 and t0p3.
+    @Test
+    public void testBuildRequestMultipleLeaders() {
+        MetadataResponseData metadataResponseData = new MetadataResponseData();
+        MetadataResponseTopic topicMetadata = new MetadataResponseTopic();
+        topicMetadata.setName("t0").setErrorCode(Errors.NONE.code());
+        topicMetadata.partitions().add(new MetadataResponsePartition()
+                .setPartitionIndex(0).setLeaderId(node1.id()).setErrorCode(Errors.NONE.code()));
+        topicMetadata.partitions().add(new MetadataResponsePartition()
+                .setPartitionIndex(1).setLeaderId(node2.id()).setErrorCode(Errors.NONE.code()));
+        topicMetadata.partitions().add(new MetadataResponsePartition()
+                .setPartitionIndex(2).setLeaderId(node1.id()).setErrorCode(Errors.NONE.code()));
+        topicMetadata.partitions().add(new MetadataResponsePartition()
+                .setPartitionIndex(3).setLeaderId(node2.id()).setErrorCode(Errors.NONE.code()));
+        metadataResponseData.topics().add(topicMetadata);
+        MetadataResponse metadataResponse = new MetadataResponse(metadataResponseData, ApiKeys.METADATA.latestVersion());
+
+        DeleteRecordsHandler handler = new DeleteRecordsHandler(recordsToDelete, logContext, timeout);
+        AdminApiLookupStrategy<TopicPartition> strategy = handler.lookupStrategy();
+        assertInstanceOf(PartitionLeaderStrategy.class, strategy);
+        PartitionLeaderStrategy specificStrategy = (PartitionLeaderStrategy) strategy;
+        MetadataRequest request = specificStrategy.buildRequest(mkSet(t0p0, t0p1, t0p2, t0p3)).build();
+        assertEquals(mkSet("t0"), new HashSet<>(request.topics()));
+
+        Set<TopicPartition> tpSet = mkSet(t0p0, t0p1, t0p2, t0p3);
+        LookupResult<TopicPartition> lookupResult = strategy.handleResponse(tpSet, metadataResponse);
+        assertEquals(emptyMap(), lookupResult.failedKeys);
+        assertEquals(tpSet, lookupResult.mappedKeys.keySet());
+
+        Map<Integer, Set<TopicPartition>> partitionsPerBroker = new HashMap<>();
+        lookupResult.mappedKeys.forEach((tp, node) -> partitionsPerBroker.computeIfAbsent(node, key -> new HashSet<>()).add(tp));
+
+        DeleteRecordsRequest deleteRequest = handler.buildBatchedRequest(node1.id(), partitionsPerBroker.get(node1.id())).build();
+        assertEquals(2, deleteRequest.data().topics().get(0).partitions().size());
+        assertEquals(mkSet(t0p0, t0p2),
+                deleteRequest.data().topics().get(0).partitions().stream()
+                        .map(drp -> new TopicPartition("t0", drp.partitionIndex()))
+                        .collect(Collectors.toSet()));
+        deleteRequest = handler.buildBatchedRequest(node2.id(), partitionsPerBroker.get(node2.id())).build();
+        assertEquals(2, deleteRequest.data().topics().get(0).partitions().size());
+        assertEquals(mkSet(t0p1, t0p3),
+                deleteRequest.data().topics().get(0).partitions().stream()
+                        .map(drp -> new TopicPartition("t0", drp.partitionIndex()))
+                        .collect(Collectors.toSet()));
+    }
+
     private DeleteRecordsResponse createResponse(Map<TopicPartition, Short> errorsByPartition) {
         return createResponse(errorsByPartition, recordsToDelete.keySet());
     }
@@ -227,7 +286,7 @@ private DeleteRecordsResponse createResponse(
     private AdminApiHandler.ApiResult<TopicPartition, DeletedRecords> handleResponse(DeleteRecordsResponse response) {
         DeleteRecordsHandler handler =
                 new DeleteRecordsHandler(recordsToDelete, logContext, timeout);
-        return handler.handleResponse(node, recordsToDelete.keySet(), response);
+        return handler.handleResponse(node1, recordsToDelete.keySet(), response);
     }
 
     private void assertResult(

From f5c4d522fd79775692441cd75f6f733324f2e7e9 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Thu, 7 Mar 2024 07:51:04 -0800
Subject: [PATCH 128/258] MINOR: Add read/write all operation (#15462)

There are a few cases in the group coordinator service where we want to read from or write to each of the known coordinators (each of __consumer_offsets partitions). The current implementation needs to get the list of the known coordinators then schedules the operation and finally aggregate the results. This patch is an attempt to streamline this by adding multi read/write to the runtime.

Reviewers: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../group/GroupCoordinatorService.java        |  60 +++----
 .../group/runtime/CoordinatorRuntime.java     |  67 ++++++--
 .../group/GroupCoordinatorServiceTest.java    | 148 ++++++------------
 .../group/runtime/CoordinatorRuntimeTest.java | 110 +++++++++++++
 .../apache/kafka/server/util/FutureUtils.java |  19 +++
 5 files changed, 254 insertions(+), 150 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
index a2363b4822274..2fdc128c7b98c 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
@@ -80,7 +80,6 @@
 import java.util.Map;
 import java.util.OptionalInt;
 import java.util.Properties;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
@@ -498,29 +497,24 @@ public CompletableFuture<ListGroupsResponseData> listGroups(
             );
         }
 
-        final Set<TopicPartition> existingPartitionSet = runtime.partitions();
-
-        if (existingPartitionSet.isEmpty()) {
-            return CompletableFuture.completedFuture(new ListGroupsResponseData());
-        }
-
-        final List<CompletableFuture<List<ListGroupsResponseData.ListedGroup>>> futures =
-            new ArrayList<>();
-
-        for (TopicPartition tp : existingPartitionSet) {
-            futures.add(runtime.scheduleReadOperation(
+        final List<CompletableFuture<List<ListGroupsResponseData.ListedGroup>>> futures = FutureUtils.mapExceptionally(
+            runtime.scheduleReadAllOperation(
                 "list-groups",
-                tp,
-                (coordinator, lastCommittedOffset) -> coordinator.listGroups(request.statesFilter(), request.typesFilter(), lastCommittedOffset)
-            ).exceptionally(exception -> {
+                (coordinator, lastCommittedOffset) -> coordinator.listGroups(
+                    request.statesFilter(),
+                    request.typesFilter(),
+                    lastCommittedOffset
+                )
+            ),
+            exception -> {
                 exception = Errors.maybeUnwrapException(exception);
                 if (exception instanceof NotCoordinatorException) {
                     return Collections.emptyList();
                 } else {
                     throw new CompletionException(exception);
                 }
-            }));
-        }
+            }
+        );
 
         return FutureUtils
             .combineFutures(futures, ArrayList::new, List::addAll)
@@ -963,23 +957,21 @@ public void onPartitionsDeleted(
     ) throws ExecutionException, InterruptedException {
         throwIfNotActive();
 
-        final Set<TopicPartition> existingPartitionSet = runtime.partitions();
-        final List<CompletableFuture<Void>> futures = new ArrayList<>(existingPartitionSet.size());
-
-        existingPartitionSet.forEach(partition -> futures.add(
-            runtime.scheduleWriteOperation(
-                "on-partition-deleted",
-                partition,
-                Duration.ofMillis(config.offsetCommitTimeoutMs),
-                coordinator -> coordinator.onPartitionsDeleted(topicPartitions)
-            ).exceptionally(exception -> {
-                log.error("Could not delete offsets for deleted partitions {} in coordinator {} due to: {}.",
-                    partition, partition, exception.getMessage(), exception);
-                return null;
-            })
-        ));
-
-        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get();
+        CompletableFuture.allOf(
+            FutureUtils.mapExceptionally(
+                runtime.scheduleWriteAllOperation(
+                    "on-partition-deleted",
+                    Duration.ofMillis(config.offsetCommitTimeoutMs),
+                    coordinator -> coordinator.onPartitionsDeleted(topicPartitions)
+                ),
+                exception -> {
+                    log.error("Could not delete offsets for deleted partitions {} due to: {}.",
+                        topicPartitions, exception.getMessage(), exception
+                    );
+                    return null;
+                }
+            ).toArray(new CompletableFuture[0])
+        ).get();
     }
 
     /**
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
index ccb4caf04f98e..6b98a51dd47b9 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
@@ -41,12 +41,11 @@
 
 import java.time.Duration;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.RejectedExecutionException;
@@ -54,6 +53,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * The CoordinatorRuntime provides a framework to implement coordinators such as the group coordinator
@@ -1446,6 +1446,32 @@ public <T> CompletableFuture<T> scheduleWriteOperation(
         return event.future;
     }
 
+    /**
+     * Schedule a write operation for each coordinator.
+     *
+     * @param name      The name of the write operation.
+     * @param timeout   The write operation timeout.
+     * @param op        The write operation.
+     *
+     * @return A list of futures where each future will be completed with the result of the write operation
+     * when the operation is completed or an exception if the write operation failed.
+     *
+     * @param <T> The type of the result.
+     */
+    public <T> List<CompletableFuture<T>> scheduleWriteAllOperation(
+        String name,
+        Duration timeout,
+        CoordinatorWriteOperation<S, T, U> op
+    ) {
+        throwIfNotRunning();
+        log.debug("Scheduled execution of write all operation {}.", name);
+        return coordinators
+            .keySet()
+            .stream()
+            .map(tp -> scheduleWriteOperation(name, tp, timeout, op))
+            .collect(Collectors.toList());
+    }
+
     /**
      * Schedules a transactional write operation.
      *
@@ -1535,12 +1561,12 @@ public CompletableFuture<Void> scheduleTransactionCompletion(
     /**
      * Schedules a read operation.
      *
-     * @param name  The name of the write operation.
+     * @param name  The name of the read operation.
      * @param tp    The address of the coordinator (aka its topic-partitions).
      * @param op    The read operation.
      *
      * @return A future that will be completed with the result of the read operation
-     * when the operation is completed or an exception if the write operation failed.
+     * when the operation is completed or an exception if the read operation failed.
      *
      * @param <T> The type of the result.
      */
@@ -1556,6 +1582,30 @@ public <T> CompletableFuture<T> scheduleReadOperation(
         return event.future;
     }
 
+    /**
+     * Schedules a read operation for each coordinator.
+     *
+     * @param name  The name of the read operation.
+     * @param op    The read operation.
+     *
+     * @return A list of futures where each future will be completed with the result of the read operation
+     * when the operation is completed or an exception if the read operation failed.
+     *
+     * @param <T> The type of the result.
+     */
+    public <T> List<CompletableFuture<T>> scheduleReadAllOperation(
+        String name,
+        CoordinatorReadOperation<S, T> op
+    ) {
+        throwIfNotRunning();
+        log.debug("Scheduled execution of read all operation {}.", name);
+        return coordinators
+            .keySet()
+            .stream()
+            .map(tp -> scheduleReadOperation(name, tp, op))
+            .collect(Collectors.toList());
+    }
+
     /**
      * Schedules an internal event.
      *
@@ -1572,15 +1622,6 @@ private void scheduleInternalOperation(
         enqueue(new CoordinatorInternalEvent(name, tp, op));
     }
 
-    /**
-     * @return The topic partitions of the coordinators currently registered in the
-     * runtime.
-     */
-    public Set<TopicPartition> partitions() {
-        throwIfNotRunning();
-        return new HashSet<>(coordinators.keySet());
-    }
-
     /**
      * Schedules the loading of a coordinator. This is called when the broker is elected as
      * the leader for a partition.
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
index 2404d304becf7..0d93f450686f5 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
@@ -77,7 +77,6 @@
 import org.junit.jupiter.params.provider.NullSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentMatchers;
-import org.mockito.internal.util.collections.Sets;
 
 import java.net.InetAddress;
 import java.time.Duration;
@@ -91,8 +90,6 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID;
@@ -759,10 +756,7 @@ public void testListGroups() throws ExecutionException, InterruptedException, Ti
             runtime,
             new GroupCoordinatorMetrics()
         );
-        int partitionCount = 3;
-        service.startup(() -> partitionCount);
-
-        ListGroupsRequestData request = new ListGroupsRequestData();
+        service.startup(() -> 3);
 
         List<ListGroupsResponseData.ListedGroup> expectedResults = Arrays.asList(
             new ListGroupsResponseData.ListedGroup()
@@ -781,26 +775,22 @@ public void testListGroups() throws ExecutionException, InterruptedException, Ti
                 .setGroupState("Dead")
                 .setGroupType("consumer")
         );
-        when(runtime.partitions()).thenReturn(Sets.newSet(
-            new TopicPartition("__consumer_offsets", 0),
-            new TopicPartition("__consumer_offsets", 1),
-            new TopicPartition("__consumer_offsets", 2)
+
+        when(runtime.scheduleReadAllOperation(
+            ArgumentMatchers.eq("list-groups"),
+            ArgumentMatchers.any()
+        )).thenReturn(Arrays.asList(
+            CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(0))),
+            CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(1))),
+            CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(2)))
         ));
-        for (int i = 0; i < partitionCount; i++) {
-            when(runtime.scheduleReadOperation(
-                ArgumentMatchers.eq("list-groups"),
-                ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", i)),
-                ArgumentMatchers.any()
-            )).thenReturn(CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(i))));
-        }
 
         CompletableFuture<ListGroupsResponseData> responseFuture = service.listGroups(
             requestContext(ApiKeys.LIST_GROUPS),
-            request
+            new ListGroupsRequestData()
         );
 
-        List<ListGroupsResponseData.ListedGroup> actualResults = responseFuture.get(5, TimeUnit.SECONDS).groups();
-        assertEquals(expectedResults, actualResults);
+        assertEquals(expectedResults, responseFuture.get(5, TimeUnit.SECONDS).groups());
     }
 
     @Test
@@ -813,8 +803,7 @@ public void testListGroupsFailedWithNotCoordinatorException()
             runtime,
             new GroupCoordinatorMetrics()
         );
-        int partitionCount = 3;
-        service.startup(() -> partitionCount);
+        service.startup(() -> 3);
 
         List<ListGroupsResponseData.ListedGroup> expectedResults = Arrays.asList(
             new ListGroupsResponseData.ListedGroup()
@@ -829,36 +818,25 @@ public void testListGroupsFailedWithNotCoordinatorException()
                 .setGroupType("consumer")
         );
 
-        ListGroupsRequestData request = new ListGroupsRequestData();
-        when(runtime.partitions()).thenReturn(Sets.newSet(
-            new TopicPartition("__consumer_offsets", 0),
-            new TopicPartition("__consumer_offsets", 1),
-            new TopicPartition("__consumer_offsets", 2)
-        ));
-        for (int i = 0; i < 2; i++) {
-            when(runtime.scheduleReadOperation(
-                ArgumentMatchers.eq("list-groups"),
-                ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", i)),
-                ArgumentMatchers.any()
-            )).thenReturn(CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(i))));
-        }
-
-        when(runtime.scheduleReadOperation(
+        when(runtime.scheduleReadAllOperation(
             ArgumentMatchers.eq("list-groups"),
-            ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", 2)),
             ArgumentMatchers.any()
-        )).thenReturn(FutureUtils.failedFuture(new NotCoordinatorException("")));
+        )).thenReturn(Arrays.asList(
+            CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(0))),
+            CompletableFuture.completedFuture(Collections.singletonList(expectedResults.get(1))),
+            FutureUtils.failedFuture(new NotCoordinatorException(""))
+        ));
 
         CompletableFuture<ListGroupsResponseData> responseFuture = service.listGroups(
             requestContext(ApiKeys.LIST_GROUPS),
-            request
+            new ListGroupsRequestData()
         );
-        List<ListGroupsResponseData.ListedGroup> actualResults = responseFuture.get(5, TimeUnit.SECONDS).groups();
-        assertEquals(expectedResults, actualResults);
+
+        assertEquals(expectedResults, responseFuture.get(5, TimeUnit.SECONDS).groups());
     }
 
     @Test
-    public void testListGroupsFailedImmediately()
+    public void testListGroupsWithFailure()
         throws InterruptedException, ExecutionException, TimeoutException {
         CoordinatorRuntime<GroupCoordinatorShard, Record> runtime = mockRuntime();
         GroupCoordinatorService service = new GroupCoordinatorService(
@@ -867,37 +845,27 @@ public void testListGroupsFailedImmediately()
             runtime,
             new GroupCoordinatorMetrics()
         );
-        int partitionCount = 3;
-        service.startup(() -> partitionCount);
-
-        ListGroupsRequestData request = new ListGroupsRequestData();
-        when(runtime.partitions()).thenReturn(Sets.newSet(
-            new TopicPartition("__consumer_offsets", 0),
-            new TopicPartition("__consumer_offsets", 1),
-            new TopicPartition("__consumer_offsets", 2)
-        ));
-        for (int i = 0; i < 2; i++) {
-            when(runtime.scheduleReadOperation(
-                ArgumentMatchers.eq("list-groups"),
-                ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", i)),
-                ArgumentMatchers.any()
-            )).thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
-        }
+        service.startup(() -> 3);
 
-        when(runtime.scheduleReadOperation(
+        when(runtime.scheduleReadAllOperation(
             ArgumentMatchers.eq("list-groups"),
-            ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", 2)),
             ArgumentMatchers.any()
-        )).thenReturn(FutureUtils.failedFuture(new CoordinatorLoadInProgressException("")));
+        )).thenReturn(Arrays.asList(
+            CompletableFuture.completedFuture(Collections.emptyList()),
+            CompletableFuture.completedFuture(Collections.emptyList()),
+            FutureUtils.failedFuture(new CoordinatorLoadInProgressException(""))
+        ));
 
         CompletableFuture<ListGroupsResponseData> responseFuture = service.listGroups(
             requestContext(ApiKeys.LIST_GROUPS),
-            request
+            new ListGroupsRequestData()
         );
-        ListGroupsResponseData listGroupsResponseData = responseFuture.get(5, TimeUnit.SECONDS);
 
-        assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS.code(), listGroupsResponseData.errorCode());
-        assertEquals(Collections.emptyList(), listGroupsResponseData.groups());
+        assertEquals(
+            new ListGroupsResponseData()
+                .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code()),
+            responseFuture.get(5, TimeUnit.SECONDS)
+        );
     }
 
     @Test
@@ -1703,12 +1671,6 @@ public void testDeleteGroups() throws Exception {
             result1.duplicate()
         ));
 
-        when(runtime.partitions()).thenReturn(Sets.newSet(
-            new TopicPartition("__consumer_offsets", 0),
-            new TopicPartition("__consumer_offsets", 1),
-            new TopicPartition("__consumer_offsets", 2)
-        ));
-
         when(runtime.scheduleWriteOperation(
             ArgumentMatchers.eq("delete-groups"),
             ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", 2)),
@@ -2067,7 +2029,6 @@ public void testCompleteTransactionWithUnexpectedPartition() {
 
     @Test
     public void testOnPartitionsDeleted() {
-        int partitionCount = 3;
         CoordinatorRuntime<GroupCoordinatorShard, Record> runtime = mockRuntime();
         GroupCoordinatorService service = new GroupCoordinatorService(
             new LogContext(),
@@ -2075,36 +2036,17 @@ public void testOnPartitionsDeleted() {
             runtime,
             new GroupCoordinatorMetrics()
         );
+        service.startup(() -> 3);
 
-        service.startup(() -> partitionCount);
-
-        when(runtime.partitions()).thenReturn(
-            IntStream
-                .range(0, partitionCount)
-                .mapToObj(i -> new TopicPartition("__consumer_offsets", i))
-                .collect(Collectors.toSet())
-        );
-
-        List<CompletableFuture<Void>> futures = IntStream
-            .range(0, partitionCount)
-            .mapToObj(__ -> new CompletableFuture<Void>())
-            .collect(Collectors.toList());
-
-        IntStream.range(0, partitionCount).forEach(i -> {
-            CompletableFuture<Void> future = futures.get(i);
-            when(runtime.scheduleWriteOperation(
-                ArgumentMatchers.eq("on-partition-deleted"),
-                ArgumentMatchers.eq(new TopicPartition("__consumer_offsets", i)),
-                ArgumentMatchers.eq(Duration.ofMillis(5000)),
-                ArgumentMatchers.any()
-            )).thenAnswer(__ -> future);
-        });
-
-        IntStream.range(0, partitionCount - 1).forEach(i -> {
-            futures.get(i).complete(null);
-        });
-
-        futures.get(partitionCount - 1).completeExceptionally(Errors.COORDINATOR_LOAD_IN_PROGRESS.exception());
+        when(runtime.scheduleWriteAllOperation(
+            ArgumentMatchers.eq("on-partition-deleted"),
+            ArgumentMatchers.eq(Duration.ofMillis(5000)),
+            ArgumentMatchers.any()
+        )).thenReturn(Arrays.asList(
+            CompletableFuture.completedFuture(null),
+            CompletableFuture.completedFuture(null),
+            FutureUtils.failedFuture(Errors.COORDINATOR_LOAD_IN_PROGRESS.exception())
+        ));
 
         // The exception is logged and swallowed.
         assertDoesNotThrow(() ->
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
index ba1a340e8e1f1..4e10978d35740 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
@@ -43,6 +43,7 @@
 import org.mockito.ArgumentMatcher;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
@@ -1145,6 +1146,63 @@ public void testScheduleWriteOpWhenWriteTimesOut() throws InterruptedException {
         assertFutureThrows(timedOutWrite, org.apache.kafka.common.errors.TimeoutException.class);
     }
 
+    @Test
+    public void testScheduleWriteAllOperation() throws ExecutionException, InterruptedException, TimeoutException {
+        MockTimer timer = new MockTimer();
+        MockPartitionWriter writer = new MockPartitionWriter();
+
+        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
+            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
+                .withTime(timer.time())
+                .withTimer(timer)
+                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
+                .withLoader(new MockCoordinatorLoader())
+                .withEventProcessor(new DirectEventProcessor())
+                .withPartitionWriter(writer)
+                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
+                .withCoordinatorRuntimeMetrics(mock(GroupCoordinatorRuntimeMetrics.class))
+                .withCoordinatorMetrics(mock(GroupCoordinatorMetrics.class))
+                .build();
+
+        TopicPartition coordinator0 = new TopicPartition("__consumer_offsets", 0);
+        TopicPartition coordinator1 = new TopicPartition("__consumer_offsets", 1);
+        TopicPartition coordinator2 = new TopicPartition("__consumer_offsets", 2);
+
+        // Load coordinators.
+        runtime.scheduleLoadOperation(coordinator0, 10);
+        runtime.scheduleLoadOperation(coordinator1, 10);
+        runtime.scheduleLoadOperation(coordinator2, 10);
+
+        // Writes.
+        AtomicInteger cnt = new AtomicInteger(0);
+        List<CompletableFuture<List<String>>> writes = runtime.scheduleWriteAllOperation("write", DEFAULT_WRITE_TIMEOUT, state -> {
+            int counter = cnt.getAndIncrement();
+            return new CoordinatorResult<>(
+                Collections.singletonList("record#" + counter),
+                Collections.singletonList("response#" + counter)
+            );
+        });
+
+        assertEquals(1L, runtime.contextOrThrow(coordinator0).coordinator.lastWrittenOffset());
+        assertEquals(1L, runtime.contextOrThrow(coordinator1).coordinator.lastWrittenOffset());
+        assertEquals(1L, runtime.contextOrThrow(coordinator2).coordinator.lastWrittenOffset());
+
+        assertEquals(Collections.singletonList(InMemoryPartitionWriter.LogEntry.value("record#0")), writer.entries(coordinator0));
+        assertEquals(Collections.singletonList(InMemoryPartitionWriter.LogEntry.value("record#1")), writer.entries(coordinator1));
+        assertEquals(Collections.singletonList(InMemoryPartitionWriter.LogEntry.value("record#2")), writer.entries(coordinator2));
+
+        // Commit.
+        writer.commit(coordinator0);
+        writer.commit(coordinator1);
+        writer.commit(coordinator2);
+
+        // Verify.
+        assertEquals(
+            Arrays.asList("response#0", "response#1", "response#2"),
+            FutureUtils.combineFutures(writes, ArrayList::new, List::addAll).get(5, TimeUnit.SECONDS)
+        );
+    }
+
     @Test
     public void testScheduleTransactionalWriteOp() {
         MockTimer timer = new MockTimer();
@@ -1743,6 +1801,58 @@ public void testScheduleReadOpWhenOpsFails() {
         assertFutureThrows(read, IllegalArgumentException.class);
     }
 
+    @Test
+    public void testScheduleReadAllOp() throws ExecutionException, InterruptedException, TimeoutException {
+        MockTimer timer = new MockTimer();
+        MockPartitionWriter writer = new MockPartitionWriter();
+
+        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
+            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
+                .withTime(timer.time())
+                .withTimer(timer)
+                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
+                .withLoader(new MockCoordinatorLoader())
+                .withEventProcessor(new DirectEventProcessor())
+                .withPartitionWriter(writer)
+                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
+                .withCoordinatorRuntimeMetrics(mock(GroupCoordinatorRuntimeMetrics.class))
+                .withCoordinatorMetrics(mock(GroupCoordinatorMetrics.class))
+                .build();
+
+        TopicPartition coordinator0 = new TopicPartition("__consumer_offsets", 0);
+        TopicPartition coordinator1 = new TopicPartition("__consumer_offsets", 1);
+        TopicPartition coordinator2 = new TopicPartition("__consumer_offsets", 2);
+
+        // Loads the coordinators.
+        runtime.scheduleLoadOperation(coordinator0, 10);
+        runtime.scheduleLoadOperation(coordinator1, 10);
+        runtime.scheduleLoadOperation(coordinator2, 10);
+
+        // Writes
+        runtime.scheduleWriteOperation("write#0", coordinator0, DEFAULT_WRITE_TIMEOUT,
+            state -> new CoordinatorResult<>(Collections.singletonList("record0"), "response0"));
+        runtime.scheduleWriteOperation("write#1", coordinator1, DEFAULT_WRITE_TIMEOUT,
+            state -> new CoordinatorResult<>(Collections.singletonList("record1"), "response1"));
+        runtime.scheduleWriteOperation("write#2", coordinator2, DEFAULT_WRITE_TIMEOUT,
+            state -> new CoordinatorResult<>(Collections.singletonList("record2"), "response2"));
+
+        // Commit writes.
+        writer.commit(coordinator0);
+        writer.commit(coordinator1);
+        writer.commit(coordinator2);
+
+        // Read.
+        List<CompletableFuture<List<String>>> responses = runtime.scheduleReadAllOperation(
+            "read",
+            (state, offset) -> new ArrayList<>(state.records)
+        );
+
+        assertEquals(
+            Arrays.asList("record0", "record1", "record2"),
+            FutureUtils.combineFutures(responses, ArrayList::new, List::addAll).get(5, TimeUnit.SECONDS)
+        );
+    }
+
     @Test
     public void testClose() throws Exception {
         MockCoordinatorLoader loader = spy(new MockCoordinatorLoader());
diff --git a/server-common/src/main/java/org/apache/kafka/server/util/FutureUtils.java b/server-common/src/main/java/org/apache/kafka/server/util/FutureUtils.java
index a3f78a84b1769..a95716407b465 100644
--- a/server-common/src/main/java/org/apache/kafka/server/util/FutureUtils.java
+++ b/server-common/src/main/java/org/apache/kafka/server/util/FutureUtils.java
@@ -19,11 +19,13 @@
 import org.apache.kafka.common.utils.Time;
 import org.slf4j.Logger;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 
@@ -125,4 +127,21 @@ public static <T> CompletableFuture<T> combineFutures(
             return res;
         });
     }
+
+    /**
+     * Applies the given exception handler to all the futures provided in the list
+     * and returns a new list of futures.
+     *
+     * @param futures   A list of futures.
+     * @param fn        A function taking an exception to handle it.
+     * @return A list of futures.
+     */
+    public static <T> List<CompletableFuture<T>> mapExceptionally(
+        List<CompletableFuture<T>> futures,
+        Function<Throwable, ? extends T> fn
+    ) {
+        final List<CompletableFuture<T>> results = new ArrayList<>(futures.size());
+        futures.forEach(future -> results.add(future.exceptionally(fn)));
+        return results;
+    }
 }

From d8dd068a626dcab538c2b234ffd8799a94b2f0ed Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Thu, 7 Mar 2024 09:17:52 -0800
Subject: [PATCH 129/258] KAFKA-15964: fix flaky StreamsAssignmentScaleTest
 (#15485)

This PR bumps some timeouts due to slow Jenkins builds.

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 .../internals/StreamsPartitionAssignor.java   |  2 +-
 .../internals/StreamsAssignmentScaleTest.java | 26 +++++++++----------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
index 3f16838dfc2d8..fb4e45cdfb1a9 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
@@ -628,7 +628,7 @@ private boolean assignTasksToClients(final Cluster fullMetadata,
 
         log.info("{} client nodes and {} consumers participating in this rebalance: \n{}.",
                  clientStates.size(),
-                 clientStates.values().stream().map(ClientState::capacity).reduce(Integer::sum),
+                 clientStates.values().stream().map(ClientState::capacity).reduce(Integer::sum).orElse(0),
                  clientStates.entrySet().stream()
                      .sorted(comparingByKey())
                      .map(entry -> entry.getKey() + ": " + entry.getValue().consumers())
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsAssignmentScaleTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsAssignmentScaleTest.java
index 834b6242194ed..e14e033363436 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsAssignmentScaleTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsAssignmentScaleTest.java
@@ -68,73 +68,73 @@
 @Category({IntegrationTest.class})
 @RunWith(MockitoJUnitRunner.StrictStubs.class)
 public class StreamsAssignmentScaleTest {
-    final static long MAX_ASSIGNMENT_DURATION = 60 * 1000L; //each individual assignment should complete within 20s
+    final static long MAX_ASSIGNMENT_DURATION = 120 * 1000L; // we should stay below `max.poll.interval.ms`
     final static String APPLICATION_ID = "streams-assignment-scale-test";
 
     private final Logger log = LoggerFactory.getLogger(StreamsAssignmentScaleTest.class);
 
     /* HighAvailabilityTaskAssignor tests */
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testHighAvailabilityTaskAssignorLargePartitionCount() {
         completeLargeAssignment(6_000, 2, 1, 1, HighAvailabilityTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testHighAvailabilityTaskAssignorLargeNumConsumers() {
         completeLargeAssignment(1_000, 1_000, 1, 1, HighAvailabilityTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testHighAvailabilityTaskAssignorManyStandbys() {
         completeLargeAssignment(1_000, 100, 1, 50, HighAvailabilityTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testHighAvailabilityTaskAssignorManyThreadsPerClient() {
         completeLargeAssignment(1_000, 10, 1000, 1, HighAvailabilityTaskAssignor.class);
     }
 
     /* StickyTaskAssignor tests */
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testStickyTaskAssignorLargePartitionCount() {
         completeLargeAssignment(2_000, 2, 1, 1, StickyTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testStickyTaskAssignorLargeNumConsumers() {
         completeLargeAssignment(1_000, 1_000, 1, 1, StickyTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testStickyTaskAssignorManyStandbys() {
         completeLargeAssignment(1_000, 100, 1, 20, StickyTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testStickyTaskAssignorManyThreadsPerClient() {
         completeLargeAssignment(1_000, 10, 1000, 1, StickyTaskAssignor.class);
     }
 
     /* FallbackPriorTaskAssignor tests */
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testFallbackPriorTaskAssignorLargePartitionCount() {
         completeLargeAssignment(2_000, 2, 1, 1, FallbackPriorTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testFallbackPriorTaskAssignorLargeNumConsumers() {
         completeLargeAssignment(1_000, 1_000, 1, 1, FallbackPriorTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testFallbackPriorTaskAssignorManyStandbys() {
         completeLargeAssignment(1_000, 100, 1, 20, FallbackPriorTaskAssignor.class);
     }
 
-    @Test(timeout = 120 * 1000)
+    @Test(timeout = 300 * 1000)
     public void testFallbackPriorTaskAssignorManyThreadsPerClient() {
         completeLargeAssignment(1_000, 10, 1000, 1, FallbackPriorTaskAssignor.class);
     }

From 5dd382ccbd1fada55cc44a5223dcae56e0e8811d Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Fri, 8 Mar 2024 03:02:22 +0800
Subject: [PATCH 130/258] MINOR: Use INFO logging for tools tests (#15487)

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 tools/src/test/resources/log4j.properties | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tools/src/test/resources/log4j.properties b/tools/src/test/resources/log4j.properties
index 5291604d49ae5..3aca07dc53016 100644
--- a/tools/src/test/resources/log4j.properties
+++ b/tools/src/test/resources/log4j.properties
@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-log4j.rootLogger=TRACE, stdout
+log4j.rootLogger=INFO, stdout
 
 log4j.appender.stdout=org.apache.log4j.ConsoleAppender
 log4j.appender.stdout.layout=org.apache.log4j.PatternLayout

From 96d026bc994fb305c366ce4de27804fcd1f7899a Mon Sep 17 00:00:00 2001
From: Dung Ha <60119105+infantlikesprogramming@users.noreply.github.com>
Date: Thu, 7 Mar 2024 17:43:57 -0500
Subject: [PATCH 131/258] KAFKA-16202 Extra dot in error message in producer
 (#15296)

The author of KAFKA-16202 noticed that there is an extra dot in the error message for KafkaStorageException message.

Looking into org.apache.kafka.clients.producer.internals.Sender, it turns out that the string for the message to be sent in completeBatch() added an extra dot. I think that the formatted component (error.exception(response.errorMessage).toString())) of the error message already has a dot in the end of its string. Thus the dot after the "{}" sign caused the extra dot.

Reviewers: "Gyeongwon, Do" <dct012@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/clients/producer/internals/Sender.java     |  5 ++---
 .../org/apache/kafka/common/protocol/Errors.java     | 12 ++++++------
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
index 2cdc3b876d0b1..99bc1d68b0b22 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
@@ -706,9 +706,8 @@ private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionRespons
                             "topic-partition may not exist or the user may not have Describe access to it",
                         batch.topicPartition);
                 } else {
-                    log.warn("Received invalid metadata error in produce request on partition {} due to {}. Going " +
-                            "to request metadata update now", batch.topicPartition,
-                            error.exception(response.errorMessage).toString());
+                    log.warn("Received invalid metadata error in produce request on partition {} due to {} Going " +
+                            "to request metadata update now", batch.topicPartition, error.exception(response.errorMessage).toString());
                 }
                 if (error.exception() instanceof NotLeaderOrFollowerException || error.exception() instanceof FencedLeaderEpochException) {
                     log.debug("For {}, received error {}, with leaderIdAndEpoch {}", batch.topicPartition, error, response.currentLeader);
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
index 838be869f6a07..610b1b66e7633 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
@@ -363,12 +363,12 @@ public enum Errors {
     DUPLICATE_RESOURCE(92, "A request illegally referred to the same resource twice.", DuplicateResourceException::new),
     UNACCEPTABLE_CREDENTIAL(93, "Requested credential would not meet criteria for acceptability.", UnacceptableCredentialException::new),
     INCONSISTENT_VOTER_SET(94, "Indicates that the either the sender or recipient of a " +
-            "voter-only request is not one of the expected voters", InconsistentVoterSetException::new),
+            "voter-only request is not one of the expected voters.", InconsistentVoterSetException::new),
     INVALID_UPDATE_VERSION(95, "The given update version was invalid.", InvalidUpdateVersionException::new),
     FEATURE_UPDATE_FAILED(96, "Unable to update finalized features due to an unexpected server error.", FeatureUpdateFailedException::new),
     PRINCIPAL_DESERIALIZATION_FAILURE(97, "Request principal deserialization failed during forwarding. " +
          "This indicates an internal error on the broker cluster security setup.", PrincipalDeserializationException::new),
-    SNAPSHOT_NOT_FOUND(98, "Requested snapshot was not found", SnapshotNotFoundException::new),
+    SNAPSHOT_NOT_FOUND(98, "Requested snapshot was not found.", SnapshotNotFoundException::new),
     POSITION_OUT_OF_RANGE(
         99,
         "Requested position is not greater than or equal to zero, and less than the size of the snapshot.",
@@ -376,10 +376,10 @@ public enum Errors {
     UNKNOWN_TOPIC_ID(100, "This server does not host this topic ID.", UnknownTopicIdException::new),
     DUPLICATE_BROKER_REGISTRATION(101, "This broker ID is already in use.", DuplicateBrokerRegistrationException::new),
     BROKER_ID_NOT_REGISTERED(102, "The given broker ID was not registered.", BrokerIdNotRegisteredException::new),
-    INCONSISTENT_TOPIC_ID(103, "The log's topic ID did not match the topic ID in the request", InconsistentTopicIdException::new),
-    INCONSISTENT_CLUSTER_ID(104, "The clusterId in the request does not match that found on the server", InconsistentClusterIdException::new),
-    TRANSACTIONAL_ID_NOT_FOUND(105, "The transactionalId could not be found", TransactionalIdNotFoundException::new),
-    FETCH_SESSION_TOPIC_ID_ERROR(106, "The fetch session encountered inconsistent topic ID usage", FetchSessionTopicIdException::new),
+    INCONSISTENT_TOPIC_ID(103, "The log's topic ID did not match the topic ID in the request.", InconsistentTopicIdException::new),
+    INCONSISTENT_CLUSTER_ID(104, "The clusterId in the request does not match that found on the server.", InconsistentClusterIdException::new),
+    TRANSACTIONAL_ID_NOT_FOUND(105, "The transactionalId could not be found.", TransactionalIdNotFoundException::new),
+    FETCH_SESSION_TOPIC_ID_ERROR(106, "The fetch session encountered inconsistent topic ID usage.", FetchSessionTopicIdException::new),
     INELIGIBLE_REPLICA(107, "The new ISR contains at least one ineligible replica.", IneligibleReplicaException::new),
     NEW_LEADER_ELECTED(108, "The AlterPartition request successfully updated the partition state but the leader has changed.", NewLeaderElectedException::new),
     OFFSET_MOVED_TO_TIERED_STORAGE(109, "The requested offset is moved to tiered storage.", OffsetMovedToTieredStorageException::new),

From 3f3cee1b21f71bcc3f478514bc996e6c9ef6ccaa Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Fri, 8 Mar 2024 07:07:21 +0800
Subject: [PATCH 132/258] KAFKA-16325 Add missing producer metrics to
 documentatio (#15466)

Add `buffer-exhausted-rate`, `buffer-exhausted-total`, `bufferpool-wait-ratio` and `metadata-wait-time-ns-total`

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 docs/ops.html | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/docs/ops.html b/docs/ops.html
index 92f676155570e..9f49b5c1b2165 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -2537,11 +2537,26 @@ <h4 class="anchor-heading"><a id="producer_monitoring" class="anchor-link"></a><
         <td>The total amount of buffer memory that is not being used (either unallocated or in the free list).</td>
         <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
       </tr>
+      <tr>
+        <td>buffer-exhausted-rate</td>
+        <td>The average per-second number of record sends that are dropped due to buffer exhaustion</td>
+        <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
+      </tr>
+      <tr>
+        <td>buffer-exhausted-total</td>
+        <td>The total number of record sends that are dropped due to buffer exhaustion</td>
+        <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
+      </tr>
       <tr>
         <td>bufferpool-wait-time</td>
         <td>The fraction of time an appender waits for space allocation.</td>
         <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
       </tr>
+      <tr>
+        <td>bufferpool-wait-ratio</td>
+        <td>The fraction of time an appender waits for space allocation.</td>
+        <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
+      </tr>
       <tr>
         <td>bufferpool-wait-time-total</td>
         <td><b>*Deprecated*</b> The total time an appender waits for space allocation in nanoseconds. Replacement is <code>bufferpool-wait-time-ns-total</code></td>
@@ -2582,7 +2597,11 @@ <h4 class="anchor-heading"><a id="producer_monitoring" class="anchor-link"></a><
         <td>The total time the Producer spent aborting transactions in nanoseconds (for EOS).</td>
         <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
       </tr>
-
+      <tr>
+        <td>metadata-wait-time-ns-total</td>
+        <td>the total time in nanoseconds that has spent waiting for metadata from the Kafka broker</td>
+        <td>kafka.producer:type=producer-metrics,client-id=([-.\w]+)</td>
+      </tr>
   </tbody></table>
 
   <h5 class="anchor-heading"><a id="producer_sender_monitoring" class="anchor-link"></a><a href="#producer_sender_monitoring">Producer Sender Metrics</a></h5>

From 80def43a3438f74344f034f6a326c02d76151c80 Mon Sep 17 00:00:00 2001
From: testn <test1@doramail.com>
Date: Fri, 8 Mar 2024 23:43:44 +0700
Subject: [PATCH 133/258] MINOR: Reduce memory allocation in
 ClientTelemetryReporter (#15402)

Reviewers: Divij Vaidya <diviv@amazon.com>
---
 checkstyle/suppressions.xml                   |  2 +-
 .../internals/ClientTelemetryReporter.java    | 19 +++++++++++--------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 32916fe6602fd..e384ad8c82243 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -93,7 +93,7 @@
               files="(Utils|Topic|KafkaLZ4BlockOutputStream|AclData|JoinGroupRequest).java"/>
 
     <suppress checks="CyclomaticComplexity"
-              files="(AbstractFetch|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/>
+              files="(AbstractFetch|ClientTelemetryReporter|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/>
 
     <suppress checks="JavaNCSS"
               files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|WorkerSinkTaskMockitoTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest|KafkaRaftClientTest).java"/>
diff --git a/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporter.java b/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporter.java
index ae92c539dafb1..483179d4c456d 100644
--- a/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporter.java
+++ b/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporter.java
@@ -265,8 +265,8 @@ class DefaultClientTelemetrySender implements ClientTelemetrySender {
          These are the lower and upper bounds of the jitter that we apply to the initial push
          telemetry API call. This helps to avoid a flood of requests all coming at the same time.
         */
-        private final static double INITIAL_PUSH_JITTER_LOWER = 0.5;
-        private final static double INITIAL_PUSH_JITTER_UPPER = 1.5;
+        private static final double INITIAL_PUSH_JITTER_LOWER = 0.5;
+        private static final double INITIAL_PUSH_JITTER_UPPER = 1.5;
 
         private final ReadWriteLock lock = new ReentrantReadWriteLock();
         private final Condition subscriptionLoaded = lock.writeLock().newCondition();
@@ -325,6 +325,7 @@ public long timeToNextUpdate(long requestTimeoutMs) {
             final long timeMs;
             final String apiName;
             final String msg;
+            final boolean isTraceEnabled = log.isTraceEnabled();
 
             switch (localState) {
                 case SUBSCRIPTION_IN_PROGRESS:
@@ -336,15 +337,15 @@ public long timeToNextUpdate(long requestTimeoutMs) {
                     */
                     apiName = (localState == ClientTelemetryState.SUBSCRIPTION_IN_PROGRESS) ? ApiKeys.GET_TELEMETRY_SUBSCRIPTIONS.name : ApiKeys.PUSH_TELEMETRY.name;
                     timeMs = requestTimeoutMs;
-                    msg = String.format("the remaining wait time for the %s network API request, as specified by %s", apiName, CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG);
+                    msg = isTraceEnabled ? "" : String.format("the remaining wait time for the %s network API request, as specified by %s", apiName, CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG);
                     break;
                 case TERMINATING_PUSH_IN_PROGRESS:
                     timeMs = Long.MAX_VALUE;
-                    msg = String.format("the terminating push is in progress, disabling telemetry for further requests");
+                    msg = isTraceEnabled ? "" : "the terminating push is in progress, disabling telemetry for further requests";
                     break;
                 case TERMINATING_PUSH_NEEDED:
                     timeMs = 0;
-                    msg = String.format("the client should try to submit the final %s network API request ASAP before closing", ApiKeys.PUSH_TELEMETRY.name);
+                    msg = isTraceEnabled ? "" : String.format("the client should try to submit the final %s network API request ASAP before closing", ApiKeys.PUSH_TELEMETRY.name);
                     break;
                 case SUBSCRIPTION_NEEDED:
                 case PUSH_NEEDED:
@@ -352,17 +353,19 @@ public long timeToNextUpdate(long requestTimeoutMs) {
                     long timeRemainingBeforeRequest = localLastRequestMs + localIntervalMs - nowMs;
                     if (timeRemainingBeforeRequest <= 0) {
                         timeMs = 0;
-                        msg = String.format("the wait time before submitting the next %s network API request has elapsed", apiName);
+                        msg = isTraceEnabled ? "" : String.format("the wait time before submitting the next %s network API request has elapsed", apiName);
                     } else {
                         timeMs = timeRemainingBeforeRequest;
-                        msg = String.format("the client will wait before submitting the next %s network API request", apiName);
+                        msg = isTraceEnabled ? "" : String.format("the client will wait before submitting the next %s network API request", apiName);
                     }
                     break;
                 default:
                     throw new IllegalStateException("Unknown telemetry state: " + localState);
             }
 
-            log.trace("For telemetry state {}, returning the value {} ms; {}", localState, timeMs, msg);
+            if (isTraceEnabled) {
+                log.trace("For telemetry state {}, returning the value {} ms; {}", localState, timeMs, msg);
+            }
             return timeMs;
         }
 

From b9a5b4a8053c1fa65e27a9f93440194b0dd5eec4 Mon Sep 17 00:00:00 2001
From: Daan Gerits <daan.gerits@gmail.com>
Date: Fri, 8 Mar 2024 19:57:56 +0100
Subject: [PATCH 134/258] KAFKA-10892: Shared Readonly State Stores ( revisited
 ) (#12742)

Implements KIP-813.

Reviewers: Matthias J. Sax <matthias@confluent.io>, Walker Carlson <wcarlson@confluent.io>
---
 checkstyle/suppressions.xml                   |   2 +-
 .../developer-guide/processor-api.html        |  13 ++
 docs/streams/upgrade-guide.html               |   8 +
 .../org/apache/kafka/streams/Topology.java    |  82 ++++++++++
 .../apache/kafka/streams/TopologyTest.java    |  63 +++++++-
 .../integration/EosIntegrationTest.java       |  24 ++-
 .../integration/RestoreIntegrationTest.java   | 145 ++++++++++++------
 .../utils/IntegrationTestUtils.java           |  12 ++
 .../streams/processor/ReadOnlyStoreTest.java  | 132 ++++++++++++++++
 .../org/apache/kafka/test/MockProcessor.java  |   5 +-
 10 files changed, 422 insertions(+), 64 deletions(-)
 create mode 100644 streams/src/test/java/org/apache/kafka/streams/processor/ReadOnlyStoreTest.java

diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index e384ad8c82243..27c252f83906c 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -234,7 +234,7 @@
 
     <!-- Streams tests -->
     <suppress checks="ClassFanOutComplexity"
-              files="(RecordCollectorTest|StreamsPartitionAssignorTest|StreamThreadTest|StreamTaskTest|TaskManagerTest|TopologyTestDriverTest|KafkaStreamsTest|EosIntegrationTest).java"/>
+              files="(RecordCollectorTest|StreamsPartitionAssignorTest|StreamThreadTest|StreamTaskTest|TaskManagerTest|TopologyTestDriverTest|KafkaStreamsTest|EosIntegrationTest|RestoreIntegrationTest).java"/>
 
     <suppress checks="MethodLength"
               files="(EosIntegrationTest|EosV2UpgradeIntegrationTest|KStreamKStreamJoinTest|RocksDBWindowStoreTest|StreamStreamJoinIntegrationTest).java"/>
diff --git a/docs/streams/developer-guide/processor-api.html b/docs/streams/developer-guide/processor-api.html
index ccb03ce7b50c1..e25ecd601ac5c 100644
--- a/docs/streams/developer-guide/processor-api.html
+++ b/docs/streams/developer-guide/processor-api.html
@@ -51,6 +51,7 @@
                         <li><a class="reference internal" href="#enable-or-disable-fault-tolerance-of-state-stores-store-changelogs" id="id6">Enable or Disable Fault Tolerance of State Stores (Store Changelogs)</a></li>
                         <li><a class="reference internal" href="#timestamped-state-stores" id="id11">Timestamped State Stores</a></li>
                         <li><a class="reference internal" href="#versioned-state-stores" id="id12">Versioned Key-Value State Stores</a></li>
+                        <li><a class="reference internal" href="#readonly-state-stores" id="id12">Readonly State Stores</a></li>
                         <li><a class="reference internal" href="#implementing-custom-state-stores" id="id7">Implementing Custom State Stores</a></li>
                     </ul>
                 </li>
@@ -466,6 +467,18 @@ <h2>
                         stores to rebuild state from changelog.</li>
                 </ul>
                 </p>
+            <div class="section" id="readonly-state-stores">
+                <span id="streams-developer-guide-state-store-readonly"></span><h3><a class="toc-backref" href="#id12">ReadOnly State Stores</a><a class="headerlink" href="#readonly-state-stores" title="Permalink to this headline"></a></h3>
+                <p>A read-only state store materialized the data from its input topic. It also uses the input topic
+                    for fault-tolerance, and thus does not have an additional changelog topic (the input topic is
+                    re-used as changelog). Thus, the input topic should be configured with <a class="reference external" href="https://kafka.apache.org/documentation.html#compaction">log compaction</a>.
+                    Note that no other processor should modify the content of the state store, and the only writer
+                    should be the associated "state update processor"; other processors may read the content of the
+                    read-only store.</p>
+
+                <p><b>note:</b> beware of the partitioning requirements when using read-only state stores for lookups during
+                    processing. You might want to make sure the original changelog topic is co-partitioned with the processors
+                    reading the read-only statestore.</p>
             </div>
             <div class="section" id="implementing-custom-state-stores">
                 <span id="streams-developer-guide-state-store-custom"></span><h3><a class="toc-backref" href="#id7">Implementing Custom State Stores</a><a class="headerlink" href="#implementing-custom-state-stores" title="Permalink to this headline"></a></h3>
diff --git a/docs/streams/upgrade-guide.html b/docs/streams/upgrade-guide.html
index 14544077403cd..0f819cb384f92 100644
--- a/docs/streams/upgrade-guide.html
+++ b/docs/streams/upgrade-guide.html
@@ -133,6 +133,14 @@ <h3 class="anchor-heading"><a id="streams_notable_changes" class="anchor-link"><
         More details about the new config <code>StreamsConfig#TOPOLOGY_OPTIMIZATION</code> can be found in <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-295%3A+Add+Streams+Configuration+Allowing+for+Optional+Topology+Optimization">KIP-295</a>.
     </p>
 
+    <h4><a id="streams_api_changes_380" href="#streams_api_changes_380">Streams API changes in 3.8.0</a></h3>
+    <p>
+        The Processor API now support so-called read-only state stores, added via 
+        <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-813%3A+Shareable+State+Stores">KIP-813</a>.
+	These stores don't have a dedicated changelog topic, but use their source topic for fault-tolerance,
+	simlar to <code>KTables</code> with source-topic optimization enabled.
+    </p>
+
     <h3><a id="streams_api_changes_370" href="#streams_api_changes_370">Streams API changes in 3.7.0</a></h3>
     <p>
         We added a new method to <code>KafkaStreams</code>, namely <code>KafkaStreams#setStandbyUpdateListener()</code> in
diff --git a/streams/src/main/java/org/apache/kafka/streams/Topology.java b/streams/src/main/java/org/apache/kafka/streams/Topology.java
index d9c810aa42599..d1f0d1eb8bc8e 100644
--- a/streams/src/main/java/org/apache/kafka/streams/Topology.java
+++ b/streams/src/main/java/org/apache/kafka/streams/Topology.java
@@ -738,6 +738,88 @@ public synchronized Topology addStateStore(final StoreBuilder<?> storeBuilder,
         return this;
     }
 
+    /**
+     * Adds a read-only {@link StateStore} to the topology.
+     * <p>
+     * A read-only {@link StateStore} does not create a dedicated changelog topic but uses it's input topic as
+     * changelog; thus, the used topic should be configured with log compaction.
+     * <p>
+     * The <code>auto.offset.reset</code> property will be set to <code>earliest</code> for this topic.
+     * <p>
+     * The provided {@link ProcessorSupplier} will be used to create a processor for all messages received
+     * from the given topic. This processor should contain logic to keep the {@link StateStore} up-to-date.
+     *
+     * @param storeBuilder          user defined store builder
+     * @param sourceName            name of the {@link SourceNode} that will be automatically added
+     * @param timestampExtractor    the stateless timestamp extractor used for this source,
+     *                              if not specified the default extractor defined in the configs will be used
+     * @param keyDeserializer       the {@link Deserializer} to deserialize keys with
+     * @param valueDeserializer     the {@link Deserializer} to deserialize values with
+     * @param topic                 the topic to source the data from
+     * @param processorName         the name of the {@link ProcessorSupplier}
+     * @param stateUpdateSupplier   the instance of {@link ProcessorSupplier}
+     * @return itself
+     * @throws TopologyException if the processor of state is already registered
+     */
+    public synchronized <KIn, VIn> Topology addReadOnlyStateStore(final StoreBuilder<?> storeBuilder,
+                                                                  final String sourceName,
+                                                                  final TimestampExtractor timestampExtractor,
+                                                                  final Deserializer<KIn> keyDeserializer,
+                                                                  final Deserializer<VIn> valueDeserializer,
+                                                                  final String topic,
+                                                                  final String processorName,
+                                                                  final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier) {
+        storeBuilder.withLoggingDisabled();
+
+        internalTopologyBuilder.addSource(AutoOffsetReset.EARLIEST, sourceName, timestampExtractor, keyDeserializer, valueDeserializer, topic);
+        internalTopologyBuilder.addProcessor(processorName, stateUpdateSupplier, sourceName);
+        internalTopologyBuilder.addStateStore(storeBuilder, processorName);
+        internalTopologyBuilder.connectSourceStoreAndTopic(storeBuilder.name(), topic);
+
+        return this;
+    }
+
+    /**
+     * Adds a read-only {@link StateStore} to the topology.
+     * <p>
+     * A read-only {@link StateStore} does not create a dedicated changelog topic but uses it's input topic as
+     * changelog; thus, the used topic should be configured with log compaction.
+     * <p>
+     * The <code>auto.offset.reset</code> property will be set to <code>earliest</code> for this topic.
+     * <p>
+     * The provided {@link ProcessorSupplier} will be used to create a processor for all messages received
+     * from the given topic. This processor should contain logic to keep the {@link StateStore} up-to-date.
+     * The default {@link TimestampExtractor} as specified in the {@link StreamsConfig config} is used.
+     *
+     * @param storeBuilder          user defined store builder
+     * @param sourceName            name of the {@link SourceNode} that will be automatically added
+     * @param keyDeserializer       the {@link Deserializer} to deserialize keys with
+     * @param valueDeserializer     the {@link Deserializer} to deserialize values with
+     * @param topic                 the topic to source the data from
+     * @param processorName         the name of the {@link ProcessorSupplier}
+     * @param stateUpdateSupplier   the instance of {@link ProcessorSupplier}
+     * @return itself
+     * @throws TopologyException if the processor of state is already registered
+     */
+    public synchronized <KIn, VIn> Topology addReadOnlyStateStore(final StoreBuilder<?> storeBuilder,
+                                                                  final String sourceName,
+                                                                  final Deserializer<KIn> keyDeserializer,
+                                                                  final Deserializer<VIn> valueDeserializer,
+                                                                  final String topic,
+                                                                  final String processorName,
+                                                                  final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier) {
+        return addReadOnlyStateStore(
+                storeBuilder,
+                sourceName,
+                null,
+                keyDeserializer,
+                valueDeserializer,
+                topic,
+                processorName,
+                stateUpdateSupplier
+        );
+    }
+
     /**
      * Adds a global {@link StateStore} to the topology.
      * The {@link StateStore} sources its data from all partitions of the provided input topic.
diff --git a/streams/src/test/java/org/apache/kafka/streams/TopologyTest.java b/streams/src/test/java/org/apache/kafka/streams/TopologyTest.java
index 5fdf5c220cfa2..ceecbc3bc2178 100644
--- a/streams/src/test/java/org/apache/kafka/streams/TopologyTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/TopologyTest.java
@@ -38,6 +38,7 @@
 import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder;
 import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder.SubtopologyDescription;
 import org.apache.kafka.streams.processor.internals.ProcessorTopology;
+import org.apache.kafka.streams.processor.internals.StoreFactory;
 import org.apache.kafka.streams.state.KeyValueStore;
 import org.apache.kafka.streams.state.SessionStore;
 import org.apache.kafka.streams.state.StoreBuilder;
@@ -58,6 +59,7 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.internal.util.collections.Sets;
 
 import java.time.Duration;
 import java.util.Arrays;
@@ -337,7 +339,7 @@ public void shouldNotAllowToAddStoreWithSameNameAndDifferentInstance() {
         mockStoreBuilder();
         topology.addStateStore(storeBuilder);
 
-        final StoreBuilder otherStoreBuilder = mock(StoreBuilder.class);
+        final StoreBuilder<?> otherStoreBuilder = mock(StoreBuilder.class);
         when(otherStoreBuilder.name()).thenReturn("store");
         when(otherStoreBuilder.logConfig()).thenReturn(Collections.emptyMap());
         when(otherStoreBuilder.loggingEnabled()).thenReturn(false);
@@ -2313,7 +2315,7 @@ private TopologyDescription.Sink addSink(final String sinkName,
 
         topology.addSink(sinkName, sinkTopic, null, null, null, parentNames);
         final TopologyDescription.Sink expectedSinkNode =
-            new InternalTopologyBuilder.Sink(sinkName, sinkTopic);
+            new InternalTopologyBuilder.Sink<>(sinkName, sinkTopic);
 
         for (final TopologyDescription.Node parent : parents) {
             ((InternalTopologyBuilder.AbstractNode) parent).addSuccessor(expectedSinkNode);
@@ -2351,6 +2353,63 @@ private void addGlobalStoreToTopologyAndExpectedDescription(final String globalS
         expectedDescription.addGlobalStore(expectedGlobalStore);
     }
 
+    @Test
+    public void readOnlyStateStoresShouldHaveTheirOwnSubTopology() {
+        final String sourceName = "source";
+        final String storeName = "store";
+        final String topicName = "topic";
+        final String processorName = "processor";
+
+        final KeyValueStoreBuilder<?, ?> storeBuilder = mock(KeyValueStoreBuilder.class);
+        when(storeBuilder.name()).thenReturn(storeName);
+        topology.addReadOnlyStateStore(
+                storeBuilder,
+                sourceName,
+                null,
+                null,
+                null,
+                topicName,
+                processorName,
+                new MockProcessorSupplier<>());
+
+        final TopologyDescription.Source expectedSource = new InternalTopologyBuilder.Source(sourceName, Sets.newSet(topicName), null);
+        final TopologyDescription.Processor expectedProcessor = new InternalTopologyBuilder.Processor(processorName, Sets.newSet(storeName));
+
+        ((InternalTopologyBuilder.AbstractNode) expectedSource).addSuccessor(expectedProcessor);
+        ((InternalTopologyBuilder.AbstractNode) expectedProcessor).addPredecessor(expectedSource);
+
+        final Set<TopologyDescription.Node> allNodes = new HashSet<>();
+        allNodes.add(expectedSource);
+        allNodes.add(expectedProcessor);
+        expectedDescription.addSubtopology(new SubtopologyDescription(0, allNodes));
+
+        assertThat(topology.describe(), equalTo(expectedDescription));
+        assertThat(topology.describe().hashCode(), equalTo(expectedDescription.hashCode()));
+    }
+
+    @Test
+    public void readOnlyStateStoresShouldNotLog() {
+        final String sourceName = "source";
+        final String storeName = "store";
+        final String topicName = "topic";
+        final String processorName = "processor";
+
+        final KeyValueStoreBuilder<?, ?> storeBuilder = mock(KeyValueStoreBuilder.class);
+        when(storeBuilder.name()).thenReturn(storeName);
+        topology.addReadOnlyStateStore(
+                storeBuilder,
+                sourceName,
+                null,
+                null,
+                null,
+                topicName,
+                processorName,
+                new MockProcessorSupplier<>());
+
+        final StoreFactory stateStoreFactory = topology.internalTopologyBuilder.stateStores().get(storeName);
+        assertThat(stateStoreFactory.loggingEnabled(), equalTo(false));
+    }
+
     private TopologyConfig overrideDefaultStore(final String defaultStore) {
         final Properties topologyOverrides = new Properties();
         // change default store as in-memory
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
index d79631f8f7a0c..7989de0e76bb7 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
@@ -123,8 +123,8 @@ public class EosIntegrationTest {
     public Timeout globalTimeout = Timeout.seconds(600);
     private static final Logger LOG = LoggerFactory.getLogger(EosIntegrationTest.class);
     private static final int NUM_BROKERS = 3;
-    private static final int MAX_POLL_INTERVAL_MS = 5 * 1000;
-    private static final int MAX_WAIT_TIME_MS = 60 * 1000;
+    private static final int MAX_POLL_INTERVAL_MS = 30_1000;
+    private static final int MAX_WAIT_TIME_MS = 120_1000;
 
     public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(
         NUM_BROKERS,
@@ -403,7 +403,7 @@ public void shouldNotViolateEosIfOneTaskFails() throws Exception {
         // -> the failure only kills one thread
         // after fail over, we should read 40 committed records (even if 50 record got written)
 
-        try (final KafkaStreams streams = getKafkaStreams("dummy", false, "appDir", 2, eosConfig, MAX_POLL_INTERVAL_MS)) {
+        try (final KafkaStreams streams = getKafkaStreams("dummy", false, "appDir", 2, eosConfig)) {
             startApplicationAndWaitUntilRunning(streams);
 
             final List<KeyValue<Long, Long>> committedDataBeforeFailure = prepareData(0L, 10L, 0L, 1L);
@@ -511,7 +511,7 @@ public void shouldNotViolateEosIfOneTaskFailsWithState() throws Exception {
 
         // We need more processing time under "with state" situation, so increasing the max.poll.interval.ms
         // to avoid unexpected rebalance during test, which will cause unexpected fail over triggered
-        try (final KafkaStreams streams = getKafkaStreams("dummy", true, "appDir", 2, eosConfig, 3 * MAX_POLL_INTERVAL_MS)) {
+        try (final KafkaStreams streams = getKafkaStreams("dummy", true, "appDir", 2, eosConfig)) {
             startApplicationAndWaitUntilRunning(streams);
 
             final List<KeyValue<Long, Long>> committedDataBeforeFailure = prepareData(0L, 10L, 0L, 1L);
@@ -624,12 +624,12 @@ public void shouldNotViolateEosIfOneTaskGetsFencedUsingIsolatedAppInstances() th
         // -> the stall only affects one thread and should trigger a rebalance
         // after rebalancing, we should read 40 committed records (even if 50 record got written)
         //
-        // afterwards, the "stalling" thread resumes, and another rebalance should get triggered
+        // afterward, the "stalling" thread resumes, and another rebalance should get triggered
         // we write the remaining 20 records and verify to read 60 result records
 
         try (
-            final KafkaStreams streams1 = getKafkaStreams("streams1", false, "appDir1", 1, eosConfig, MAX_POLL_INTERVAL_MS);
-            final KafkaStreams streams2 = getKafkaStreams("streams2", false, "appDir2", 1, eosConfig, MAX_POLL_INTERVAL_MS)
+            final KafkaStreams streams1 = getKafkaStreams("streams1", false, "appDir1", 1, eosConfig);
+            final KafkaStreams streams2 = getKafkaStreams("streams2", false, "appDir2", 1, eosConfig)
         ) {
             startApplicationAndWaitUntilRunning(streams1);
             startApplicationAndWaitUntilRunning(streams2);
@@ -778,7 +778,7 @@ public void shouldWriteLatestOffsetsToCheckpointOnShutdown() throws Exception {
         final List<KeyValue<Long, Long>> writtenData = prepareData(0L, 10, 0L, 1L);
         final List<KeyValue<Long, Long>> expectedResult = computeExpectedResult(writtenData);
 
-        try (final KafkaStreams streams = getKafkaStreams("streams", true, "appDir", 1, eosConfig, MAX_POLL_INTERVAL_MS)) {
+        try (final KafkaStreams streams = getKafkaStreams("streams", true, "appDir", 1, eosConfig)) {
             writeInputData(writtenData);
 
             startApplicationAndWaitUntilRunning(streams);
@@ -1004,8 +1004,7 @@ private KafkaStreams getKafkaStreams(final String dummyHostName,
                                          final boolean withState,
                                          final String appDir,
                                          final int numberOfStreamsThreads,
-                                         final String eosConfig,
-                                         final int maxPollIntervalMs) {
+                                         final String eosConfig) {
         commitRequested = new AtomicInteger(0);
         errorInjected = new AtomicBoolean(false);
         stallInjected = new AtomicBoolean(false);
@@ -1112,9 +1111,8 @@ public void close() { }
         properties.put(StreamsConfig.producerPrefix(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG), (int) commitIntervalMs);
         properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.METADATA_MAX_AGE_CONFIG), "1000");
         properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG), "earliest");
-        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG), maxPollIntervalMs);
-        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), maxPollIntervalMs - 1);
-        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG), maxPollIntervalMs);
+        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), MAX_POLL_INTERVAL_MS - 1);
+        properties.put(StreamsConfig.consumerPrefix(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG), MAX_POLL_INTERVAL_MS);
         properties.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
         properties.put(StreamsConfig.STATE_DIR_CONFIG, stateTmpDir + appDir);
         properties.put(StreamsConfig.APPLICATION_SERVER_CONFIG, dummyHostName + ":2142");
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
index 812da30074775..b41e0c3d29f2f 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.serialization.IntegerDeserializer;
 import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Utils;
@@ -155,8 +156,8 @@ private Properties props(final Properties extraProperties) {
         streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
         streamsConfiguration.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
         streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory(appId).getPath());
-        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
-        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
+        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
+        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.IntegerSerde.class);
         streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000L);
         streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
         streamsConfiguration.putAll(extraProperties);
@@ -196,8 +197,8 @@ public void shouldRestoreNullRecord() throws Exception {
         final Properties streamsConfiguration = StreamsTestUtils.getStreamsConfig(
                 applicationId,
                 CLUSTER.bootstrapServers(),
-                Serdes.Integer().getClass().getName(),
-                Serdes.ByteArray().getClass().getName(),
+                Serdes.IntegerSerde.class.getName(),
+                Serdes.BytesSerde.class.getName(),
                 props);
 
         CLUSTER.createTopics(inputTopic);
@@ -249,7 +250,63 @@ public void shouldRestoreNullRecord() throws Exception {
 
     @ParameterizedTest
     @MethodSource("parameters")
-    public void shouldRestoreStateFromSourceTopic(final boolean stateUpdaterEnabled) throws Exception {
+    public void shouldRestoreStateFromSourceTopicForReadOnlyStore(final boolean stateUpdaterEnabled) throws Exception {
+        final AtomicInteger numReceived = new AtomicInteger(0);
+        final Topology topology = new Topology();
+
+        final Properties props = props(stateUpdaterEnabled);
+
+        // restoring from 1000 to 4000 (committed), and then process from 4000 to 5000 on each of the two partitions
+        final int offsetLimitDelta = 1000;
+        final int offsetCheckpointed = 1000;
+        createStateForRestoration(inputStream, 0);
+        setCommittedOffset(inputStream, offsetLimitDelta);
+
+        final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true, false);
+        // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
+        new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 0)), ".checkpoint"))
+            .write(Collections.singletonMap(new TopicPartition(inputStream, 0), (long) offsetCheckpointed - 1));
+        new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 1)), ".checkpoint"))
+            .write(Collections.singletonMap(new TopicPartition(inputStream, 1), (long) offsetCheckpointed - 1));
+
+        final CountDownLatch startupLatch = new CountDownLatch(1);
+        final CountDownLatch shutdownLatch = new CountDownLatch(1);
+
+        topology.addReadOnlyStateStore(
+            Stores.keyValueStoreBuilder(
+                Stores.persistentKeyValueStore("store"),
+                new Serdes.IntegerSerde(),
+                new Serdes.StringSerde()
+            ),
+            "readOnlySource",
+            new IntegerDeserializer(),
+            new StringDeserializer(),
+            inputStream,
+            "readOnlyProcessor",
+            () -> new ReadOnlyStoreProcessor(numReceived, offsetLimitDelta, shutdownLatch)
+        );
+
+        kafkaStreams = new KafkaStreams(topology, props);
+        kafkaStreams.setStateListener((newState, oldState) -> {
+            if (newState == KafkaStreams.State.RUNNING && oldState == KafkaStreams.State.REBALANCING) {
+                startupLatch.countDown();
+            }
+        });
+
+        final AtomicLong restored = new AtomicLong(0);
+        kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
+        kafkaStreams.start();
+
+        assertTrue(startupLatch.await(30, TimeUnit.SECONDS));
+        assertThat(restored.get(), equalTo((long) numberOfKeys - offsetLimitDelta * 2 - offsetCheckpointed * 2));
+
+        assertTrue(shutdownLatch.await(30, TimeUnit.SECONDS));
+        assertThat(numReceived.get(), equalTo(offsetLimitDelta * 2));
+    }
+
+    @ParameterizedTest
+    @MethodSource("parameters")
+    public void shouldRestoreStateFromSourceTopicForGlobalTable(final boolean stateUpdaterEnabled) throws Exception {
         final AtomicInteger numReceived = new AtomicInteger(0);
         final StreamsBuilder builder = new StreamsBuilder();
 
@@ -265,9 +322,9 @@ public void shouldRestoreStateFromSourceTopic(final boolean stateUpdaterEnabled)
         final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true, false);
         // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
         new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 0)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(inputStream, 0), (long) offsetCheckpointed - 1));
+            .write(Collections.singletonMap(new TopicPartition(inputStream, 0), (long) offsetCheckpointed - 1));
         new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 1)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(inputStream, 1), (long) offsetCheckpointed - 1));
+            .write(Collections.singletonMap(new TopicPartition(inputStream, 1), (long) offsetCheckpointed - 1));
 
         final CountDownLatch startupLatch = new CountDownLatch(1);
         final CountDownLatch shutdownLatch = new CountDownLatch(1);
@@ -288,22 +345,7 @@ public void shouldRestoreStateFromSourceTopic(final boolean stateUpdaterEnabled)
         });
 
         final AtomicLong restored = new AtomicLong(0);
-        kafkaStreams.setGlobalStateRestoreListener(new StateRestoreListener() {
-            @Override
-            public void onRestoreStart(final TopicPartition topicPartition, final String storeName, final long startingOffset, final long endingOffset) {
-
-            }
-
-            @Override
-            public void onBatchRestored(final TopicPartition topicPartition, final String storeName, final long batchEndOffset, final long numRestored) {
-
-            }
-
-            @Override
-            public void onRestoreEnd(final TopicPartition topicPartition, final String storeName, final long totalRestored) {
-                restored.addAndGet(totalRestored);
-            }
-        });
+        kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
         kafkaStreams.start();
 
         assertTrue(startupLatch.await(30, TimeUnit.SECONDS));
@@ -332,9 +374,9 @@ public void shouldRestoreStateFromChangelogTopic(final boolean stateUpdaterEnabl
         final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true, false);
         // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
         new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 0)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(changelog, 0), (long) offsetCheckpointed - 1));
+            .write(Collections.singletonMap(new TopicPartition(changelog, 0), (long) offsetCheckpointed - 1));
         new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 1)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(changelog, 1), (long) offsetCheckpointed - 1));
+            .write(Collections.singletonMap(new TopicPartition(changelog, 1), (long) offsetCheckpointed - 1));
 
         final CountDownLatch startupLatch = new CountDownLatch(1);
         final CountDownLatch shutdownLatch = new CountDownLatch(1);
@@ -355,22 +397,7 @@ public void shouldRestoreStateFromChangelogTopic(final boolean stateUpdaterEnabl
         });
 
         final AtomicLong restored = new AtomicLong(0);
-        kafkaStreams.setGlobalStateRestoreListener(new StateRestoreListener() {
-            @Override
-            public void onRestoreStart(final TopicPartition topicPartition, final String storeName, final long startingOffset, final long endingOffset) {
-
-            }
-
-            @Override
-            public void onBatchRestored(final TopicPartition topicPartition, final String storeName, final long batchEndOffset, final long numRestored) {
-
-            }
-
-            @Override
-            public void onRestoreEnd(final TopicPartition topicPartition, final String storeName, final long totalRestored) {
-                restored.addAndGet(totalRestored);
-            }
-        });
+        kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
         kafkaStreams.start();
 
         assertTrue(startupLatch.await(30, TimeUnit.SECONDS));
@@ -386,10 +413,12 @@ public void shouldSuccessfullyStartWhenLoggingDisabled(final boolean stateUpdate
         final StreamsBuilder builder = new StreamsBuilder();
 
         final KStream<Integer, Integer> stream = builder.stream(inputStream);
-        stream.groupByKey()
-                .reduce(
-                        (value1, value2) -> value1 + value2,
-                        Materialized.<Integer, Integer, KeyValueStore<Bytes, byte[]>>as("reduce-store").withLoggingDisabled());
+        stream
+            .groupByKey()
+            .reduce(
+                Integer::sum,
+                Materialized.<Integer, Integer, KeyValueStore<Bytes, byte[]>>as("reduce-store").withLoggingDisabled()
+            );
 
         final CountDownLatch startupLatch = new CountDownLatch(1);
         kafkaStreams = new KafkaStreams(builder.build(), props(stateUpdaterEnabled));
@@ -821,4 +850,30 @@ private void waitForTransitionTo(final Set<KafkaStreams.State> observed, final K
             () -> "Client did not transition to " + state + " on time. Observed transitions: " + observed
         );
     }
+
+    private static class ReadOnlyStoreProcessor implements Processor<Integer, String, Void, Void> {
+        private final AtomicInteger numReceived;
+        private final int offsetLimitDelta;
+        private final CountDownLatch shutdownLatch;
+        KeyValueStore<Integer, String> store;
+
+        public ReadOnlyStoreProcessor(final AtomicInteger numReceived, final int offsetLimitDelta, final CountDownLatch shutdownLatch) {
+            this.numReceived = numReceived;
+            this.offsetLimitDelta = offsetLimitDelta;
+            this.shutdownLatch = shutdownLatch;
+        }
+
+        @Override
+        public void init(final ProcessorContext<Void, Void> context) {
+            store = context.getStateStore("store");
+        }
+
+        @Override
+        public void process(final Record<Integer, String> record) {
+            store.put(record.key(), record.value());
+            if (numReceived.incrementAndGet() == offsetLimitDelta * 2) {
+                shutdownLatch.countDown();
+            }
+        }
+    }
 }
\ No newline at end of file
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
index 4f1d8d3d4266f..7c5734d7b96af 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
@@ -1518,6 +1518,15 @@ public static class TrackingStateRestoreListener implements StateRestoreListener
         public final Map<TopicPartition, AtomicLong> changelogToStartOffset = new ConcurrentHashMap<>();
         public final Map<TopicPartition, AtomicLong> changelogToEndOffset = new ConcurrentHashMap<>();
         public final Map<TopicPartition, AtomicLong> changelogToTotalNumRestored = new ConcurrentHashMap<>();
+        private final AtomicLong restored;
+
+        public TrackingStateRestoreListener() {
+            restored = null;
+        }
+
+        public TrackingStateRestoreListener(final AtomicLong restored) {
+            this.restored = restored;
+        }
 
         @Override
         public void onRestoreStart(final TopicPartition topicPartition,
@@ -1541,6 +1550,9 @@ public void onBatchRestored(final TopicPartition topicPartition,
         public void onRestoreEnd(final TopicPartition topicPartition,
                                  final String storeName,
                                  final long totalRestored) {
+            if (restored != null) {
+                restored.addAndGet(totalRestored);
+            }
         }
 
         public long totalNumRestored() {
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/ReadOnlyStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/ReadOnlyStoreTest.java
new file mode 100644
index 0000000000000..a786c82bf25bd
--- /dev/null
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/ReadOnlyStoreTest.java
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor;
+
+import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.TestInputTopic;
+import org.apache.kafka.streams.TestOutputTopic;
+import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.TopologyTestDriver;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.streams.state.KeyValueStore;
+import org.apache.kafka.streams.state.Stores;
+import org.junit.Test;
+
+import java.util.LinkedList;
+import java.util.List;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+public class ReadOnlyStoreTest {
+
+    @Test
+    public void shouldConnectProcessorAndWriteDataToReadOnlyStore() {
+        final Topology topology = new Topology();
+        topology.addReadOnlyStateStore(
+            Stores.keyValueStoreBuilder(
+                Stores.inMemoryKeyValueStore("readOnlyStore"),
+                new Serdes.IntegerSerde(),
+                new Serdes.StringSerde()
+            ),
+            "readOnlySource",
+            new IntegerDeserializer(),
+            new StringDeserializer(),
+            "storeTopic",
+            "readOnlyProcessor",
+            () -> new Processor<Integer, String, Void, Void>() {
+                KeyValueStore<Integer, String> store;
+
+                @Override
+                public void init(final ProcessorContext<Void, Void> context) {
+                    store = context.getStateStore("readOnlyStore");
+                }
+                @Override
+                public void process(final Record<Integer, String> record) {
+                    store.put(record.key(), record.value());
+                }
+            }
+        );
+
+        topology.addSource("source", new IntegerDeserializer(), new StringDeserializer(), "inputTopic");
+        topology.addProcessor(
+            "processor",
+            () -> new Processor<Integer, String, Integer, String>() {
+                ProcessorContext<Integer, String> context;
+                KeyValueStore<Integer, String> store;
+
+                @Override
+                public void init(final ProcessorContext<Integer, String> context) {
+                    this.context = context;
+                    store = context.getStateStore("readOnlyStore");
+                }
+
+                @Override
+                public void process(final Record<Integer, String> record) {
+                    context.forward(record.withValue(
+                        record.value() + " -- " + store.get(record.key())
+                    ));
+                }
+            },
+            "source"
+        );
+        topology.connectProcessorAndStateStores("processor", "readOnlyStore");
+        topology.addSink("sink", "outputTopic", new IntegerSerializer(), new StringSerializer(), "processor");
+
+        try (final TopologyTestDriver driver = new TopologyTestDriver(topology)) {
+            final TestInputTopic<Integer, String> readOnlyStoreTopic =
+                driver.createInputTopic("storeTopic", new IntegerSerializer(), new StringSerializer());
+            final TestInputTopic<Integer, String> input =
+                driver.createInputTopic("inputTopic", new IntegerSerializer(), new StringSerializer());
+            final TestOutputTopic<Integer, String> output =
+                driver.createOutputTopic("outputTopic", new IntegerDeserializer(), new StringDeserializer());
+
+            readOnlyStoreTopic.pipeInput(1, "foo");
+            readOnlyStoreTopic.pipeInput(2, "bar");
+
+            input.pipeInput(1, "bar");
+            input.pipeInput(2, "foo");
+
+            final KeyValueStore<Integer, String> store = driver.getKeyValueStore("readOnlyStore");
+
+            try (final KeyValueIterator<Integer, String> it = store.all()) {
+                final List<KeyValue<Integer, String>> storeContent = new LinkedList<>();
+                it.forEachRemaining(storeContent::add);
+
+                final List<KeyValue<Integer, String>> expectedResult = new LinkedList<>();
+                expectedResult.add(KeyValue.pair(1, "foo"));
+                expectedResult.add(KeyValue.pair(2, "bar"));
+
+                assertThat(storeContent, equalTo(expectedResult));
+            }
+
+            final List<KeyValue<Integer, String>> expectedResult = new LinkedList<>();
+            expectedResult.add(KeyValue.pair(1, "bar -- foo"));
+            expectedResult.add(KeyValue.pair(2, "foo -- bar"));
+
+            assertThat(output.readKeyValuesToList(), equalTo(expectedResult));
+        }
+    }
+}
diff --git a/streams/src/test/java/org/apache/kafka/test/MockProcessor.java b/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
index 9766d1c0fe85c..8c1812ccf14a2 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockProcessor.java
@@ -32,7 +32,6 @@
 public class MockProcessor<KIn, VIn, KOut, VOut> implements Processor<KIn, VIn, KOut, VOut> {
     private final MockApiProcessor<KIn, VIn, KOut, VOut> delegate;
 
-
     public MockProcessor(final PunctuationType punctuationType,
                          final long scheduleInterval) {
         delegate = new MockApiProcessor<>(punctuationType, scheduleInterval);
@@ -43,12 +42,12 @@ public MockProcessor() {
     }
 
     @Override
-    public void init(ProcessorContext<KOut, VOut> context) {
+    public void init(final ProcessorContext<KOut, VOut> context) {
         delegate.init(context);
     }
 
     @Override
-    public void process(Record<KIn, VIn> record) {
+    public void process(final Record<KIn, VIn> record) {
         delegate.process(record);
     }
 

From 414365979e960d0705f955579836faff6881f46a Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Fri, 8 Mar 2024 22:54:39 +0300
Subject: [PATCH 135/258] KAFKA-14589 [4/4] Tests of ConsoleGroupCommand
 rewritten in java (#15465)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../admin/ConsumerGroupCommandTest.scala      | 228 -------
 .../admin/ResetConsumerGroupOffsetTest.scala  | 514 ---------------
 .../group/ConsumerGroupCommandTest.java       |   5 +-
 .../group/ResetConsumerGroupOffsetTest.java   | 589 ++++++++++++++++++
 4 files changed, 590 insertions(+), 746 deletions(-)
 delete mode 100644 core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
 delete mode 100644 core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
 create mode 100644 tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java

diff --git a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
deleted file mode 100644
index f682df1f1dcdb..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
+++ /dev/null
@@ -1,228 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.admin
-
-import java.time.Duration
-import java.util.concurrent.{ExecutorService, Executors, TimeUnit}
-import java.util.{Collections, Properties, stream}
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
-import kafka.api.BaseConsumerTest
-import kafka.integration.KafkaServerTestHarness
-import kafka.server.KafkaConfig
-import kafka.utils.TestUtils
-import org.apache.kafka.clients.admin.AdminClientConfig
-import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig, GroupProtocol, KafkaConsumer, RangeAssignor}
-import org.apache.kafka.common.{PartitionInfo, TopicPartition}
-import org.apache.kafka.common.errors.WakeupException
-import org.apache.kafka.common.serialization.StringDeserializer
-import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo}
-import org.junit.jupiter.params.provider.Arguments
-
-import scala.jdk.CollectionConverters._
-import scala.collection.mutable.ArrayBuffer
-
-class ConsumerGroupCommandTest extends KafkaServerTestHarness {
-  import ConsumerGroupCommandTest._
-
-  val topic = "foo"
-  val group = "test.group"
-
-  private var consumerGroupService: List[ConsumerGroupService] = List()
-  private var consumerGroupExecutors: List[AbstractConsumerGroupExecutor] = List()
-
-  // configure the servers and clients
-  override def generateConfigs = {
-    val configs = TestUtils.createBrokerConfigs(1, zkConnectOrNull, enableControlledShutdown = false)
-
-    if (isNewGroupCoordinatorEnabled()) {
-      configs.foreach(_.setProperty(KafkaConfig.NewGroupCoordinatorEnableProp, "true"))
-    }
-
-    configs.map(KafkaConfig.fromProps)
-  }
-
-  @BeforeEach
-  override def setUp(testInfo: TestInfo): Unit = {
-    super.setUp(testInfo)
-    createTopic(topic, 1, 1)
-  }
-
-  @AfterEach
-  override def tearDown(): Unit = {
-    consumerGroupService.foreach(_.close())
-    consumerGroupExecutors.foreach(_.shutdown())
-    super.tearDown()
-  }
-
-  def committedOffsets(topic: String = topic, group: String = group): collection.Map[TopicPartition, Long] = {
-    val consumer = createNoAutoCommitConsumer(group)
-    try {
-      val partitions: Set[TopicPartition] = consumer.partitionsFor(topic)
-        .asScala.toSet.map {partitionInfo : PartitionInfo => new TopicPartition(partitionInfo.topic, partitionInfo.partition)}
-      consumer.committed(partitions.asJava).asScala.filter(_._2 != null).map { case (k, v) => k -> v.offset }
-    } finally {
-      consumer.close()
-    }
-  }
-
-  def createNoAutoCommitConsumer(group: String): Consumer[String, String] = {
-    val props = new Properties
-    props.put("bootstrap.servers", bootstrapServers())
-    props.put("group.id", group)
-    props.put("enable.auto.commit", "false")
-    new KafkaConsumer(props, new StringDeserializer, new StringDeserializer)
-  }
-
-  def getConsumerGroupService(args: Array[String]): ConsumerGroupService = {
-    val opts = new ConsumerGroupCommandOptions(args)
-    val service = new ConsumerGroupService(opts, Map(AdminClientConfig.RETRIES_CONFIG -> Int.MaxValue.toString))
-    consumerGroupService = service :: consumerGroupService
-    service
-  }
-
-  def addConsumerGroupExecutor(numConsumers: Int,
-                               topic: String = topic,
-                               group: String = group,
-                               strategy: String = classOf[RangeAssignor].getName,
-                               remoteAssignor: Option[String] = None,
-                               customPropsOpt: Option[Properties] = None,
-                               syncCommit: Boolean = false,
-                               groupProtocol: String = GroupProtocol.CLASSIC.toString): ConsumerGroupExecutor = {
-    val executor = new ConsumerGroupExecutor(bootstrapServers(), numConsumers, group, groupProtocol, topic, strategy, remoteAssignor, customPropsOpt, syncCommit)
-    addExecutor(executor)
-    executor
-  }
-
-  def addSimpleGroupExecutor(partitions: Iterable[TopicPartition] = Seq(new TopicPartition(topic, 0)),
-                             group: String = group): SimpleConsumerGroupExecutor = {
-    val executor = new SimpleConsumerGroupExecutor(bootstrapServers(), group, partitions)
-    addExecutor(executor)
-    executor
-  }
-
-  private def addExecutor(executor: AbstractConsumerGroupExecutor): AbstractConsumerGroupExecutor = {
-    consumerGroupExecutors = executor :: consumerGroupExecutors
-    executor
-  }
-
-}
-
-object ConsumerGroupCommandTest {
-   def getTestQuorumAndGroupProtocolParametersAll(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
-   def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly()
-   def getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly(): stream.Stream[Arguments] = BaseConsumerTest.getTestQuorumAndGroupProtocolParametersConsumerGroupProtocolOnly()
-
-  abstract class AbstractConsumerRunnable(broker: String, groupId: String, customPropsOpt: Option[Properties] = None,
-                                          syncCommit: Boolean = false) extends Runnable {
-    val props = new Properties
-    configure(props)
-    customPropsOpt.foreach(props.asScala ++= _.asScala)
-    val consumer = new KafkaConsumer(props)
-
-    def configure(props: Properties): Unit = {
-      props.put("bootstrap.servers", broker)
-      props.put("group.id", groupId)
-      props.put("key.deserializer", classOf[StringDeserializer].getName)
-      props.put("value.deserializer", classOf[StringDeserializer].getName)
-    }
-
-    def subscribe(): Unit
-
-    def run(): Unit = {
-      try {
-        subscribe()
-        while (true) {
-          consumer.poll(Duration.ofMillis(Long.MaxValue))
-          if (syncCommit)
-            consumer.commitSync()
-        }
-      } catch {
-        case _: WakeupException => // OK
-      } finally {
-        consumer.close()
-      }
-    }
-
-    def shutdown(): Unit = {
-      consumer.wakeup()
-    }
-  }
-
-  class ConsumerRunnable(broker: String, groupId: String, groupProtocol: String, topic: String, strategy: String,
-                         remoteAssignor: Option[String], customPropsOpt: Option[Properties] = None, syncCommit: Boolean = false)
-    extends AbstractConsumerRunnable(broker, groupId, customPropsOpt, syncCommit) {
-
-    override def configure(props: Properties): Unit = {
-      super.configure(props)
-      props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol)
-      if (groupProtocol.toUpperCase == GroupProtocol.CONSUMER.toString) {
-        remoteAssignor.foreach { assignor =>
-          props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, assignor)
-        }
-      } else {
-        props.put("partition.assignment.strategy", strategy)
-      }
-    }
-
-    override def subscribe(): Unit = {
-      consumer.subscribe(Collections.singleton(topic))
-    }
-  }
-
-  class SimpleConsumerRunnable(broker: String, groupId: String, partitions: Iterable[TopicPartition])
-    extends AbstractConsumerRunnable(broker, groupId) {
-
-    override def subscribe(): Unit = {
-      consumer.assign(partitions.toList.asJava)
-    }
-  }
-
-  class AbstractConsumerGroupExecutor(numThreads: Int) {
-    private val executor: ExecutorService = Executors.newFixedThreadPool(numThreads)
-    private val consumers = new ArrayBuffer[AbstractConsumerRunnable]()
-
-    def submit(consumerThread: AbstractConsumerRunnable): Unit = {
-      consumers += consumerThread
-      executor.submit(consumerThread)
-    }
-
-    def shutdown(): Unit = {
-      consumers.foreach(_.shutdown())
-      executor.shutdown()
-      executor.awaitTermination(5000, TimeUnit.MILLISECONDS)
-    }
-  }
-
-  class ConsumerGroupExecutor(broker: String, numConsumers: Int, groupId: String, groupProtocol: String, topic: String, strategy: String,
-                              remoteAssignor: Option[String], customPropsOpt: Option[Properties] = None, syncCommit: Boolean = false)
-    extends AbstractConsumerGroupExecutor(numConsumers) {
-
-    for (_ <- 1 to numConsumers) {
-      submit(new ConsumerRunnable(broker, groupId, groupProtocol, topic, strategy, remoteAssignor, customPropsOpt, syncCommit))
-    }
-
-  }
-
-  class SimpleConsumerGroupExecutor(broker: String, groupId: String, partitions: Iterable[TopicPartition])
-    extends AbstractConsumerGroupExecutor(1) {
-
-    submit(new SimpleConsumerRunnable(broker, groupId, partitions))
-  }
-
-}
-
diff --git a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala b/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
deleted file mode 100644
index 4aafedf0ab319..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
+++ /dev/null
@@ -1,514 +0,0 @@
-/**
-  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
-  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
-  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
-  * License. You may obtain a copy of the License at
-  *
-  * http://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
-  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
-  * specific language governing permissions and limitations under the License.
-  */
-package kafka.admin
-
-import java.io.{BufferedWriter, FileWriter}
-import java.text.{SimpleDateFormat}
-import java.util.{Calendar, Date, Properties}
-
-import joptsimple.OptionException
-import kafka.admin.ConsumerGroupCommand.ConsumerGroupService
-import kafka.server.KafkaConfig
-import kafka.utils.TestUtils
-import org.apache.kafka.clients.producer.ProducerRecord
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.test
-import org.junit.jupiter.api.Assertions._
-import org.junit.jupiter.api.Test
-
-import scala.jdk.CollectionConverters._
-import scala.collection.Seq
-
-
-
-/**
-  * Test cases by:
-  * - Non-existing consumer group
-  * - One for each scenario, with scope=all-topics
-  * - scope=one topic, scenario=to-earliest
-  * - scope=one topic+partitions, scenario=to-earliest
-  * - scope=topics, scenario=to-earliest
-  * - scope=topics+partitions, scenario=to-earliest
-  * - export/import
-  */
-class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
-
-  val overridingProps = new Properties()
-  val topic1 = "foo1"
-  val topic2 = "foo2"
-
-  override def generateConfigs: Seq[KafkaConfig] = {
-    TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false)
-      .map(KafkaConfig.fromProps(_, overridingProps))
-  }
-
-  private def basicArgs: Array[String] = {
-    Array("--reset-offsets",
-      "--bootstrap-server", bootstrapServers(),
-      "--timeout", test.TestUtils.DEFAULT_MAX_WAIT_MS.toString)
-  }
-
-  private def buildArgsForGroups(groups: Seq[String], args: String*): Array[String] = {
-    val groupArgs = groups.flatMap(group => Seq("--group", group)).toArray
-    basicArgs ++ groupArgs ++ args
-  }
-
-  private def buildArgsForGroup(group: String, args: String*): Array[String] = {
-    buildArgsForGroups(Seq(group), args: _*)
-  }
-
-  private def buildArgsForAllGroups(args: String*): Array[String] = {
-    basicArgs ++ Array("--all-groups") ++ args
-  }
-
-  @Test
-  def testResetOffsetsNotExistingGroup(): Unit = {
-    val group = "missing.group"
-    val args = buildArgsForGroup(group, "--all-topics", "--to-current", "--execute")
-    val consumerGroupCommand = getConsumerGroupService(args)
-    // Make sure we got a coordinator
-    TestUtils.waitUntilTrue(() => {
-      consumerGroupCommand.collectGroupState(group).coordinator.host() == "localhost"
-    }, "Can't find a coordinator")
-    val resetOffsets = consumerGroupCommand.resetOffsets()(group)
-    assertEquals(Map.empty, resetOffsets)
-    assertEquals(resetOffsets, committedOffsets(group = group))
-  }
-
-  @Test
-  def testResetOffsetsExistingTopic(): Unit = {
-    val group = "new.group"
-    val args = buildArgsForGroup(group, "--topic", topic, "--to-offset", "50")
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
-  }
-
-  @Test
-  def testResetOffsetsExistingTopicSelectedGroups(): Unit = {
-    produceMessages(topic, 100)
-    val groups =
-      for (id <- 1 to 3) yield {
-        val group = this.group + id
-        val executor = addConsumerGroupExecutor(numConsumers = 1, topic = topic, group = group)
-        awaitConsumerProgress(count = 100L, group = group)
-        executor.shutdown()
-        group
-      }
-    val args = buildArgsForGroups(groups,"--topic", topic, "--to-offset", "50")
-    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
-  }
-
-  @Test
-  def testResetOffsetsExistingTopicAllGroups(): Unit = {
-    val args = buildArgsForAllGroups("--topic", topic, "--to-offset", "50")
-    produceMessages(topic, 100)
-    for (group <- 1 to 3 map (group + _)) {
-      val executor = addConsumerGroupExecutor(numConsumers = 1, topic = topic, group = group)
-      awaitConsumerProgress(count = 100L, group = group)
-      executor.shutdown()
-    }
-    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true)
-    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
-  }
-
-  @Test
-  def testResetOffsetsAllTopicsAllGroups(): Unit = {
-    val args = buildArgsForAllGroups("--all-topics", "--to-offset", "50")
-    val topics = 1 to 3 map (topic + _)
-    val groups = 1 to 3 map (group + _)
-    topics foreach (topic => produceMessages(topic, 100))
-    for {
-      topic <- topics
-      group <- groups
-    } {
-      val executor = addConsumerGroupExecutor(numConsumers = 3, topic = topic, group = group)
-      awaitConsumerProgress(topic = topic, count = 100L, group = group)
-      executor.shutdown()
-    }
-    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true, topics = topics)
-    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true, topics = topics)
-    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50, topics = topics)
-  }
-
-  @Test
-  def testResetOffsetsToLocalDateTime(): Unit = {
-    val format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS")
-    val calendar = Calendar.getInstance()
-    calendar.add(Calendar.DATE, -1)
-
-    produceMessages(topic, 100)
-
-    val executor = addConsumerGroupExecutor(numConsumers = 1, topic)
-    awaitConsumerProgress(count = 100L)
-    executor.shutdown()
-
-    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(calendar.getTime), "--execute")
-    resetAndAssertOffsets(args, expectedOffset = 0)
-  }
-
-  @Test
-  def testResetOffsetsToZonedDateTime(): Unit = {
-    val format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX")
-
-    produceMessages(topic, 50)
-    val checkpoint = new Date()
-    produceMessages(topic, 50)
-
-    val executor = addConsumerGroupExecutor(numConsumers = 1, topic)
-    awaitConsumerProgress(count = 100L)
-    executor.shutdown()
-
-    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute")
-    resetAndAssertOffsets(args, expectedOffset = 50)
-  }
-
-  @Test
-  def testResetOffsetsByDuration(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--by-duration", "PT1M", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    resetAndAssertOffsets(args, expectedOffset = 0)
-  }
-
-  @Test
-  def testResetOffsetsByDurationToEarliest(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--by-duration", "PT0.1S", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    resetAndAssertOffsets(args, expectedOffset = 100)
-  }
-
-  @Test
-  def testResetOffsetsByDurationFallbackToLatestWhenNoRecords(): Unit = {
-    val topic = "foo2"
-    val args = buildArgsForGroup(group, "--topic", topic, "--by-duration", "PT1M", "--execute")
-    createTopic(topic)
-    resetAndAssertOffsets(args, expectedOffset = 0, topics = Seq("foo2"))
-
-    adminZkClient.deleteTopic(topic)
-  }
-
-  @Test
-  def testResetOffsetsToEarliest(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--to-earliest", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    resetAndAssertOffsets(args, expectedOffset = 0)
-  }
-
-  @Test
-  def testResetOffsetsToLatest(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--to-latest", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 200)
-  }
-
-  @Test
-  def testResetOffsetsToCurrentOffset(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--to-current", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 100)
-  }
-
-  @Test
-  def testResetOffsetsToSpecificOffset(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--to-offset", "1", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    resetAndAssertOffsets(args, expectedOffset = 1)
-  }
-
-  @Test
-  def testResetOffsetsShiftPlus(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--shift-by", "50", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 150)
-  }
-
-  @Test
-  def testResetOffsetsShiftMinus(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--shift-by", "-50", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 50)
-  }
-
-  @Test
-  def testResetOffsetsShiftByLowerThanEarliest(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--shift-by", "-150", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 0)
-  }
-
-  @Test
-  def testResetOffsetsShiftByHigherThanLatest(): Unit = {
-    val args = buildArgsForGroup(group, "--all-topics", "--shift-by", "150", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    produceMessages(topic, 100)
-    resetAndAssertOffsets(args, expectedOffset = 200)
-  }
-
-  @Test
-  def testResetOffsetsToEarliestOnOneTopic(): Unit = {
-    val args = buildArgsForGroup(group, "--topic", topic, "--to-earliest", "--execute")
-    produceConsumeAndShutdown(topic, group, totalMessages = 100)
-    resetAndAssertOffsets(args, expectedOffset = 0)
-  }
-
-  @Test
-  def testResetOffsetsToEarliestOnOneTopicAndPartition(): Unit = {
-    val topic = "bar"
-    createTopic(topic, 2, 1)
-
-    val args = buildArgsForGroup(group, "--topic", s"$topic:1", "--to-earliest", "--execute")
-    val consumerGroupCommand = getConsumerGroupService(args)
-
-    produceConsumeAndShutdown(topic, group, totalMessages = 100, numConsumers = 2)
-    val priorCommittedOffsets = committedOffsets(topic = topic)
-
-    val tp0 = new TopicPartition(topic, 0)
-    val tp1 = new TopicPartition(topic, 1)
-    val expectedOffsets = Map(tp0 -> priorCommittedOffsets(tp0), tp1 -> 0L)
-    resetAndAssertOffsetsCommitted(consumerGroupCommand, expectedOffsets, topic)
-
-    adminZkClient.deleteTopic(topic)
-  }
-
-  @Test
-  def testResetOffsetsToEarliestOnTopics(): Unit = {
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-    createTopic(topic1, 1, 1)
-    createTopic(topic2, 1, 1)
-
-    val args = buildArgsForGroup(group, "--topic", topic1, "--topic", topic2, "--to-earliest", "--execute")
-    val consumerGroupCommand = getConsumerGroupService(args)
-
-    produceConsumeAndShutdown(topic1, group, 100, 1)
-    produceConsumeAndShutdown(topic2, group, 100, 1)
-
-    val tp1 = new TopicPartition(topic1, 0)
-    val tp2 = new TopicPartition(topic2, 0)
-
-    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).map { case (k, v) => k -> v.offset }
-    assertEquals(Map(tp1 -> 0L, tp2 -> 0L), allResetOffsets)
-    assertEquals(Map(tp1 -> 0L), committedOffsets(topic1))
-    assertEquals(Map(tp2 -> 0L), committedOffsets(topic2))
-
-    adminZkClient.deleteTopic(topic1)
-    adminZkClient.deleteTopic(topic2)
-  }
-
-  @Test
-  def testResetOffsetsToEarliestOnTopicsAndPartitions(): Unit = {
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-
-    createTopic(topic1, 2, 1)
-    createTopic(topic2, 2, 1)
-
-    val args = buildArgsForGroup(group, "--topic", s"$topic1:1", "--topic", s"$topic2:1", "--to-earliest", "--execute")
-    val consumerGroupCommand = getConsumerGroupService(args)
-
-    produceConsumeAndShutdown(topic1, group, 100, 2)
-    produceConsumeAndShutdown(topic2, group, 100, 2)
-
-    val priorCommittedOffsets1 = committedOffsets(topic1)
-    val priorCommittedOffsets2 = committedOffsets(topic2)
-
-    val tp1 = new TopicPartition(topic1, 1)
-    val tp2 = new TopicPartition(topic2, 1)
-    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).map { case (k, v) => k -> v.offset }
-    assertEquals(Map(tp1 -> 0, tp2 -> 0), allResetOffsets)
-
-    assertEquals(priorCommittedOffsets1.toMap + (tp1 -> 0L), committedOffsets(topic1))
-    assertEquals(priorCommittedOffsets2.toMap + (tp2 -> 0L), committedOffsets(topic2))
-
-    adminZkClient.deleteTopic(topic1)
-    adminZkClient.deleteTopic(topic2)
-  }
-
-  @Test
-  // This one deals with old CSV export/import format for a single --group arg: "topic,partition,offset" to support old behavior
-  def testResetOffsetsExportImportPlanSingleGroupArg(): Unit = {
-    val topic = "bar"
-    val tp0 = new TopicPartition(topic, 0)
-    val tp1 = new TopicPartition(topic, 1)
-    createTopic(topic, 2, 1)
-
-    val cgcArgs = buildArgsForGroup(group, "--all-topics", "--to-offset", "2", "--export")
-    val consumerGroupCommand = getConsumerGroupService(cgcArgs)
-
-    produceConsumeAndShutdown(topic = topic, group = group, totalMessages = 100, numConsumers = 2)
-
-    val file = TestUtils.tempFile("reset", ".csv")
-
-    val exportedOffsets = consumerGroupCommand.resetOffsets()
-    val bw = new BufferedWriter(new FileWriter(file))
-    bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets))
-    bw.close()
-    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), exportedOffsets(group).map { case (k, v) => k -> v.offset })
-
-    val cgcArgsExec = buildArgsForGroup(group, "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
-    val consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec)
-    val importedOffsets = consumerGroupCommandExec.resetOffsets()
-    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), importedOffsets(group).map { case (k, v) => k -> v.offset })
-
-    adminZkClient.deleteTopic(topic)
-  }
-
-  @Test
-  // This one deals with universal CSV export/import file format "group,topic,partition,offset",
-  // supporting multiple --group args or --all-groups arg
-  def testResetOffsetsExportImportPlan(): Unit = {
-    val group1 = group + "1"
-    val group2 = group + "2"
-    val topic1 = "bar1"
-    val topic2 = "bar2"
-    val t1p0 = new TopicPartition(topic1, 0)
-    val t1p1 = new TopicPartition(topic1, 1)
-    val t2p0 = new TopicPartition(topic2, 0)
-    val t2p1 = new TopicPartition(topic2, 1)
-    createTopic(topic1, 2, 1)
-    createTopic(topic2, 2, 1)
-
-    val cgcArgs = buildArgsForGroups(Seq(group1, group2), "--all-topics", "--to-offset", "2", "--export")
-    val consumerGroupCommand = getConsumerGroupService(cgcArgs)
-
-    produceConsumeAndShutdown(topic = topic1, group = group1, totalMessages = 100)
-    produceConsumeAndShutdown(topic = topic2, group = group2, totalMessages = 100)
-
-    awaitConsumerGroupInactive(consumerGroupCommand, group1)
-    awaitConsumerGroupInactive(consumerGroupCommand, group2)
-
-    val file = TestUtils.tempFile("reset", ".csv")
-
-    val exportedOffsets = consumerGroupCommand.resetOffsets()
-    val bw = new BufferedWriter(new FileWriter(file))
-    bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets))
-    bw.close()
-    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), exportedOffsets(group1).map { case (k, v) => k -> v.offset })
-    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), exportedOffsets(group2).map { case (k, v) => k -> v.offset })
-
-    // Multiple --group's offset import
-    val cgcArgsExec = buildArgsForGroups(Seq(group1, group2), "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
-    val consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec)
-    val importedOffsets = consumerGroupCommandExec.resetOffsets()
-    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), importedOffsets(group1).map { case (k, v) => k -> v.offset })
-    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), importedOffsets(group2).map { case (k, v) => k -> v.offset })
-
-    // Single --group offset import using "group,topic,partition,offset" csv format
-    val cgcArgsExec2 = buildArgsForGroup(group1, "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
-    val consumerGroupCommandExec2 = getConsumerGroupService(cgcArgsExec2)
-    val importedOffsets2 = consumerGroupCommandExec2.resetOffsets()
-    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), importedOffsets2(group1).map { case (k, v) => k -> v.offset })
-
-    adminZkClient.deleteTopic(topic)
-  }
-
-  @Test
-  def testResetWithUnrecognizedNewConsumerOption(): Unit = {
-    val cgcArgs = Array("--new-consumer", "--bootstrap-server", bootstrapServers(), "--reset-offsets", "--group", group, "--all-topics",
-      "--to-offset", "2", "--export")
-    assertThrows(classOf[OptionException], () => getConsumerGroupService(cgcArgs))
-  }
-
-  private def produceMessages(topic: String, numMessages: Int): Unit = {
-    val records = (0 until numMessages).map(_ => new ProducerRecord[Array[Byte], Array[Byte]](topic,
-      new Array[Byte](100 * 1000)))
-    TestUtils.produceMessages(servers, records, acks = 1)
-  }
-
-  private def produceConsumeAndShutdown(topic: String, group: String, totalMessages: Int, numConsumers: Int = 1): Unit = {
-    produceMessages(topic, totalMessages)
-    val executor = addConsumerGroupExecutor(numConsumers = numConsumers, topic = topic, group = group)
-    awaitConsumerProgress(topic, group, totalMessages)
-    executor.shutdown()
-  }
-
-  private def awaitConsumerProgress(topic: String = topic,
-                                    group: String = group,
-                                    count: Long): Unit = {
-    val consumer = createNoAutoCommitConsumer(group)
-    try {
-      val partitions = consumer.partitionsFor(topic).asScala.map { partitionInfo =>
-        new TopicPartition(partitionInfo.topic, partitionInfo.partition)
-      }.toSet
-
-      TestUtils.waitUntilTrue(() => {
-        val committed = consumer.committed(partitions.asJava).values.asScala
-        val total = committed.foldLeft(0L) { case (currentSum, offsetAndMetadata) =>
-          currentSum + Option(offsetAndMetadata).map(_.offset).getOrElse(0L)
-        }
-        total == count
-      }, "Expected that consumer group has consumed all messages from topic/partition. " +
-        s"Expected offset: $count. Actual offset: ${committedOffsets(topic, group).values.sum}")
-
-    } finally {
-      consumer.close()
-    }
-
-  }
-
-  private def awaitConsumerGroupInactive(consumerGroupService: ConsumerGroupService, group: String): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val state = consumerGroupService.collectGroupState(group).state
-      state == "Empty" || state == "Dead"
-    }, s"Expected that consumer group is inactive. Actual state: ${consumerGroupService.collectGroupState(group).state}")
-  }
-
-  private def resetAndAssertOffsets(args: Array[String],
-                                    expectedOffset: Long,
-                                    dryRun: Boolean = false,
-                                    topics: Seq[String] = Seq(topic)): Unit = {
-    val consumerGroupCommand = getConsumerGroupService(args)
-    val expectedOffsets = topics.map(topic => topic -> Map(new TopicPartition(topic, 0) -> expectedOffset)).toMap
-    val resetOffsetsResultByGroup = resetOffsets(consumerGroupCommand)
-
-    try {
-      for {
-        topic <- topics
-        (group, partitionInfo) <- resetOffsetsResultByGroup
-      } {
-        val priorOffsets = committedOffsets(topic = topic, group = group)
-        assertEquals(expectedOffsets(topic),
-          partitionInfo.filter(partitionInfo => partitionInfo._1.topic() == topic).map { case (k, v) => k -> v.offset })
-        assertEquals(if (dryRun) priorOffsets else expectedOffsets(topic), committedOffsets(topic = topic, group = group))
-      }
-    } finally {
-      consumerGroupCommand.close()
-    }
-  }
-
-  private def resetAndAssertOffsetsCommitted(consumerGroupService: ConsumerGroupService,
-                                             expectedOffsets: Map[TopicPartition, Long],
-                                             topic: String): Unit = {
-    val allResetOffsets = resetOffsets(consumerGroupService)
-    for {
-      (group, offsetsInfo) <- allResetOffsets
-      (tp, offsetMetadata) <- offsetsInfo
-    } {
-      assertEquals(offsetMetadata.offset(), expectedOffsets(tp))
-      assertEquals(expectedOffsets, committedOffsets(topic, group))
-    }
-  }
-
-  private def resetOffsets(consumerGroupService: ConsumerGroupService) = {
-    consumerGroupService.resetOffsets()
-  }
-}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
index bde3af37a1d70..e3ee39f6b1d04 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
@@ -87,10 +87,7 @@ public Seq<kafka.server.KafkaConfig> generateConfigs() {
             0,
             false
         ).foreach(props -> {
-            if (isNewGroupCoordinatorEnabled()) {
-                props.setProperty(KafkaConfig.NewGroupCoordinatorEnableProp(), "true");
-            }
-
+            props.setProperty(KafkaConfig.NewGroupCoordinatorEnableProp(), isNewGroupCoordinatorEnabled() + "");
             cfgs.add(KafkaConfig.fromProps(props));
             return null;
         });
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java
new file mode 100644
index 0000000000000..dca263452231f
--- /dev/null
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java
@@ -0,0 +1,589 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import joptsimple.OptionException;
+import kafka.admin.ConsumerGroupCommand;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.GroupProtocol;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.Test;
+import scala.Option;
+
+import java.io.BufferedWriter;
+import java.io.File;
+import java.io.FileWriter;
+import java.text.SimpleDateFormat;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Calendar;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Test cases by:
+ * - Non-existing consumer group
+ * - One for each scenario, with scope=all-topics
+ * - scope=one topic, scenario=to-earliest
+ * - scope=one topic+partitions, scenario=to-earliest
+ * - scope=topics, scenario=to-earliest
+ * - scope=topics+partitions, scenario=to-earliest
+ * - export/import
+ */
+public class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
+    private String[] basicArgs() {
+        return new String[]{"--reset-offsets",
+            "--bootstrap-server", bootstrapServers(listenerName()),
+            "--timeout", Long.toString(DEFAULT_MAX_WAIT_MS)};
+    }
+
+    private String[] buildArgsForGroups(List<String> groups, String...args) {
+        List<String> res = new ArrayList<>(Arrays.asList(basicArgs()));
+        for (String group : groups) {
+            res.add("--group");
+            res.add(group);
+        }
+        res.addAll(Arrays.asList(args));
+        return res.toArray(new String[0]);
+    }
+
+    private String[] buildArgsForGroup(String group, String...args) {
+        return buildArgsForGroups(Collections.singletonList(group), args);
+    }
+
+    private String[] buildArgsForAllGroups(String...args) {
+        List<String> res = new ArrayList<>(Arrays.asList(basicArgs()));
+        res.add("--all-groups");
+        res.addAll(Arrays.asList(args));
+        return res.toArray(new String[0]);
+    }
+
+    @Test
+    public void testResetOffsetsNotExistingGroup() throws Exception {
+        String group = "missing.group";
+        String[] args = buildArgsForGroup(group, "--all-topics", "--to-current", "--execute");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
+        // Make sure we got a coordinator
+        TestUtils.waitForCondition(
+            () -> Objects.equals(consumerGroupCommand.collectGroupState(group).coordinator().host(), "localhost"),
+            "Can't find a coordinator");
+        Option<scala.collection.Map<TopicPartition, OffsetAndMetadata>> resetOffsets = consumerGroupCommand.resetOffsets().get(group);
+        assertTrue(resetOffsets.isDefined() && resetOffsets.get().isEmpty());
+        assertTrue(committedOffsets(TOPIC, group).isEmpty());
+    }
+
+    @Test
+    public void testResetOffsetsExistingTopic() {
+        String group = "new.group";
+        String[] args = buildArgsForGroup(group, "--topic", TOPIC, "--to-offset", "50");
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--dry-run"), 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--execute"), 50, false, Collections.singletonList(TOPIC));
+    }
+
+    @Test
+    public void testResetOffsetsExistingTopicSelectedGroups() throws Exception {
+        produceMessages(TOPIC, 100);
+        List<String> groups = IntStream.rangeClosed(1, 3).mapToObj(id -> GROUP + id).collect(Collectors.toList());
+        for (String group : groups) {
+            ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, group, GroupProtocol.CLASSIC.name);
+            awaitConsumerProgress(TOPIC, group, 100L);
+            executor.shutdown();
+        }
+        String[] args = buildArgsForGroups(groups, "--topic", TOPIC, "--to-offset", "50");
+        resetAndAssertOffsets(args, 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--dry-run"), 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--execute"), 50, false, Collections.singletonList(TOPIC));
+    }
+
+    @Test
+    public void testResetOffsetsExistingTopicAllGroups() throws Exception {
+        String[] args = buildArgsForAllGroups("--topic", TOPIC, "--to-offset", "50");
+        produceMessages(TOPIC, 100);
+        for (int i = 1; i <= 3; i++) {
+            String group = GROUP + i;
+            ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, group, GroupProtocol.CLASSIC.name);
+            awaitConsumerProgress(TOPIC, group, 100L);
+            executor.shutdown();
+        }
+        resetAndAssertOffsets(args, 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--dry-run"), 50, true, Collections.singletonList(TOPIC));
+        resetAndAssertOffsets(addTo(args, "--execute"), 50, false, Collections.singletonList(TOPIC));
+    }
+
+    @Test
+    public void testResetOffsetsAllTopicsAllGroups() throws Exception {
+        String[] args = buildArgsForAllGroups("--all-topics", "--to-offset", "50");
+        List<String> topics = IntStream.rangeClosed(1, 3).mapToObj(i -> TOPIC + i).collect(Collectors.toList());
+        List<String> groups = IntStream.rangeClosed(1, 3).mapToObj(i -> GROUP + i).collect(Collectors.toList());
+        topics.forEach(topic -> produceMessages(topic, 100));
+
+        for (String topic : topics) {
+            for (String group : groups) {
+                ConsumerGroupExecutor executor = addConsumerGroupExecutor(3, topic, group, GroupProtocol.CLASSIC.name);
+                awaitConsumerProgress(topic, group, 100);
+                executor.shutdown();
+            }
+        }
+        resetAndAssertOffsets(args, 50, true, topics);
+        resetAndAssertOffsets(addTo(args, "--dry-run"), 50, true, topics);
+        resetAndAssertOffsets(addTo(args, "--execute"), 50, false, topics);
+    }
+
+    @Test
+    public void testResetOffsetsToLocalDateTime() throws Exception {
+        SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS");
+        Calendar calendar = Calendar.getInstance();
+        calendar.add(Calendar.DATE, -1);
+
+        produceMessages(TOPIC, 100);
+
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, GROUP, GroupProtocol.CLASSIC.name);
+        awaitConsumerProgress(TOPIC, GROUP, 100L);
+        executor.shutdown();
+
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-datetime", format.format(calendar.getTime()), "--execute");
+        resetAndAssertOffsets(args, 0);
+    }
+
+    @Test
+    public void testResetOffsetsToZonedDateTime() throws Exception {
+        SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
+
+        produceMessages(TOPIC, 50);
+        Date checkpoint = new Date();
+        produceMessages(TOPIC, 50);
+
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(1, TOPIC, GROUP, GroupProtocol.CLASSIC.name);
+        awaitConsumerProgress(TOPIC, GROUP, 100L);
+        executor.shutdown();
+
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute");
+        resetAndAssertOffsets(args, 50);
+    }
+
+    @Test
+    public void testResetOffsetsByDuration() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--by-duration", "PT1M", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        resetAndAssertOffsets(args, 0);
+    }
+
+    @Test
+    public void testResetOffsetsByDurationToEarliest() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--by-duration", "PT0.1S", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        resetAndAssertOffsets(args, 100);
+    }
+
+    @Test
+    public void testResetOffsetsByDurationFallbackToLatestWhenNoRecords() throws Exception {
+        String topic = "foo2";
+        String[] args = buildArgsForGroup(GROUP, "--topic", topic, "--by-duration", "PT1M", "--execute");
+        createTopic(topic, 1, 1, new Properties(), listenerName(), new Properties());
+        resetAndAssertOffsets(args, 0, false, Collections.singletonList("foo2"));
+
+        adminZkClient().deleteTopic(topic);
+    }
+
+    @Test
+    public void testResetOffsetsToEarliest() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-earliest", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        resetAndAssertOffsets(args, 0);
+    }
+
+    @Test
+    public void testResetOffsetsToLatest() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-latest", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 200);
+    }
+
+    @Test
+    public void testResetOffsetsToCurrentOffset() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-current", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 100);
+    }
+
+    @Test
+    public void testResetOffsetsToSpecificOffset() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--to-offset", "1", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        resetAndAssertOffsets(args, 1);
+    }
+
+    @Test
+    public void testResetOffsetsShiftPlus() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--shift-by", "50", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 150);
+    }
+
+    @Test
+    public void testResetOffsetsShiftMinus() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--shift-by", "-50", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 50);
+    }
+
+    @Test
+    public void testResetOffsetsShiftByLowerThanEarliest() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--shift-by", "-150", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 0);
+    }
+
+    @Test
+    public void testResetOffsetsShiftByHigherThanLatest() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--all-topics", "--shift-by", "150", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        produceMessages(TOPIC, 100);
+        resetAndAssertOffsets(args, 200);
+    }
+
+    @Test
+    public void testResetOffsetsToEarliestOnOneTopic() throws Exception {
+        String[] args = buildArgsForGroup(GROUP, "--topic", TOPIC, "--to-earliest", "--execute");
+        produceConsumeAndShutdown(TOPIC, GROUP, 100, 1);
+        resetAndAssertOffsets(args, 0);
+    }
+
+    @Test
+    public void testResetOffsetsToEarliestOnOneTopicAndPartition() throws Exception {
+        String topic = "bar";
+        createTopic(topic, 2, 1, new Properties(), listenerName(), new Properties());
+
+        String[] args = buildArgsForGroup(GROUP, "--topic", topic + ":1", "--to-earliest", "--execute");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
+
+        produceConsumeAndShutdown(topic, GROUP, 100, 2);
+        Map<TopicPartition, Long> priorCommittedOffsets = committedOffsets(topic, GROUP);
+
+        TopicPartition tp0 = new TopicPartition(topic, 0);
+        TopicPartition tp1 = new TopicPartition(topic, 1);
+        Map<TopicPartition, Long> expectedOffsets = new HashMap<>();
+        expectedOffsets.put(tp0, priorCommittedOffsets.get(tp0));
+        expectedOffsets.put(tp1, 0L);
+        resetAndAssertOffsetsCommitted(consumerGroupCommand, expectedOffsets, topic);
+
+        adminZkClient().deleteTopic(topic);
+    }
+
+    @Test
+    public void testResetOffsetsToEarliestOnTopics() throws Exception {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        createTopic(topic1, 1, 1, new Properties(), listenerName(), new Properties());
+        createTopic(topic2, 1, 1, new Properties(), listenerName(), new Properties());
+
+        String[] args = buildArgsForGroup(GROUP, "--topic", topic1, "--topic", topic2, "--to-earliest", "--execute");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
+
+        produceConsumeAndShutdown(topic1, GROUP, 100, 1);
+        produceConsumeAndShutdown(topic2, GROUP, 100, 1);
+
+        TopicPartition tp1 = new TopicPartition(topic1, 0);
+        TopicPartition tp2 = new TopicPartition(topic2, 0);
+
+        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(consumerGroupCommand.resetOffsets().get(GROUP));
+        Map<TopicPartition, Long> expMap = new HashMap<>();
+        expMap.put(tp1, 0L);
+        expMap.put(tp2, 0L);
+        assertEquals(expMap, allResetOffsets);
+        assertEquals(Collections.singletonMap(tp1, 0L), committedOffsets(topic1, GROUP));
+        assertEquals(Collections.singletonMap(tp2, 0L), committedOffsets(topic2, GROUP));
+
+        adminZkClient().deleteTopic(topic1);
+        adminZkClient().deleteTopic(topic2);
+    }
+
+    @Test
+    public void testResetOffsetsToEarliestOnTopicsAndPartitions() throws Exception {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+
+        createTopic(topic1, 2, 1, new Properties(), listenerName(), new Properties());
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        String[] args = buildArgsForGroup(GROUP, "--topic", topic1 + ":1", "--topic", topic2 + ":1", "--to-earliest", "--execute");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
+
+        produceConsumeAndShutdown(topic1, GROUP, 100, 2);
+        produceConsumeAndShutdown(topic2, GROUP, 100, 2);
+
+        Map<TopicPartition, Long> priorCommittedOffsets1 = committedOffsets(topic1, GROUP);
+        Map<TopicPartition, Long> priorCommittedOffsets2 = committedOffsets(topic2, GROUP);
+
+        TopicPartition tp1 = new TopicPartition(topic1, 1);
+        TopicPartition tp2 = new TopicPartition(topic2, 1);
+        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(consumerGroupCommand.resetOffsets().get(GROUP));
+        Map<TopicPartition, Long> expMap = new HashMap<>();
+        expMap.put(tp1, 0L);
+        expMap.put(tp2, 0L);
+        assertEquals(expMap, allResetOffsets);
+        priorCommittedOffsets1.put(tp1, 0L);
+        assertEquals(priorCommittedOffsets1, committedOffsets(topic1, GROUP));
+        priorCommittedOffsets2.put(tp2, 0L);
+        assertEquals(priorCommittedOffsets2, committedOffsets(topic2, GROUP));
+
+        adminZkClient().deleteTopic(topic1);
+        adminZkClient().deleteTopic(topic2);
+    }
+
+    @Test
+    // This one deals with old CSV export/import format for a single --group arg: "topic,partition,offset" to support old behavior
+    public void testResetOffsetsExportImportPlanSingleGroupArg() throws Exception {
+        String topic = "bar";
+        TopicPartition tp0 = new TopicPartition(topic, 0);
+        TopicPartition tp1 = new TopicPartition(topic, 1);
+        createTopic(topic, 2, 1, new Properties(), listenerName(), new Properties());
+
+        String[] cgcArgs = buildArgsForGroup(GROUP, "--all-topics", "--to-offset", "2", "--export");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(cgcArgs);
+
+        produceConsumeAndShutdown(topic, GROUP, 100, 2);
+
+        File file = TestUtils.tempFile("reset", ".csv");
+
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
+        BufferedWriter bw = new BufferedWriter(new FileWriter(file));
+        bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets));
+        bw.close();
+
+        Map<TopicPartition, Long> exp1 = new HashMap<>();
+        exp1.put(tp0, 2L);
+        exp1.put(tp1, 2L);
+        assertEquals(exp1, toOffsetMap(exportedOffsets.get(GROUP)));
+
+        String[] cgcArgsExec = buildArgsForGroup(GROUP, "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec);
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
+        assertEquals(exp1, toOffsetMap(importedOffsets.get(GROUP)));
+
+        adminZkClient().deleteTopic(topic);
+    }
+
+    @Test
+    // This one deals with universal CSV export/import file format "group,topic,partition,offset",
+    // supporting multiple --group args or --all-groups arg
+    public void testResetOffsetsExportImportPlan() throws Exception {
+        String group1 = GROUP + "1";
+        String group2 = GROUP + "2";
+        String topic1 = "bar1";
+        String topic2 = "bar2";
+        TopicPartition t1p0 = new TopicPartition(topic1, 0);
+        TopicPartition t1p1 = new TopicPartition(topic1, 1);
+        TopicPartition t2p0 = new TopicPartition(topic2, 0);
+        TopicPartition t2p1 = new TopicPartition(topic2, 1);
+        createTopic(topic1, 2, 1, new Properties(), listenerName(), new Properties());
+        createTopic(topic2, 2, 1, new Properties(), listenerName(), new Properties());
+
+        String[] cgcArgs = buildArgsForGroups(Arrays.asList(group1, group2), "--all-topics", "--to-offset", "2", "--export");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(cgcArgs);
+
+        produceConsumeAndShutdown(topic1, group1, 100, 1);
+        produceConsumeAndShutdown(topic2, group2, 100, 1);
+
+        awaitConsumerGroupInactive(consumerGroupCommand, group1);
+        awaitConsumerGroupInactive(consumerGroupCommand, group2);
+
+        File file = TestUtils.tempFile("reset", ".csv");
+
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
+        BufferedWriter bw = new BufferedWriter(new FileWriter(file));
+        bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets));
+        bw.close();
+        Map<TopicPartition, Long> exp1 = new HashMap<>();
+        exp1.put(t1p0, 2L);
+        exp1.put(t1p1, 2L);
+        Map<TopicPartition, Long> exp2 = new HashMap<>();
+        exp2.put(t2p0, 2L);
+        exp2.put(t2p1, 2L);
+
+        assertEquals(exp1, toOffsetMap(exportedOffsets.get(group1)));
+        assertEquals(exp2, toOffsetMap(exportedOffsets.get(group2)));
+
+        // Multiple --group's offset import
+        String[] cgcArgsExec = buildArgsForGroups(Arrays.asList(group1, group2), "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec);
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
+        assertEquals(exp1, toOffsetMap(importedOffsets.get(group1)));
+        assertEquals(exp2, toOffsetMap(importedOffsets.get(group2)));
+
+        // Single --group offset import using "group,topic,partition,offset" csv format
+        String[] cgcArgsExec2 = buildArgsForGroup(group1, "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec2 = getConsumerGroupService(cgcArgsExec2);
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets2 = consumerGroupCommandExec2.resetOffsets();
+        assertEquals(exp1, toOffsetMap(importedOffsets2.get(group1)));
+
+        adminZkClient().deleteTopic(TOPIC);
+    }
+
+    @Test
+    public void testResetWithUnrecognizedNewConsumerOption() {
+        String[] cgcArgs = new String[]{"--new-consumer", "--bootstrap-server", bootstrapServers(listenerName()), "--reset-offsets",
+            "--group", GROUP, "--all-topics", "--to-offset", "2", "--export"};
+        assertThrows(OptionException.class, () -> getConsumerGroupService(cgcArgs));
+    }
+
+    private void produceMessages(String topic, int numMessages) {
+        List<ProducerRecord<byte[], byte[]>> records = IntStream.range(0, numMessages)
+            .mapToObj(i -> new ProducerRecord<byte[], byte[]>(topic, new byte[100 * 1000]))
+            .collect(Collectors.toList());
+        kafka.utils.TestUtils.produceMessages(servers(), seq(records), 1);
+    }
+
+    private void produceConsumeAndShutdown(String topic, String group, int totalMessages, int numConsumers) throws Exception {
+        produceMessages(topic, totalMessages);
+        ConsumerGroupExecutor executor = addConsumerGroupExecutor(numConsumers, topic, group, GroupProtocol.CLASSIC.name);
+        awaitConsumerProgress(topic, group, totalMessages);
+        executor.shutdown();
+    }
+
+    private void awaitConsumerProgress(String topic,
+                                       String group,
+                                       long count) throws Exception {
+        try (Consumer<String, String> consumer = createNoAutoCommitConsumer(group)) {
+            Set<TopicPartition> partitions = consumer.partitionsFor(topic).stream()
+                .map(partitionInfo -> new TopicPartition(partitionInfo.topic(), partitionInfo.partition()))
+                .collect(Collectors.toSet());
+
+            TestUtils.waitForCondition(() -> {
+                Collection<OffsetAndMetadata> committed = consumer.committed(partitions).values();
+                long total = committed.stream()
+                    .mapToLong(offsetAndMetadata -> Optional.ofNullable(offsetAndMetadata).map(OffsetAndMetadata::offset).orElse(0L))
+                    .sum();
+
+                return total == count;
+            }, "Expected that consumer group has consumed all messages from topic/partition. " +
+                "Expected offset: " + count + ". Actual offset: " + committedOffsets(topic, group).values().stream().mapToLong(Long::longValue).sum());
+        }
+    }
+
+    private void awaitConsumerGroupInactive(ConsumerGroupCommand.ConsumerGroupService consumerGroupService, String group) throws Exception {
+        TestUtils.waitForCondition(() -> {
+            String state = consumerGroupService.collectGroupState(group).state();
+            return Objects.equals(state, "Empty") || Objects.equals(state, "Dead");
+        }, "Expected that consumer group is inactive. Actual state: " + consumerGroupService.collectGroupState(group).state());
+    }
+
+    private void resetAndAssertOffsets(String[] args,
+                                       long expectedOffset) {
+        resetAndAssertOffsets(args, expectedOffset, false, Collections.singletonList(TOPIC));
+    }
+
+    private void resetAndAssertOffsets(String[] args,
+                                      long expectedOffset,
+                                      boolean dryRun,
+                                      List<String> topics) {
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
+        Map<String, Map<TopicPartition, Long>> expectedOffsets = topics.stream().collect(Collectors.toMap(
+            Function.identity(),
+            topic -> Collections.singletonMap(new TopicPartition(topic, 0), expectedOffset)));
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> resetOffsetsResultByGroup = consumerGroupCommand.resetOffsets();
+
+        try {
+            for (final String topic : topics) {
+                resetOffsetsResultByGroup.foreach(entry -> {
+                    String group = entry._1;
+                    scala.collection.Map<TopicPartition, OffsetAndMetadata> partitionInfo = entry._2;
+                    Map<TopicPartition, Long> priorOffsets = committedOffsets(topic, group);
+                    Map<TopicPartition, Long> offsets = new HashMap<>();
+                    partitionInfo.foreach(partitionInfoEntry -> {
+                        TopicPartition tp = partitionInfoEntry._1;
+                        OffsetAndMetadata offsetAndMetadata = partitionInfoEntry._2;
+                        if (Objects.equals(tp.topic(), topic))
+                            offsets.put(tp, offsetAndMetadata.offset());
+                        return null;
+                    });
+                    assertEquals(expectedOffsets.get(topic), offsets);
+                    assertEquals(dryRun ? priorOffsets : expectedOffsets.get(topic), committedOffsets(topic, group));
+                    return null;
+                });
+            }
+        } finally {
+            consumerGroupCommand.close();
+        }
+    }
+
+    private void resetAndAssertOffsetsCommitted(ConsumerGroupCommand.ConsumerGroupService consumerGroupService,
+                                                Map<TopicPartition, Long> expectedOffsets,
+                                                String topic) {
+        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> allResetOffsets = consumerGroupService.resetOffsets();
+
+        allResetOffsets.foreach(entry -> {
+            String group = entry._1;
+            scala.collection.Map<TopicPartition, OffsetAndMetadata> offsetsInfo = entry._2;
+            offsetsInfo.foreach(offsetInfoEntry -> {
+                TopicPartition tp = offsetInfoEntry._1;
+                OffsetAndMetadata offsetMetadata = offsetInfoEntry._2;
+                assertEquals(offsetMetadata.offset(), expectedOffsets.get(tp));
+                assertEquals(expectedOffsets, committedOffsets(topic, group));
+                return null;
+            });
+            return null;
+        });
+    }
+
+    Map<TopicPartition, Long> toOffsetMap(Option<scala.collection.Map<TopicPartition, OffsetAndMetadata>> map) {
+        assertTrue(map.isDefined());
+        Map<TopicPartition, Long> res = new HashMap<>();
+        map.foreach(m -> {
+            m.foreach(entry -> {
+                TopicPartition tp = entry._1;
+                OffsetAndMetadata offsetAndMetadata = entry._2;
+                res.put(tp, offsetAndMetadata.offset());
+                return null;
+            });
+            return null;
+        });
+        return res;
+    }
+
+    private String[] addTo(String[] args, String...extra) {
+        List<String> res = new ArrayList<>(Arrays.asList(args));
+        res.addAll(Arrays.asList(extra));
+        return res.toArray(new String[0]);
+    }
+}

From 861fe68cee5ae052531a34a7c40196d2ffdda055 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Fri, 8 Mar 2024 13:51:32 -0800
Subject: [PATCH 136/258] TRIVIAL: fix typo

---
 .../src/main/java/org/apache/kafka/streams/StreamsConfig.java   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 54b30fea4d5d2..9616ed81aac32 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -519,7 +519,7 @@ public class StreamsConfig extends AbstractConfig {
     /** {@code enable.metrics.push} */
     @SuppressWarnings("WeakerAccess")
     public static  final String ENABLE_METRICS_PUSH_CONFIG = CommonClientConfigs.ENABLE_METRICS_PUSH_CONFIG;
-    public static final String ENABLE_METRICS_PUSH_DOC = "Whether to enable pushing of internal [main-|restore-|global]consumer, producer, and admin client metrics to the cluster, if the cluster has a client metrics subscription which matches a client.";
+    public static final String ENABLE_METRICS_PUSH_DOC = "Whether to enable pushing of internal [main-|restore-|global-]consumer, producer, and admin client metrics to the cluster, if the cluster has a client metrics subscription which matches a client.";
 
     /** {@code commit.interval.ms} */
     @SuppressWarnings("WeakerAccess")

From 2fcafbd497f1a49bcb9915147c2c7f2de05e5f59 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Fri, 8 Mar 2024 14:28:39 -0800
Subject: [PATCH 137/258] HOTFIX: fix html markup

---
 .../src/main/java/org/apache/kafka/streams/StreamsConfig.java   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 9616ed81aac32..c166b3a43f04f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -514,7 +514,7 @@ public class StreamsConfig extends AbstractConfig {
     @SuppressWarnings("WeakerAccess")
     public static final String CLIENT_ID_CONFIG = CommonClientConfigs.CLIENT_ID_CONFIG;
     private static final String CLIENT_ID_DOC = "An ID prefix string used for the client IDs of internal [main-|restore-|global-]consumer, producer, and admin clients" +
-        " with pattern <code>&lt;client.id&gt;-[Global]StreamThread[-&lt;threadSequenceNumber$gt;]-&lt;consumer|producer|restore-consumer|global-consumer&gt;</code>.";
+        " with pattern <code>&lt;client.id&gt;-[Global]StreamThread[-&lt;threadSequenceNumber&gt;]-&lt;consumer|producer|restore-consumer|global-consumer&gt;</code>.";
 
     /** {@code enable.metrics.push} */
     @SuppressWarnings("WeakerAccess")

From 432825df95fe9824da05cf86cd40b7741210deb2 Mon Sep 17 00:00:00 2001
From: Joel Hamill <11722533+joel-hamill@users.noreply.github.com>
Date: Fri, 8 Mar 2024 15:15:43 -0800
Subject: [PATCH 138/258] MINOR: Fix incorrect syntax for config (#15500)

Fix incorrect syntax for config.

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../org/apache/kafka/clients/producer/ProducerConfig.java   | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java
index d74b5e4c1dd41..9471b48aa42ec 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java
@@ -294,8 +294,10 @@ public class ProducerConfig extends AbstractConfig {
             "<ul>" +
             "<li>If not set, the default partitioning logic is used. " + 
             "This strategy send records to a partition until at least " + BATCH_SIZE_CONFIG + " bytes is produced to the partition. It works with the strategy:" + 
-            "<p> 1) If no partition is specified but a key is present, choose a partition based on a hash of the key." +
-            "<p> 2) If no partition or key is present, choose the sticky partition that changes when at least " + BATCH_SIZE_CONFIG + " bytes are produced to the partition." +
+            "<ol>" +
+            "<li>If no partition is specified but a key is present, choose a partition based on a hash of the key.</li>" +
+            "<li>If no partition or key is present, choose the sticky partition that changes when at least " + BATCH_SIZE_CONFIG + " bytes are produced to the partition.</li>" +
+            "</ol>" +
             "</li>" +
             "<li><code>org.apache.kafka.clients.producer.RoundRobinPartitioner</code>: A partitioning strategy where " +
             "each record in a series of consecutive records is sent to a different partition, regardless of whether the 'key' is provided or not, " +

From 3fcaa9ccc0e36d05626e6c122c67258773bb6a9b Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Sun, 10 Mar 2024 03:06:41 +0800
Subject: [PATCH 139/258] MINOR: remove the copy constructor of LogSegment
 (#15488)

In the LogSegment, the copy constructor is only used in LogLoaderTest

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../scala/unit/kafka/log/LogLoaderTest.scala  | 26 ++++++++-----------
 .../storage/internals/log/LogSegment.java     |  6 -----
 2 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala b/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
index 5dcca54961bb8..c32e59e71569b 100644
--- a/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogLoaderTest.scala
@@ -34,14 +34,14 @@ import org.apache.kafka.server.common.MetadataVersion.IBP_0_11_0_IV0
 import org.apache.kafka.server.config.Defaults
 import org.apache.kafka.server.util.{MockTime, Scheduler}
 import org.apache.kafka.storage.internals.epoch.LeaderEpochFileCache
-import org.apache.kafka.storage.internals.log.{AbortedTxn, CleanerConfig, EpochEntry, FetchDataInfo, LogConfig, LogDirFailureChannel, LogFileUtils, LogOffsetMetadata, LogSegment, LogSegments, LogStartOffsetIncrementReason, OffsetIndex, ProducerStateManager, ProducerStateManagerConfig, SnapshotFile}
+import org.apache.kafka.storage.internals.log.{AbortedTxn, CleanerConfig, EpochEntry, LogConfig, LogDirFailureChannel, LogFileUtils, LogOffsetMetadata, LogSegment, LogSegments, LogStartOffsetIncrementReason, OffsetIndex, ProducerStateManager, ProducerStateManagerConfig, SnapshotFile}
 import org.apache.kafka.storage.internals.checkpoint.CleanShutdownFileHandler
 import org.junit.jupiter.api.Assertions.{assertDoesNotThrow, assertEquals, assertFalse, assertNotEquals, assertThrows, assertTrue}
 import org.junit.jupiter.api.function.Executable
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test}
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.CsvSource
-import org.mockito.ArgumentMatchers
+import org.mockito.{ArgumentMatchers, Mockito}
 import org.mockito.ArgumentMatchers.{any, anyLong}
 import org.mockito.Mockito.{mock, reset, times, verify, when}
 
@@ -352,19 +352,15 @@ class LogLoaderTest {
       // Intercept all segment read calls
       val interceptedLogSegments = new LogSegments(topicPartition) {
         override def add(segment: LogSegment): LogSegment = {
-          val wrapper = new LogSegment(segment) {
-
-            override def read(startOffset: Long, maxSize: Int, maxPosition: Long, minOneMessage: Boolean): FetchDataInfo = {
-              segmentsWithReads += this
-              super.read(startOffset, maxSize, maxPosition, minOneMessage)
-            }
-
-            override def recover(producerStateManager: ProducerStateManager,
-                                 leaderEpochCache: Optional[LeaderEpochFileCache]): Int = {
-              recoveredSegments += this
-              super.recover(producerStateManager, leaderEpochCache)
-            }
-          }
+          val wrapper = Mockito.spy(segment)
+          Mockito.doAnswer { in =>
+            segmentsWithReads += wrapper
+            segment.read(in.getArgument(0, classOf[java.lang.Long]), in.getArgument(1, classOf[java.lang.Integer]), in.getArgument(2, classOf[java.lang.Long]), in.getArgument(3, classOf[java.lang.Boolean]))
+          }.when(wrapper).read(ArgumentMatchers.any(), ArgumentMatchers.any(), ArgumentMatchers.any(), ArgumentMatchers.any())
+          Mockito.doAnswer { in =>
+            recoveredSegments += wrapper
+            segment.recover(in.getArgument(0, classOf[ProducerStateManager]), in.getArgument(1, classOf[Optional[LeaderEpochFileCache]]))
+          }.when(wrapper).recover(ArgumentMatchers.any(), ArgumentMatchers.any())
           super.add(wrapper)
         }
       }
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
index cdb03e72b3495..d42bf1124560d 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
@@ -128,12 +128,6 @@ public LogSegment(FileRecords log,
         this.created = time.milliseconds();
     }
 
-    // Visible for testing
-    public LogSegment(LogSegment segment) {
-        this(segment.log, segment.lazyOffsetIndex, segment.lazyTimeIndex, segment.txnIndex, segment.baseOffset,
-                segment.indexIntervalBytes, segment.rollJitterMs, segment.time);
-    }
-
     public OffsetIndex offsetIndex() throws IOException {
         return lazyOffsetIndex.get();
     }

From 808bf07fb71e83432a09abe4e3fcc09d4ddc28eb Mon Sep 17 00:00:00 2001
From: Gaurav Narula <gaurav_narula2@apple.com>
Date: Sat, 9 Mar 2024 19:12:39 +0000
Subject: [PATCH 140/258] MINOR: Cleanup log.dirs in ReplicaManagerTest on JVM
 exit (#15289)

- Scala TestUtils now delegates to the function in JTestUtils
- The function is modified such that we delete the rootDir on JVM exit if it didn't exist prior to the function being invoked.

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .gitignore                                       |  1 -
 .../java/org/apache/kafka/test/TestUtils.java    | 16 ++++++++++++++++
 .../test/scala/unit/kafka/utils/TestUtils.scala  |  7 +------
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4ac36a815ba54..7dfe61c38cccf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -34,7 +34,6 @@ Vagrantfile.local
 
 config/server-*
 config/zookeeper-*
-core/data/*
 gradle/wrapper/*.jar
 gradlew.bat
 
diff --git a/clients/src/test/java/org/apache/kafka/test/TestUtils.java b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
index 3a7dcfc9305d6..2052240d38f78 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -235,6 +235,22 @@ public static File tempDirectory() {
         return tempDirectory(null);
     }
 
+    /**
+     * Create a temporary directory under the given root directory.
+     * The root directory is removed on JVM exit if it doesn't already exist
+     * when this function is invoked.
+     *
+     * @param root path to create temporary directory under
+     * @return the temporary directory created within {@code root}
+     */
+    public static File tempRelativeDir(String root) {
+        File rootFile = new File(root);
+        if (rootFile.mkdir()) {
+            rootFile.deleteOnExit();
+        }
+        return tempDirectory(rootFile.toPath(), null);
+    }
+
     /**
      * Create a temporary relative directory in the specified parent directory with the given prefix.
      *
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 17ee3139a51ff..44243d39ce86a 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -133,12 +133,7 @@ object TestUtils extends Logging {
   /**
    * Create a temporary relative directory
    */
-  def tempRelativeDir(parent: String): File = {
-    val parentFile = new File(parent)
-    parentFile.mkdirs()
-
-    JTestUtils.tempDirectory(parentFile.toPath, null)
-  }
+  def tempRelativeDir(root: String): File = JTestUtils.tempRelativeDir(root)
 
   /**
    * Create a random log directory in the format <string>-<int> used for Kafka partition logs.

From 1bb36c2bc12efb537a0f0b160b7013e83f8b34a7 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Sun, 10 Mar 2024 03:28:06 +0800
Subject: [PATCH 141/258] MINOR: change "inter.broker.protocol version" to
 inter.broker.protocol.version (#15504)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 docs/ops.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/ops.html b/docs/ops.html
index 9f49b5c1b2165..6e440b6f3b510 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -3826,7 +3826,7 @@ <h3>Migration Phases</h3>
   <h3>Limitations</h3>
   <ul>
     <li>While a cluster is being migrated from ZK mode to KRaft mode, we do not support changing the <i>metadata
-      version</i> (also known as the <i>inter.broker.protocol</i> version.) Please do not attempt to do this during
+      version</i> (also known as the <i>inter.broker.protocol.version</i>.) Please do not attempt to do this during
       a migration, or you may break the cluster.</li>
     <li>After the migration has been finalized, it is not possible to revert back to ZooKeeper mode.</li>
     <li><a href="#kraft_missing">As noted above</a>, some features are not fully implemented in KRaft mode. If you are

From 66ed6d3b268cf481852297ee08ecc67d0f919aec Mon Sep 17 00:00:00 2001
From: Kamal Chandraprakash <kchandraprakash@uber.com>
Date: Mon, 11 Mar 2024 14:50:29 +0530
Subject: [PATCH 142/258] KAFKA-16146: Checkpoint log-start-offset for remote
 log enabled topics (#15201)

The log-start-offset was not getting flushed to the checkpoint file due to the check where we compare the log-start-offset with the localLog first segment base offset only. This change makes sure that tiered storage enabled topics will always try to add their entries in the log-start-offset checkpoint file.

Reviewers: Jun Rao <junrao@gmail.com>, Satish Duggana <satishd@apache.org>
---
 .../src/main/scala/kafka/log/LogManager.scala |  3 +-
 .../scala/unit/kafka/log/LogManagerTest.scala | 80 ++++++++++++++++++-
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala
index db96c497c5537..7c8bbec5292e9 100755
--- a/core/src/main/scala/kafka/log/LogManager.scala
+++ b/core/src/main/scala/kafka/log/LogManager.scala
@@ -865,7 +865,8 @@ class LogManager(logDirs: Seq[File],
     try {
       logStartOffsetCheckpoints.get(logDir).foreach { checkpoint =>
         val logStartOffsets = logsToCheckpoint.collect {
-          case (tp, log) if log.logStartOffset > log.logSegments.asScala.head.baseOffset => tp -> log.logStartOffset
+          case (tp, log) if log.remoteLogEnabled() || log.logStartOffset > log.logSegments.asScala.head.baseOffset =>
+            tp -> log.logStartOffset
         }
         checkpoint.write(logStartOffsets)
       }
diff --git a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
index 2dcc3e5a35cc0..ce6e2a3b04ee3 100755
--- a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
@@ -46,7 +46,7 @@ import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Future}
 import java.util.{Collections, Optional, OptionalLong, Properties}
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.kafka.server.util.MockTime
-import org.apache.kafka.storage.internals.log.{FetchDataInfo, FetchIsolation, LogConfig, LogDirFailureChannel, ProducerStateManagerConfig, RemoteIndexCache}
+import org.apache.kafka.storage.internals.log.{FetchDataInfo, FetchIsolation, LogConfig, LogDirFailureChannel, LogStartOffsetIncrementReason, ProducerStateManagerConfig, RemoteIndexCache}
 import org.apache.kafka.storage.internals.checkpoint.CleanShutdownFileHandler
 
 import scala.collection.{Map, mutable}
@@ -1117,6 +1117,84 @@ class LogManagerTest {
     assertEquals(2, logManager.directoryIdsSet.size)
   }
 
+  @Test
+  def testCheckpointLogStartOffsetForRemoteTopic(): Unit = {
+    logManager.shutdown()
+
+    val props = new Properties()
+    props.putAll(logProps)
+    props.put(TopicConfig.REMOTE_LOG_STORAGE_ENABLE_CONFIG, "true")
+    val logConfig = new LogConfig(props)
+    logManager = TestUtils.createLogManager(
+      defaultConfig = logConfig,
+      configRepository = new MockConfigRepository,
+      logDirs = Seq(this.logDir),
+      time = this.time,
+      recoveryThreadsPerDataDir = 1,
+      remoteStorageSystemEnable = true
+    )
+
+    val checkpointFile = new File(logDir, LogManager.LogStartOffsetCheckpointFile)
+    val checkpoint = new OffsetCheckpointFile(checkpointFile)
+    val topicPartition = new TopicPartition("test", 0)
+    val log = logManager.getOrCreateLog(topicPartition, topicId = None)
+    var offset = 0L
+    for(_ <- 0 until 50) {
+      val set = TestUtils.singletonRecords("test".getBytes())
+      val info = log.appendAsLeader(set, leaderEpoch = 0)
+      offset = info.lastOffset
+      if (offset != 0 && offset % 10 == 0)
+        log.roll()
+    }
+    assertEquals(5, log.logSegments.size())
+    log.updateHighWatermark(49)
+    // simulate calls to upload 3 segments to remote storage and remove them from local-log.
+    log.updateHighestOffsetInRemoteStorage(30)
+    log.maybeIncrementLocalLogStartOffset(31L, LogStartOffsetIncrementReason.SegmentDeletion)
+    log.deleteOldSegments()
+    assertEquals(2, log.logSegments.size())
+
+    // simulate two remote-log segment deletion
+    val logStartOffset = 21L
+    log.maybeIncrementLogStartOffset(logStartOffset, LogStartOffsetIncrementReason.SegmentDeletion)
+    logManager.checkpointLogStartOffsets()
+
+    assertEquals(logStartOffset, log.logStartOffset)
+    assertEquals(logStartOffset, checkpoint.read().getOrElse(topicPartition, -1L))
+  }
+
+  @Test
+  def testCheckpointLogStartOffsetForNormalTopic(): Unit = {
+    val checkpointFile = new File(logDir, LogManager.LogStartOffsetCheckpointFile)
+    val checkpoint = new OffsetCheckpointFile(checkpointFile)
+    val topicPartition = new TopicPartition("test", 0)
+    val log = logManager.getOrCreateLog(topicPartition, topicId = None)
+    var offset = 0L
+    for(_ <- 0 until 50) {
+      val set = TestUtils.singletonRecords("test".getBytes())
+      val info = log.appendAsLeader(set, leaderEpoch = 0)
+      offset = info.lastOffset
+      if (offset != 0 && offset % 10 == 0)
+        log.roll()
+    }
+    assertEquals(5, log.logSegments.size())
+    log.updateHighWatermark(49)
+
+    val logStartOffset = 31L
+    log.maybeIncrementLogStartOffset(logStartOffset, LogStartOffsetIncrementReason.SegmentDeletion)
+    logManager.checkpointLogStartOffsets()
+    assertEquals(5, log.logSegments.size())
+    assertEquals(logStartOffset, checkpoint.read().getOrElse(topicPartition, -1L))
+
+    log.deleteOldSegments()
+    assertEquals(2, log.logSegments.size())
+    assertEquals(logStartOffset, log.logStartOffset)
+
+    // When you checkpoint log-start-offset after removing the segments, then there should not be any checkpoint
+    logManager.checkpointLogStartOffsets()
+    assertEquals(-1L, checkpoint.read().getOrElse(topicPartition, -1L))
+  }
+
   def writeMetaProperties(dir: File, directoryId: Optional[Uuid] = Optional.empty()): Unit = {
     val metaProps = new MetaProperties.Builder().
       setVersion(MetaPropertiesVersion.V0).

From 8b72a2c72f09838fdd2e7416c98d30fe876b4078 Mon Sep 17 00:00:00 2001
From: Christo Lolov <lolovc@amazon.com>
Date: Mon, 11 Mar 2024 11:51:20 +0000
Subject: [PATCH 143/258] KAFKA-14133: Move consumer mock in TaskManagerTest to
 Mockito - part 3 (#15497)

The previous pull request in this series was #15261.

This pull request continues the migration of the consumer mock in TaskManagerTest test by test for easier reviews.

The next pull request in the series will be #15254 which ought to complete the Mockito migration for the TaskManagerTest class

Reviewer: Bruno Cadonna <cadonna@apache.org>
---
 .../processor/internals/TaskManagerTest.java  | 318 ++++++------------
 1 file changed, 96 insertions(+), 222 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index 681e69d300487..36d2a3e378633 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -194,10 +194,8 @@ public class TaskManagerTest {
     private StateDirectory stateDirectory;
     @org.mockito.Mock
     private ChangelogReader changeLogReader;
-    @Mock(type = MockType.STRICT)
-    private Consumer<byte[], byte[]> consumer;
     @org.mockito.Mock
-    private Consumer<byte[], byte[]> mockitoConsumer;
+    private Consumer<byte[], byte[]> consumer;
     @org.mockito.Mock
     private ActiveTaskCreator activeTaskCreator;
     @org.mockito.Mock
@@ -311,7 +309,6 @@ public void shouldLockAllTasksOnCorruptionWithProcessingThreads() {
             .withInputPartitions(taskId00Partitions).build();
         final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
-        taskManager.setMainConsumer(mockitoConsumer);
         when(tasks.activeTaskIds()).thenReturn(mkSet(taskId00, taskId01));
         when(tasks.task(taskId00)).thenReturn(activeTask1);
         final KafkaFuture<Void> mockFuture = KafkaFuture.completedFuture(null);
@@ -319,7 +316,7 @@ public void shouldLockAllTasksOnCorruptionWithProcessingThreads() {
 
         taskManager.handleCorruption(mkSet(taskId00));
 
-        Mockito.verify(mockitoConsumer).assignment();
+        Mockito.verify(consumer).assignment();
         Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
         Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
     }
@@ -1252,13 +1249,12 @@ public void shouldSuspendRevokedTaskRemovedFromStateUpdater() {
         when(stateUpdater.hasRemovedTasks()).thenReturn(true);
         when(stateUpdater.drainRemovedTasks()).thenReturn(mkSet(statefulTask));
         taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
         Mockito.verify(statefulTask).suspend();
         Mockito.verify(tasks).addTask(statefulTask);
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1284,7 +1280,7 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
         when(stateUpdater.drainRemovedTasks())
             .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions, taskToCloseReviveAndUpdateInputPartitions));
         when(stateUpdater.restoresActiveTasks()).thenReturn(true);
-        when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, mockitoConsumer))
+        when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, consumer))
             .thenReturn(convertedTask1);
         when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle0, taskId00Partitions))
             .thenReturn(convertedTask0);
@@ -1302,7 +1298,6 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
             argThat(taskId -> !taskId.equals(taskToCloseReviveAndUpdateInputPartitions.id()))
         )).thenReturn(null);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter -> { });
 
@@ -1320,7 +1315,7 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
         Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId05Partitions), anyMap());
         Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded();
         Mockito.verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions);
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1456,14 +1451,13 @@ public void shouldTransitRestoredTaskToRunning() {
             .withInputPartitions(taskId00Partitions).build();
         final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
         Mockito.verify(task).completeRestoration(noOpResetter);
         Mockito.verify(task).clearTaskTimeout();
         Mockito.verify(tasks).addTask(task);
-        Mockito.verify(mockitoConsumer).resume(task.inputPartitions());
+        Mockito.verify(consumer).resume(task.inputPartitions());
     }
 
     @Test
@@ -1473,7 +1467,6 @@ public void shouldHandleTimeoutExceptionInTransitRestoredTaskToRunning() {
             .withInputPartitions(taskId00Partitions).build();
         final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks);
-        taskManager.setMainConsumer(mockitoConsumer);
         final TimeoutException timeoutException = new TimeoutException();
         doThrow(timeoutException).when(task).completeRestoration(noOpResetter);
 
@@ -1482,7 +1475,7 @@ public void shouldHandleTimeoutExceptionInTransitRestoredTaskToRunning() {
         Mockito.verify(task).maybeInitTaskTimeoutOrThrow(anyLong(), Mockito.eq(timeoutException));
         Mockito.verify(tasks, never()).addTask(task);
         Mockito.verify(task, never()).clearTaskTimeout();
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     private TaskManager setUpTransitionToRunningOfRestoredTask(final StreamTask statefulTask,
@@ -1672,11 +1665,10 @@ public void shouldUpdateInputPartitionsOfRestoredTask() {
         when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(mkSet(statefulTask));
         when(stateUpdater.restoresActiveTasks()).thenReturn(true);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(mockitoConsumer).resume(statefulTask.inputPartitions());
+        Mockito.verify(consumer).resume(statefulTask.inputPartitions());
         Mockito.verify(statefulTask).updateInputPartitions(Mockito.eq(taskId01Partitions), anyMap());
         Mockito.verify(statefulTask).completeRestoration(noOpResetter);
         Mockito.verify(statefulTask).clearTaskTimeout();
@@ -1717,13 +1709,12 @@ public void shouldSuspendRestoredTaskIfRevoked() {
         when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(mkSet(statefulTask));
         when(stateUpdater.restoresActiveTasks()).thenReturn(true);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
         Mockito.verify(statefulTask).suspend();
         Mockito.verify(tasks).addTask(statefulTask);
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1974,12 +1965,11 @@ public void shouldUnlockEmptyDirsAtRebalanceStart() throws Exception {
     @Test
     public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() {
         final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
-        taskManager.setMainConsumer(mockitoConsumer);
-        when(mockitoConsumer.assignment()).thenReturn(assigned);
+        when(consumer.assignment()).thenReturn(assigned);
 
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(mockitoConsumer).pause(assigned);
+        Mockito.verify(consumer).pause(assigned);
     }
 
     @Test
@@ -1989,14 +1979,13 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
             .withInputPartitions(taskId00Partitions).build();
         final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
         when(tasks.allTasks()).thenReturn(mkSet(statefulTask0));
         final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
-        when(mockitoConsumer.assignment()).thenReturn(assigned);
+        when(consumer.assignment()).thenReturn(assigned);
 
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(mockitoConsumer).pause(mkSet(t1p1));
+        Mockito.verify(consumer).pause(mkSet(t1p1));
     }
 
     @Test
@@ -2021,7 +2010,7 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01)));
         verify(stateDirectory);
 
-        Mockito.verify(mockitoConsumer).pause(assignment);
+        Mockito.verify(consumer).pause(assignment);
     }
 
     @Test
@@ -2040,7 +2029,6 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
             .withInputPartitions(taskId03Partitions).build();
         final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId00, runningStatefulTask)));
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask, restoringStatefulTask));
         when(tasks.allTasks()).thenReturn(mkSet(runningStatefulTask));
@@ -2056,12 +2044,12 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
         replay(stateDirectory);
 
         final Set<TopicPartition> assigned = mkSet(t1p0, t1p1, t1p2);
-        when(mockitoConsumer.assignment()).thenReturn(assigned);
+        when(consumer.assignment()).thenReturn(assigned);
 
         taskManager.handleRebalanceStart(singleton("topic"));
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(mockitoConsumer).pause(mkSet(t1p1, t1p2));
+        Mockito.verify(consumer).pause(mkSet(t1p1, t1p2));
         verify(stateDirectory);
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
     }
@@ -2331,12 +2319,10 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         // first `handleAssignment`
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -2362,8 +2348,6 @@ public void closeClean() {
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         taskManager.handleRevocation(taskId00Partitions);
 
@@ -2387,7 +2371,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
         // `handleAssignment`
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
@@ -2400,8 +2384,6 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
         expectLockObtainedFor();
         replay(stateDirectory);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleRebalanceStart(emptySet());
         assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01)));
 
@@ -2443,14 +2425,12 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         // `handleAssignment`
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         // `handleAssignment`
         doThrow(new RuntimeException("KABOOM!")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -2480,7 +2460,6 @@ public void shouldReAddRevivedTasksToStateUpdater() {
             .withInputPartitions(taskId02Partitions).build();
         final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
         when(tasks.task(taskId03)).thenReturn(corruptedActiveTask);
         when(tasks.task(taskId02)).thenReturn(corruptedStandbyTask);
 
@@ -2496,7 +2475,7 @@ public void shouldReAddRevivedTasksToStateUpdater() {
         Mockito.verify(tasks).removeTask(corruptedStandbyTask);
         Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedActiveTask));
         Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedStandbyTask));
-        Mockito.verify(mockitoConsumer).assignment();
+        Mockito.verify(consumer).assignment();
     }
 
     @Test
@@ -2515,13 +2494,11 @@ public void postCommit(final boolean enforceCheckpoint) {
         };
 
         // `handleAssignment`
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(taskId00Partitions);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -2551,13 +2528,11 @@ public void suspend() {
             }
         };
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(taskId00Partitions);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -2587,12 +2562,10 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedTask, nonCorruptedTask));
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(taskId00Partitions);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
 
@@ -2607,7 +2580,7 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
         // check that we should not commit empty map either
-        Mockito.verify(mockitoConsumer, never()).commitSync(emptyMap());
+        Mockito.verify(consumer, never()).commitSync(emptyMap());
         Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
@@ -2626,9 +2599,7 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         // `handleAssignment`
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
             .thenReturn(asList(corruptedTask, nonRunningNonCorruptedTask));
-        when(mockitoConsumer.assignment()).thenReturn(taskId00Partitions);
-
-        taskManager.setMainConsumer(mockitoConsumer);
+        when(consumer.assignment()).thenReturn(taskId00Partitions);
 
         taskManager.handleAssignment(assignment, emptyMap());
 
@@ -2658,9 +2629,8 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndNotCommitRunningSt
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId02, corruptedTask)));
         when(tasks.task(taskId02)).thenReturn(corruptedTask);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        taskManager.setMainConsumer(mockitoConsumer);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeRestoringTask, standbyTask));
-        when(mockitoConsumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
+        when(consumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         taskManager.handleCorruption(mkSet(taskId02));
 
@@ -2692,8 +2662,7 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndCommitRunningStand
         ));
         when(tasks.task(taskId02)).thenReturn(corruptedTask);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, false);
-        taskManager.setMainConsumer(mockitoConsumer);
-        when(mockitoConsumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
+        when(consumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         taskManager.handleCorruption(mkSet(taskId02));
 
@@ -2721,9 +2690,7 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             .thenReturn(singleton(runningNonCorruptedActive));
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singleton(corruptedStandby));
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
-
-        taskManager.setMainConsumer(mockitoConsumer);
+        when(consumer.assignment()).thenReturn(assignment);
 
         taskManager.handleAssignment(taskId01Assignment, taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2762,7 +2729,7 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
         when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignement)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
@@ -2770,8 +2737,6 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
 
         uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(firstAssignement, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -2813,13 +2778,11 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
-        doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(offsets);
-
-        taskManager.setMainConsumer(mockitoConsumer);
+        doThrow(new TimeoutException()).when(consumer).commitSync(offsets);
 
         taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2886,17 +2849,15 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask));
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
+        when(consumer.groupMetadata()).thenReturn(groupMetadata);
 
         doThrow(new TimeoutException()).when(producer).commitTransaction(offsets, groupMetadata);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -2964,16 +2925,14 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             mkEntry(taskId02, taskId02Partitions)
         );
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded));
 
-        doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(expectedCommittedOffsets);
-
-        taskManager.setMainConsumer(mockitoConsumer);
+        doThrow(new TimeoutException()).when(consumer).commitSync(expectedCommittedOffsets);
 
         taskManager.handleAssignment(assignmentActive, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3024,7 +2983,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             mkEntry(taskId02, taskId02Partitions)
             );
 
-        when(mockitoConsumer.assignment())
+        when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
@@ -3032,12 +2991,10 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
+        when(consumer.groupMetadata()).thenReturn(groupMetadata);
 
         doThrow(new TimeoutException()).when(producer).commitTransaction(expectedCommittedOffsets, groupMetadata);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(revokedActiveTask.state(), is(Task.State.RUNNING));
@@ -3063,11 +3020,9 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
     public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(emptyMap(), taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3083,12 +3038,10 @@ public void shouldAddNonResumedSuspendedTasks() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3101,19 +3054,17 @@ public void shouldAddNonResumedSuspendedTasks() {
 
         // expect these calls twice (because we're going to tryToCompleteRestoration twice)
         Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
-        Mockito.verify(mockitoConsumer, times(2)).assignment();
-        Mockito.verify(mockitoConsumer, times(2)).resume(assignment);
+        Mockito.verify(consumer, times(2)).assignment();
+        Mockito.verify(consumer, times(2)).resume(assignment);
     }
 
     @Test
     public void shouldUpdateInputPartitionsAfterRebalance() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3125,8 +3076,8 @@ public void shouldUpdateInputPartitionsAfterRebalance() {
         assertThat(task00.state(), is(Task.State.RUNNING));
         assertEquals(newPartitionsSet, task00.inputPartitions());
         // expect these calls twice (because we're going to tryToCompleteRestoration twice)
-        Mockito.verify(mockitoConsumer, times(2)).resume(assignment);
-        Mockito.verify(mockitoConsumer, times(2)).assignment();
+        Mockito.verify(consumer, times(2)).resume(assignment);
+        Mockito.verify(consumer, times(2)).assignment();
         Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
     }
 
@@ -3135,8 +3086,6 @@ public void shouldAddNewActiveTasks() {
         final Map<TaskId, Set<TopicPartition>> assignment = taskId00Assignment;
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3149,8 +3098,8 @@ public void shouldAddNewActiveTasks() {
         assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(mockitoConsumer).assignment();
-        Mockito.verify(mockitoConsumer).resume(Mockito.eq(emptySet()));
+        Mockito.verify(consumer).assignment();
+        Mockito.verify(consumer).resume(Mockito.eq(emptySet()));
     }
 
     @Test
@@ -3172,8 +3121,6 @@ public void initializeIfNeeded() {
             }
         };
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01));
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3191,7 +3138,7 @@ public void initializeIfNeeded() {
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -3206,8 +3153,6 @@ public void completeRestoration(final java.util.function.Consumer<Set<TopicParti
             }
         };
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3223,7 +3168,7 @@ public void completeRestoration(final java.util.function.Consumer<Set<TopicParti
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -3232,11 +3177,9 @@ public void shouldSuspendActiveTasksDuringRevocation() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task00.setCommittableOffsetsAndMetadata(offsets);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3279,7 +3222,7 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = mkMap(
             mkEntry(taskId10, taskId10Partitions)
         );
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
@@ -3289,15 +3232,13 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
             .thenReturn(singletonList(task10));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
-        when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata);
+        when(consumer.groupMetadata()).thenReturn(groupMetadata);
 
         task00.committedOffsets();
         task01.committedOffsets();
         task02.committedOffsets();
         task10.committedOffsets();
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3346,15 +3287,13 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = mkMap(
             mkEntry(taskId10, taskId10Partitions)
         );
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task10));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3371,7 +3310,7 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         assertThat(task02.commitPrepared, is(false));
         assertThat(task10.commitPrepared, is(false));
 
-        Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets);
+        Mockito.verify(consumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -3386,13 +3325,11 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() {
         final Map<TaskId, Set<TopicPartition>> assignmentActive = singletonMap(taskId00, taskId00Partitions);
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = singletonMap(taskId10, taskId10Partitions);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3416,13 +3353,11 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() {
         final Map<TaskId, Set<TopicPartition>> assignmentActive = singletonMap(taskId00, taskId00Partitions);
         final Map<TaskId, Set<TopicPartition>> assignmentStandby = singletonMap(taskId10, taskId10Partitions);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3460,9 +3395,8 @@ public void suspend() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
-        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3810,8 +3744,6 @@ public void shouldCloseStandbyTasksOnShutdown() {
         final Map<TaskId, Set<TopicPartition>> assignment = singletonMap(taskId00, taskId00Partitions);
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         // `handleAssignment`
         when(standbyTaskCreator.createTasks(assignment)).thenReturn(singletonList(task00));
 
@@ -3830,8 +3762,8 @@ public void shouldCloseStandbyTasksOnShutdown() {
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
         Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
         // `tryToCompleteRestoration`
-        Mockito.verify(mockitoConsumer).assignment();
-        Mockito.verify(mockitoConsumer).resume(Mockito.eq(emptySet()));
+        Mockito.verify(consumer).assignment();
+        Mockito.verify(consumer).resume(Mockito.eq(emptySet()));
     }
 
     @Test
@@ -3918,13 +3850,11 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() {
     @Test
     public void shouldInitializeNewActiveTasks() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -3932,18 +3862,16 @@ public void shouldInitializeNewActiveTasks() {
         assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         // verifies that we actually resume the assignment at the end of restoration.
-        Mockito.verify(mockitoConsumer).resume(assignment);
+        Mockito.verify(consumer).resume(assignment);
     }
 
     @Test
     public void shouldInitialiseNewStandbyTasks() {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(emptyMap(), taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -3954,8 +3882,7 @@ public void shouldInitialiseNewStandbyTasks() {
 
     @Test
     public void shouldHandleRebalanceEvents() {
-        taskManager.setMainConsumer(mockitoConsumer);
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>());
         replay(stateDirectory);
         assertThat(taskManager.rebalanceInProgress(), is(false));
@@ -3963,7 +3890,7 @@ public void shouldHandleRebalanceEvents() {
         assertThat(taskManager.rebalanceInProgress(), is(true));
         taskManager.handleRebalanceComplete();
         assertThat(taskManager.rebalanceInProgress(), is(false));
-        Mockito.verify(mockitoConsumer).pause(assignment);
+        Mockito.verify(consumer).pause(assignment);
     }
 
     @Test
@@ -3973,14 +3900,12 @@ public void shouldCommitActiveAndStandbyTasks() {
         task00.setCommittableOffsetsAndMetadata(offsets);
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
             .thenReturn(singletonList(task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -3994,7 +3919,7 @@ public void shouldCommitActiveAndStandbyTasks() {
         assertThat(task00.commitNeeded, is(false));
         assertThat(task01.commitNeeded, is(false));
 
-        Mockito.verify(mockitoConsumer).commitSync(offsets);
+        Mockito.verify(consumer).commitSync(offsets);
     }
 
     @Test
@@ -4017,14 +3942,12 @@ public void shouldCommitProvidedTasksIfNeeded() {
             mkEntry(taskId05, taskId05Partitions)
         );
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(Arrays.asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(Arrays.asList(task03, task04, task05));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4049,11 +3972,9 @@ public void shouldCommitProvidedTasksIfNeeded() {
     public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4073,7 +3994,7 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
         makeTaskFolders(taskId00.toString(), taskId01.toString());
         expectDirectoryNotEmpty(taskId00, taskId01);
         expectLockObtainedFor(taskId00, taskId01);
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
@@ -4081,8 +4002,6 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
 
         replay(stateDirectory);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4111,12 +4030,11 @@ public void shouldCommitViaConsumerIfEosDisabled() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p1, new OffsetAndMetadata(0L, null));
         task01.setCommittableOffsetsAndMetadata(offsets);
         task01.setCommitNeeded();
-        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.addTask(task01);
 
         taskManager.commitAll();
 
-        Mockito.verify(mockitoConsumer).commitSync(offsets);
+        Mockito.verify(consumer).commitSync(offsets);
     }
 
     @Test
@@ -4156,7 +4074,6 @@ private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processing
                                                      final Map<TopicPartition, OffsetAndMetadata> offsetsT01,
                                                      final Map<TopicPartition, OffsetAndMetadata> offsetsT02) {
         final TaskManager taskManager = setUpTaskManager(processingMode, false);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
         task01.setCommittableOffsetsAndMetadata(offsetsT01);
@@ -4167,7 +4084,7 @@ private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processing
         task02.setCommitNeeded();
         taskManager.addTask(task02);
 
-        when(mockitoConsumer.groupMetadata()).thenReturn(new ConsumerGroupMetadata("appId"));
+        when(consumer.groupMetadata()).thenReturn(new ConsumerGroupMetadata("appId"));
 
         taskManager.commitAll();
     }
@@ -4181,11 +4098,9 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4207,11 +4122,9 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(emptyMap(), taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4241,11 +4154,9 @@ public Map<TopicPartition, Long> purgeableOffsets() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4276,11 +4187,9 @@ public Map<TopicPartition, Long> purgeableOffsets() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4300,15 +4209,13 @@ public Map<TopicPartition, Long> purgeableOffsets() {
     public void shouldIgnorePurgeDataErrors() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
 
         final KafkaFutureImpl<DeletedRecords> futureDeletedRecords = new KafkaFutureImpl<>();
         final DeleteRecordsResult deleteRecordsResult = new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords));
         futureDeletedRecords.completeExceptionally(new Exception("KABOOM!"));
         when(adminClient.deleteRecords(any())).thenReturn(deleteRecordsResult);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.addTask(task00);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4349,14 +4256,12 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
             mkEntry(taskId10, taskId10Partitions)
         );
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02, task03));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task04));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4381,7 +4286,7 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
 
         assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3));
 
-        Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets);
+        Mockito.verify(consumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -4393,12 +4298,10 @@ public void shouldProcessActiveTasks() {
         firstAssignment.put(taskId00, taskId00Partitions);
         firstAssignment.put(taskId01, taskId01Partitions);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
             .thenReturn(Arrays.asList(task00, task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(firstAssignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4507,11 +4410,9 @@ public boolean process(final long wallClockTime) {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4532,12 +4433,10 @@ public boolean process(final long wallClockTime) {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4561,11 +4460,9 @@ public boolean maybePunctuateStreamTime() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4583,11 +4480,9 @@ public boolean maybePunctuateStreamTime() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4610,11 +4505,9 @@ public boolean maybePunctuateSystemTime() {
             }
         };
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4633,14 +4526,12 @@ public Set<TopicPartition> changelogPartitions() {
             }
         };
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(false));
         assertThat(task00.state(), is(Task.State.RESTORING));
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -4649,11 +4540,9 @@ public void shouldHaveRemainingPartitionsUncleared() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task00.setCommittableOffsetsAndMetadata(offsets);
 
-        when(mockitoConsumer.assignment()).thenReturn(assignment);
+        when(consumer.assignment()).thenReturn(assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) {
             appender.setClassLoggerToDebug(TaskManager.class);
             taskManager.handleAssignment(taskId00Assignment, emptyMap());
@@ -4808,9 +4697,7 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
         when(standbyTaskCreator.createTasks(standbyAssignment)).thenReturn(standbyTasks);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(allActiveTasksAssignment))).thenReturn(allActiveTasks);
 
-        lenient().when(mockitoConsumer.assignment()).thenReturn(assignment);
-
-        taskManager.setMainConsumer(mockitoConsumer);
+        lenient().when(consumer.assignment()).thenReturn(assignment);
 
         taskManager.handleAssignment(allActiveTasksAssignment, standbyAssignment);
         taskManager.tryToCompleteRestoration(time.milliseconds(), null);
@@ -4864,10 +4751,9 @@ public void shouldThrowTaskMigratedExceptionOnCommitFailed() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task01.setCommittableOffsetsAndMetadata(offsets);
         task01.setCommitNeeded();
-        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.addTask(task01);
 
-        doThrow(new CommitFailedException()).when(mockitoConsumer).commitSync(offsets);
+        doThrow(new CommitFailedException()).when(consumer).commitSync(offsets);
 
         final TaskMigratedException thrown = assertThrows(
             TaskMigratedException.class,
@@ -4892,9 +4778,7 @@ public void shouldNotFailForTimeoutExceptionOnConsumerCommit() {
         task00.setCommittableOffsetsAndMetadata(taskId00Partitions.stream().collect(Collectors.toMap(p -> p, p -> new OffsetAndMetadata(0))));
         task01.setCommittableOffsetsAndMetadata(taskId00Partitions.stream().collect(Collectors.toMap(p -> p, p -> new OffsetAndMetadata(0))));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
-        doThrow(new TimeoutException("KABOOM!")).doNothing().when(mockitoConsumer).commitSync(any(Map.class));
+        doThrow(new TimeoutException("KABOOM!")).doNothing().when(consumer).commitSync(any(Map.class));
 
         task00.setCommitNeeded();
 
@@ -4906,14 +4790,13 @@ public void shouldNotFailForTimeoutExceptionOnConsumerCommit() {
         assertNull(task00.timeout);
         assertNull(task01.timeout);
 
-        Mockito.verify(mockitoConsumer, times(2)).commitSync(any(Map.class));
+        Mockito.verify(consumer, times(2)).commitSync(any(Map.class));
     }
 
     @Test
     public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() {
         final Tasks tasks = mock(Tasks.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_ALPHA, tasks, false);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.streamsProducerForTask(any(TaskId.class))).thenReturn(producer);
@@ -4949,13 +4832,12 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() {
             equalTo(Collections.singleton(taskId00))
         );
 
-        Mockito.verify(mockitoConsumer, times(2)).groupMetadata();
+        Mockito.verify(consumer, times(2)).groupMetadata();
     }
 
     @Test
     public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV2() {
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false);
-        taskManager.setMainConsumer(mockitoConsumer);
 
         final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.threadProducer()).thenReturn(producer);
@@ -4985,7 +4867,7 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV
             equalTo(mkSet(taskId00, taskId01))
         );
 
-        Mockito.verify(mockitoConsumer).groupMetadata();
+        Mockito.verify(consumer).groupMetadata();
     }
 
     @Test
@@ -4994,10 +4876,9 @@ public void shouldStreamsExceptionOnCommitError() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task01.setCommittableOffsetsAndMetadata(offsets);
         task01.setCommitNeeded();
-        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.addTask(task01);
 
-        doThrow(new KafkaException()).when(mockitoConsumer).commitSync(offsets);
+        doThrow(new KafkaException()).when(consumer).commitSync(offsets);
 
         final StreamsException thrown = assertThrows(
             StreamsException.class,
@@ -5015,10 +4896,9 @@ public void shouldFailOnCommitFatal() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task01.setCommittableOffsetsAndMetadata(offsets);
         task01.setCommitNeeded();
-        taskManager.setMainConsumer(mockitoConsumer);
         taskManager.addTask(task01);
 
-        doThrow(new RuntimeException("KABOOM")).when(mockitoConsumer).commitSync(offsets);
+        doThrow(new RuntimeException("KABOOM")).when(consumer).commitSync(offsets);
 
         final RuntimeException thrown = assertThrows(
             RuntimeException.class,
@@ -5044,8 +4924,6 @@ public void suspend() {
         assignment.putAll(taskId01Assignment);
         when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01));
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(assignment, Collections.emptyMap());
 
         final RuntimeException thrown = assertThrows(
@@ -5055,7 +4933,7 @@ public void suspend() {
         assertThat(thrown.getCause().getMessage(), is("KABOOM!"));
         assertThat(task00.state(), is(Task.State.SUSPENDED));
         assertThat(task01.state(), is(Task.State.SUSPENDED));
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -5071,15 +4949,13 @@ public void shouldConvertActiveTaskToStandbyTask() {
         when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(activeTask));
         when(standbyTaskCreator.createStandbyTaskFromActive(Mockito.any(), Mockito.eq(taskId00Partitions))).thenReturn(standbyTask);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap());
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
 
         Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
         Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
         Mockito.verify(standbyTaskCreator, times(2)).createTasks(Collections.emptyMap());
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test
@@ -5096,14 +4972,12 @@ public void shouldConvertStandbyTaskToActiveTask() {
         when(activeTaskCreator.createActiveTaskFromStandby(Mockito.eq(standbyTask), Mockito.eq(taskId00Partitions), any()))
             .thenReturn(activeTask);
 
-        taskManager.setMainConsumer(mockitoConsumer);
-
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
         taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap());
 
         Mockito.verify(activeTaskCreator, times(2)).createTasks(any(), Mockito.eq(emptyMap()));
         Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
-        Mockito.verifyNoInteractions(mockitoConsumer);
+        Mockito.verifyNoInteractions(consumer);
     }
 
     @Test

From 58ddd693e69599b177d09c2e384f31e7f5e11171 Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <cadonna@apache.org>
Date: Tue, 12 Mar 2024 11:04:37 +0100
Subject: [PATCH 144/258] KAFKA-16227: Avoid IllegalStateException during fetch
 initialization (#15491)

The AsyncKafkaConsumer might throw an IllegalStateException during
the initialization of a new fetch. The exception is caused by
 the partition being unassigned by the background thread before
the subscription state is accessed during initialisation.

This commit avoids the IllegalStateException by verifying that
the partition was not unassigned each time the subscription state
is accessed.

Reviewer: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../consumer/internals/FetchCollector.java    |  62 ++--
 .../consumer/internals/SubscriptionState.java |  65 +++-
 .../internals/FetchCollectorTest.java         | 308 ++++++++++++++++--
 .../internals/SubscriptionStateTest.java      |  72 ++++
 4 files changed, 460 insertions(+), 47 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchCollector.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchCollector.java
index 692a67b9c301e..998136797b573 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchCollector.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchCollector.java
@@ -248,10 +248,10 @@ private CompletedFetch handleInitializeSuccess(final CompletedFetch completedFet
 
         // we are interested in this fetch only if the beginning offset matches the
         // current consumed position
-        SubscriptionState.FetchPosition position = subscriptions.position(tp);
+        SubscriptionState.FetchPosition position = subscriptions.positionOrNull(tp);
         if (position == null || position.offset != fetchOffset) {
             log.debug("Discarding stale fetch response for partition {} since its offset {} does not match " +
-                    "the expected offset {}", tp, fetchOffset, position);
+                "the expected offset {} or the partition has been unassigned", tp, fetchOffset, position);
             return null;
         }
 
@@ -278,32 +278,48 @@ private CompletedFetch handleInitializeSuccess(final CompletedFetch completedFet
             }
         }
 
-        if (partition.highWatermark() >= 0) {
-            log.trace("Updating high watermark for partition {} to {}", tp, partition.highWatermark());
-            subscriptions.updateHighWatermark(tp, partition.highWatermark());
+        if (!updatePartitionState(partition, tp)) {
+            return null;
+        }
+
+        completedFetch.setInitialized();
+        return completedFetch;
+    }
+
+    private boolean updatePartitionState(final FetchResponseData.PartitionData partitionData,
+                                         final TopicPartition tp) {
+        if (partitionData.highWatermark() >= 0) {
+            log.trace("Updating high watermark for partition {} to {}", tp, partitionData.highWatermark());
+            if (!subscriptions.tryUpdatingHighWatermark(tp, partitionData.highWatermark())) {
+                return false;
+            }
         }
 
-        if (partition.logStartOffset() >= 0) {
-            log.trace("Updating log start offset for partition {} to {}", tp, partition.logStartOffset());
-            subscriptions.updateLogStartOffset(tp, partition.logStartOffset());
+        if (partitionData.logStartOffset() >= 0) {
+            log.trace("Updating log start offset for partition {} to {}", tp, partitionData.logStartOffset());
+            if (!subscriptions.tryUpdatingLogStartOffset(tp, partitionData.logStartOffset())) {
+                return false;
+            }
         }
 
-        if (partition.lastStableOffset() >= 0) {
-            log.trace("Updating last stable offset for partition {} to {}", tp, partition.lastStableOffset());
-            subscriptions.updateLastStableOffset(tp, partition.lastStableOffset());
+        if (partitionData.lastStableOffset() >= 0) {
+            log.trace("Updating last stable offset for partition {} to {}", tp, partitionData.lastStableOffset());
+            if (!subscriptions.tryUpdatingLastStableOffset(tp, partitionData.lastStableOffset())) {
+                return false;
+            }
         }
 
-        if (FetchResponse.isPreferredReplica(partition)) {
-            subscriptions.updatePreferredReadReplica(completedFetch.partition, partition.preferredReadReplica(), () -> {
-                long expireTimeMs = time.milliseconds() + metadata.metadataExpireMs();
-                log.debug("Updating preferred read replica for partition {} to {}, set to expire at {}",
-                        tp, partition.preferredReadReplica(), expireTimeMs);
-                return expireTimeMs;
-            });
+        if (FetchResponse.isPreferredReplica(partitionData)) {
+            return subscriptions.tryUpdatingPreferredReadReplica(
+                tp, partitionData.preferredReadReplica(), () -> {
+                    long expireTimeMs = time.milliseconds() + metadata.metadataExpireMs();
+                    log.debug("Updating preferred read replica for partition {} to {}, set to expire at {}",
+                        tp, partitionData.preferredReadReplica(), expireTimeMs);
+                    return expireTimeMs;
+                });
         }
 
-        completedFetch.setInitialized();
-        return completedFetch;
+        return true;
     }
 
     private void handleInitializeErrors(final CompletedFetch completedFetch, final Errors error) {
@@ -331,17 +347,17 @@ private void handleInitializeErrors(final CompletedFetch completedFetch, final E
 
             if (!clearedReplicaId.isPresent()) {
                 // If there's no preferred replica to clear, we're fetching from the leader so handle this error normally
-                SubscriptionState.FetchPosition position = subscriptions.position(tp);
+                SubscriptionState.FetchPosition position = subscriptions.positionOrNull(tp);
 
                 if (position == null || fetchOffset != position.offset) {
                     log.debug("Discarding stale fetch response for partition {} since the fetched offset {} " +
-                            "does not match the current offset {}", tp, fetchOffset, position);
+                            "does not match the current offset {} or the partition has been unassigned", tp, fetchOffset, position);
                 } else {
                     String errorMessage = "Fetch position " + position + " is out of range for partition " + tp;
 
                     if (subscriptions.hasDefaultOffsetResetPolicy()) {
                         log.info("{}, resetting offset", errorMessage);
-                        subscriptions.requestOffsetReset(tp);
+                        subscriptions.requestOffsetResetIfPartitionAssigned(tp);
                     } else {
                         log.info("{}, raising error to the application since no reset policy is configured", errorMessage);
                         throw new OffsetOutOfRangeException(errorMessage,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
index d466a4a1b29dd..87d471d2e150d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
@@ -540,6 +540,14 @@ public synchronized FetchPosition position(TopicPartition tp) {
         return assignedState(tp).position;
     }
 
+    public synchronized FetchPosition positionOrNull(TopicPartition tp) {
+        final TopicPartitionState state = assignedStateOrNull(tp);
+        if (state == null) {
+            return null;
+        }
+        return assignedState(tp).position;
+    }
+
     public synchronized Long partitionLag(TopicPartition tp, IsolationLevel isolationLevel) {
         TopicPartitionState topicPartitionState = assignedState(tp);
         if (topicPartitionState.position == null) {
@@ -579,14 +587,37 @@ synchronized void updateHighWatermark(TopicPartition tp, long highWatermark) {
         assignedState(tp).highWatermark(highWatermark);
     }
 
-    synchronized void updateLogStartOffset(TopicPartition tp, long logStartOffset) {
-        assignedState(tp).logStartOffset(logStartOffset);
+    synchronized boolean tryUpdatingHighWatermark(TopicPartition tp, long highWatermark) {
+        final TopicPartitionState state = assignedStateOrNull(tp);
+        if (state != null) {
+            assignedState(tp).highWatermark(highWatermark);
+            return true;
+        }
+        return false;
+    }
+
+    synchronized boolean tryUpdatingLogStartOffset(TopicPartition tp, long highWatermark) {
+        final TopicPartitionState state = assignedStateOrNull(tp);
+        if (state != null) {
+            assignedState(tp).logStartOffset(highWatermark);
+            return true;
+        }
+        return false;
     }
 
     synchronized void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
         assignedState(tp).lastStableOffset(lastStableOffset);
     }
 
+    synchronized boolean tryUpdatingLastStableOffset(TopicPartition tp, long lastStableOffset) {
+        final TopicPartitionState state = assignedStateOrNull(tp);
+        if (state != null) {
+            assignedState(tp).lastStableOffset(lastStableOffset);
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Set the preferred read replica with a lease timeout. After this time, the replica will no longer be valid and
      * {@link #preferredReadReplica(TopicPartition, long)} will return an empty result.
@@ -599,6 +630,28 @@ public synchronized void updatePreferredReadReplica(TopicPartition tp, int prefe
         assignedState(tp).updatePreferredReadReplica(preferredReadReplicaId, timeMs);
     }
 
+    /**
+     * Tries to set the preferred read replica with a lease timeout. After this time, the replica will no longer be valid and
+     * {@link #preferredReadReplica(TopicPartition, long)} will return an empty result. If the preferred replica of
+     * the partition could not be updated (e.g. because the partition is not assigned) this method will return
+     * {@code false}, otherwise it will return {@code true}.
+     *
+     * @param tp The topic partition
+     * @param preferredReadReplicaId The preferred read replica
+     * @param timeMs The time at which this preferred replica is no longer valid
+     * @return {@code true} if the preferred read replica was updated, {@code false} otherwise.
+     */
+    public synchronized boolean tryUpdatingPreferredReadReplica(TopicPartition tp,
+                                                             int preferredReadReplicaId,
+                                                             LongSupplier timeMs) {
+        final TopicPartitionState state = assignedStateOrNull(tp);
+        if (state != null) {
+            assignedState(tp).updatePreferredReadReplica(preferredReadReplicaId, timeMs);
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Get the preferred read replica
      *
@@ -655,6 +708,14 @@ public void requestOffsetReset(TopicPartition partition) {
         requestOffsetReset(partition, defaultResetStrategy);
     }
 
+    public synchronized void requestOffsetResetIfPartitionAssigned(TopicPartition partition) {
+        final TopicPartitionState state = assignedStateOrNull(partition);
+        if (state != null) {
+            state.reset(defaultResetStrategy);
+        }
+    }
+
+
     synchronized void setNextAllowedRetry(Set<TopicPartition> partitions, long nextAllowResetTimeMs) {
         for (TopicPartition partition : partitions) {
             assignedState(partition).setNextAllowedRetry(nextAllowResetTimeMs);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchCollectorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchCollectorTest.java
index 0ca4a18a48ae5..094e346e02bdf 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchCollectorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchCollectorTest.java
@@ -40,6 +40,8 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -61,11 +63,17 @@
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 /**
  * This tests the {@link FetchCollector} functionality in addition to what {@link FetcherTest} tests during the course
  * of its tests.
  */
+@MockitoSettings(strictness = Strictness.STRICT_STUBS)
 public class FetchCollectorTest {
 
     private final static int DEFAULT_RECORD_COUNT = 10;
@@ -75,7 +83,7 @@ public class FetchCollectorTest {
     private final TopicPartition topicAPartition1 = new TopicPartition("topic-a", 1);
     private final TopicPartition topicAPartition2 = new TopicPartition("topic-a", 2);
     private final Set<TopicPartition> allPartitions = partitions(topicAPartition0, topicAPartition1, topicAPartition2);
-    private LogContext logContext;
+    private final LogContext logContext = new LogContext();
 
     private SubscriptionState subscriptions;
     private FetchConfig fetchConfig;
@@ -406,6 +414,228 @@ public void testFetchWithOtherErrors(final Errors error) {
         assertThrows(IllegalStateException.class, () -> fetchCollector.collectFetch(fetchBuffer));
     }
 
+    @Test
+    public void testCollectFetchInitializationWithNullPosition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(null);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        final Records records = createRecords();
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setHighWatermark(1000)
+            .setRecords(records);
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationWithUpdateHighWatermarkOnNotAssignedPartition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final long fetchOffset = 42;
+        final long highWatermark = 1000;
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(new SubscriptionState.FetchPosition(fetchOffset));
+        when(subscriptions.tryUpdatingHighWatermark(topicPartition0, highWatermark)).thenReturn(false);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        final Records records = createRecords();
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setHighWatermark(highWatermark)
+            .setRecords(records);
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0)
+            .fetchOffset(fetchOffset).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationWithUpdateLogStartOffsetOnNotAssignedPartition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final long fetchOffset = 42;
+        final long highWatermark = 1000;
+        final long logStartOffset = 10;
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(new SubscriptionState.FetchPosition(fetchOffset));
+        when(subscriptions.tryUpdatingHighWatermark(topicPartition0, highWatermark)).thenReturn(true);
+        when(subscriptions.tryUpdatingLogStartOffset(topicPartition0, logStartOffset)).thenReturn(false);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        final Records records = createRecords();
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setHighWatermark(highWatermark)
+            .setRecords(records)
+            .setLogStartOffset(logStartOffset);
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0)
+            .fetchOffset(fetchOffset).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationWithUpdateLastStableOffsetOnNotAssignedPartition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final long fetchOffset = 42;
+        final long highWatermark = 1000;
+        final long logStartOffset = 10;
+        final long lastStableOffset = 900;
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(new SubscriptionState.FetchPosition(fetchOffset));
+        when(subscriptions.tryUpdatingHighWatermark(topicPartition0, highWatermark)).thenReturn(true);
+        when(subscriptions.tryUpdatingLogStartOffset(topicPartition0, logStartOffset)).thenReturn(true);
+        when(subscriptions.tryUpdatingLastStableOffset(topicPartition0, lastStableOffset)).thenReturn(false);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        final Records records = createRecords();
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setHighWatermark(highWatermark)
+            .setRecords(records)
+            .setLogStartOffset(logStartOffset)
+            .setLastStableOffset(lastStableOffset);
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0)
+            .fetchOffset(fetchOffset).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationWithUpdatePreferredReplicaOnNotAssignedPartition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final long fetchOffset = 42;
+        final long highWatermark = 1000;
+        final long logStartOffset = 10;
+        final long lastStableOffset = 900;
+        final int preferredReadReplicaId = 21;
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(new SubscriptionState.FetchPosition(fetchOffset));
+        when(subscriptions.tryUpdatingHighWatermark(topicPartition0, highWatermark)).thenReturn(true);
+        when(subscriptions.tryUpdatingLogStartOffset(topicPartition0, logStartOffset)).thenReturn(true);
+        when(subscriptions.tryUpdatingLastStableOffset(topicPartition0, lastStableOffset)).thenReturn(true);
+        when(subscriptions.tryUpdatingPreferredReadReplica(eq(topicPartition0), eq(preferredReadReplicaId), any())).thenReturn(false);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        final Records records = createRecords();
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setHighWatermark(highWatermark)
+            .setRecords(records)
+            .setLogStartOffset(logStartOffset)
+            .setLastStableOffset(lastStableOffset)
+            .setPreferredReadReplica(preferredReadReplicaId);
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0)
+            .fetchOffset(fetchOffset).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationOffsetOutOfRangeErrorWithNullPosition() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(null);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setErrorCode(Errors.OFFSET_OUT_OF_RANGE.code());
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    @Test
+    public void testCollectFetchInitializationOffsetOutOfRangeErrorWithOffsetReset() {
+        final TopicPartition topicPartition0 = new TopicPartition("topic", 0);
+        final long fetchOffset = 42;
+        final SubscriptionState subscriptions = mock(SubscriptionState.class);
+        when(subscriptions.hasValidPosition(topicPartition0)).thenReturn(true);
+        when(subscriptions.positionOrNull(topicPartition0)).thenReturn(new SubscriptionState.FetchPosition(fetchOffset));
+        when(subscriptions.hasDefaultOffsetResetPolicy()).thenReturn(true);
+        final FetchCollector<String, String> fetchCollector = createFetchCollector(subscriptions);
+        FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            .setPartitionIndex(topicPartition0.partition())
+            .setErrorCode(Errors.OFFSET_OUT_OF_RANGE.code());
+        final CompletedFetch completedFetch = new CompletedFetchBuilder()
+            .partitionData(partitionData)
+            .partition(topicPartition0)
+            .fetchOffset(fetchOffset).build();
+        final FetchBuffer fetchBuffer = mock(FetchBuffer.class);
+        when(fetchBuffer.nextInLineFetch()).thenReturn(null);
+        when(fetchBuffer.peek()).thenReturn(completedFetch).thenReturn(null);
+
+        final Fetch<String, String> fetch = fetchCollector.collectFetch(fetchBuffer);
+
+        assertTrue(fetch.isEmpty());
+        verify(subscriptions).requestOffsetResetIfPartitionAssigned(topicPartition0);
+        verify(fetchBuffer).setNextInLineFetch(null);
+    }
+
+    private FetchCollector<String, String> createFetchCollector(final SubscriptionState subscriptions) {
+        final Properties consumerProps = consumerProps();
+        return new FetchCollector<>(
+            logContext,
+            mock(ConsumerMetadata.class),
+            subscriptions,
+            new FetchConfig(new ConsumerConfig(consumerProps)),
+            new Deserializers<>(new StringDeserializer(), new StringDeserializer()),
+            mock(FetchMetricsManager.class),
+            new MockTime()
+        );
+    }
+
     /**
      * This is a handy utility method for returning a set from a varargs array.
      */
@@ -418,14 +648,7 @@ private void buildDependencies() {
     }
 
     private void buildDependencies(int maxPollRecords) {
-        logContext = new LogContext();
-
-        Properties p = new Properties();
-        p.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
-        p.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
-        p.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
-        p.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, String.valueOf(maxPollRecords));
-
+        Properties p = consumerProperties(maxPollRecords);
         ConsumerConfig config = new ConsumerConfig(p);
 
         deserializers = new Deserializers<>(new StringDeserializer(), new StringDeserializer());
@@ -456,6 +679,20 @@ private void buildDependencies(int maxPollRecords) {
         completedFetchBuilder = new CompletedFetchBuilder();
     }
 
+    private Properties consumerProps() {
+        return consumerProperties(DEFAULT_MAX_POLL_RECORDS);
+    }
+
+    private Properties consumerProperties(final int maxPollRecords) {
+        Properties p = new Properties();
+        p.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        p.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
+        p.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
+        p.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, String.valueOf(maxPollRecords));
+
+        return p;
+    }
+
     private void assign(TopicPartition... partitions) {
         subscriptions.assignFromUser(partitions(partitions));
     }
@@ -527,6 +764,10 @@ private class CompletedFetchBuilder {
 
         private int recordCount = DEFAULT_RECORD_COUNT;
 
+        private TopicPartition topicPartition = topicAPartition0;
+
+        private FetchResponseData.PartitionData partitionData = null;
+
         private Errors error = null;
 
         private CompletedFetchBuilder fetchOffset(long fetchOffset) {
@@ -544,24 +785,29 @@ private CompletedFetchBuilder error(Errors error) {
             return this;
         }
 
-        private CompletedFetch build() {
-            Records records;
-            ByteBuffer allocate = ByteBuffer.allocate(1024);
+        private CompletedFetchBuilder partitionData(FetchResponseData.PartitionData partitionData) {
+            this.partitionData = partitionData;
+            return this;
+        }
 
-            try (MemoryRecordsBuilder builder = MemoryRecords.builder(allocate,
-                    CompressionType.NONE,
-                    TimestampType.CREATE_TIME,
-                    0)) {
-                for (int i = 0; i < recordCount; i++)
-                    builder.append(0L, "key".getBytes(), ("value-" + i).getBytes());
+        private CompletedFetchBuilder partition(TopicPartition topicPartition) {
+            this.topicPartition = topicPartition;
+            return this;
+        }
 
-                records = builder.build();
-            }
+        private CompletedFetch build() {
+            Records records = createRecords(recordCount);
 
-            FetchResponseData.PartitionData partitionData = new FetchResponseData.PartitionData()
+            if (partitionData == null) {
+                partitionData = new FetchResponseData.PartitionData()
                     .setPartitionIndex(topicAPartition0.partition())
                     .setHighWatermark(1000)
                     .setRecords(records);
+            }
+
+            if (topicPartition != null) {
+                partitionData.setPartitionIndex(topicPartition.partition());
+            }
 
             if (error != null)
                 partitionData.setErrorCode(error.code());
@@ -571,11 +817,29 @@ private CompletedFetch build() {
                     logContext,
                     subscriptions,
                     BufferSupplier.create(),
-                    topicAPartition0,
+                    topicPartition,
                     partitionData,
                     metricsAggregator,
                     fetchOffset,
                     ApiKeys.FETCH.latestVersion());
         }
     }
+
+    private Records createRecords() {
+        return createRecords(DEFAULT_RECORD_COUNT);
+    }
+
+    private Records createRecords(final int recordCount) {
+        ByteBuffer allocate = ByteBuffer.allocate(1024);
+
+        try (MemoryRecordsBuilder builder = MemoryRecords.builder(allocate,
+            CompressionType.NONE,
+            TimestampType.CREATE_TIME,
+            0)) {
+            for (int i = 0; i < recordCount; i++)
+                builder.append(0L, "key".getBytes(), ("value-" + i).getBytes());
+
+            return builder.build();
+        }
+    }
 }
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
index fe6cc5f028171..e54bb36e179a6 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
@@ -39,6 +39,7 @@
 import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.LongSupplier;
 import java.util.regex.Pattern;
 
 import static java.util.Collections.singleton;
@@ -895,4 +896,75 @@ public void nullPositionLagOnNoPosition() {
         assertNull(state.partitionLag(tp0, IsolationLevel.READ_COMMITTED));
     }
 
+    @Test
+    public void testPositionOrNull() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+        state.seek(tp0, 5);
+
+        assertEquals(5, state.positionOrNull(tp0).offset);
+        assertNull(state.positionOrNull(unassignedPartition));
+    }
+
+    @Test
+    public void testTryUpdatingHighWatermark() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+
+        final long highWatermark = 10L;
+        assertTrue(state.tryUpdatingHighWatermark(tp0, highWatermark));
+        assertEquals(highWatermark, state.partitionEndOffset(tp0, IsolationLevel.READ_UNCOMMITTED));
+        assertFalse(state.tryUpdatingHighWatermark(unassignedPartition, highWatermark));
+    }
+
+    @Test
+    public void testTryUpdatingLogStartOffset() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+        final long position = 25;
+        state.seek(tp0, position);
+
+        final long logStartOffset = 10L;
+        assertTrue(state.tryUpdatingLogStartOffset(tp0, logStartOffset));
+        assertEquals(position - logStartOffset, state.partitionLead(tp0));
+        assertFalse(state.tryUpdatingLogStartOffset(unassignedPartition, logStartOffset));
+    }
+
+    @Test
+    public void testTryUpdatingLastStableOffset() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+
+        final long lastStableOffset = 10L;
+        assertTrue(state.tryUpdatingLastStableOffset(tp0, lastStableOffset));
+        assertEquals(lastStableOffset, state.partitionEndOffset(tp0, IsolationLevel.READ_COMMITTED));
+        assertFalse(state.tryUpdatingLastStableOffset(unassignedPartition, lastStableOffset));
+    }
+
+    @Test
+    public void testTryUpdatingPreferredReadReplica() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+
+        final int preferredReadReplicaId = 10;
+        final LongSupplier expirationTimeMs = () -> System.currentTimeMillis() + 60000L;
+        assertTrue(state.tryUpdatingPreferredReadReplica(tp0, preferredReadReplicaId, expirationTimeMs));
+        assertEquals(Optional.of(preferredReadReplicaId), state.preferredReadReplica(tp0, System.currentTimeMillis()));
+        assertFalse(state.tryUpdatingPreferredReadReplica(unassignedPartition, preferredReadReplicaId, expirationTimeMs));
+        assertEquals(Optional.empty(), state.preferredReadReplica(unassignedPartition, System.currentTimeMillis()));
+    }
+
+    @Test
+    public void testRequestOffsetResetIfPartitionAssigned() {
+        state.assignFromUser(Collections.singleton(tp0));
+        final TopicPartition unassignedPartition = new TopicPartition("unassigned", 0);
+
+        state.requestOffsetResetIfPartitionAssigned(tp0);
+
+        assertTrue(state.isOffsetResetNeeded(tp0));
+
+        state.requestOffsetResetIfPartitionAssigned(unassignedPartition);
+
+        assertThrows(IllegalStateException.class, () -> state.isOffsetResetNeeded(unassignedPartition));
+    }
 }

From 2c613b2d42fc3d6f7dfd89d23c08ef51dd90ba4b Mon Sep 17 00:00:00 2001
From: Cheryl Simmons <csimmons@confluent.io>
Date: Tue, 12 Mar 2024 15:21:43 -0700
Subject: [PATCH 145/258] MINOR: Tweak streams config doc (#15518)

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../main/java/org/apache/kafka/streams/StreamsConfig.java    | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index c166b3a43f04f..9500e315800e2 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -513,13 +513,14 @@ public class StreamsConfig extends AbstractConfig {
     /** {@code client.id} */
     @SuppressWarnings("WeakerAccess")
     public static final String CLIENT_ID_CONFIG = CommonClientConfigs.CLIENT_ID_CONFIG;
-    private static final String CLIENT_ID_DOC = "An ID prefix string used for the client IDs of internal [main-|restore-|global-]consumer, producer, and admin clients" +
+    private static final String CLIENT_ID_DOC = "An ID prefix string used for the client IDs of internal (main, restore, and global) consumers , producers, and admin clients" +
         " with pattern <code>&lt;client.id&gt;-[Global]StreamThread[-&lt;threadSequenceNumber&gt;]-&lt;consumer|producer|restore-consumer|global-consumer&gt;</code>.";
 
     /** {@code enable.metrics.push} */
     @SuppressWarnings("WeakerAccess")
     public static  final String ENABLE_METRICS_PUSH_CONFIG = CommonClientConfigs.ENABLE_METRICS_PUSH_CONFIG;
-    public static final String ENABLE_METRICS_PUSH_DOC = "Whether to enable pushing of internal [main-|restore-|global-]consumer, producer, and admin client metrics to the cluster, if the cluster has a client metrics subscription which matches a client.";
+    public static final String ENABLE_METRICS_PUSH_DOC = "Whether to enable pushing of internal client metrics for (main, restore, and global) consumers, producers, and admin clients." + 
+        " The cluster must have a client metrics subscription which corresponds to a client.";
 
     /** {@code commit.interval.ms} */
     @SuppressWarnings("WeakerAccess")

From aa7bef414e2e08fbe35923979e0c8fc8438e594e Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Wed, 13 Mar 2024 11:45:38 -0700
Subject: [PATCH 146/258] MINOR: Resolve SSLContextFactory.getNeedClientAuth
 deprecation (#15468)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 .../kafka/connect/runtime/rest/RestClient.java      |  4 ++--
 .../kafka/connect/runtime/rest/RestServer.java      |  2 +-
 .../kafka/connect/runtime/rest/util/SSLUtils.java   |  8 ++++----
 .../integration/RestForwardingIntegrationTest.java  |  2 +-
 .../connect/runtime/rest/util/SSLUtilsTest.java     | 13 ++++---------
 5 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestClient.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestClient.java
index 8558009ee9dd7..9a47a0e7530bb 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestClient.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestClient.java
@@ -61,7 +61,7 @@ public RestClient(AbstractConfig config) {
     }
 
     // VisibleForTesting
-    HttpClient httpClient(SslContextFactory sslContextFactory) {
+    HttpClient httpClient(SslContextFactory.Client sslContextFactory) {
         return sslContextFactory != null ? new HttpClient(sslContextFactory) : new HttpClient();
     }
 
@@ -120,7 +120,7 @@ public <T> HttpResponse<T> httpRequest(String url, String method, HttpHeaders he
         Objects.requireNonNull(method, "method must be non-null");
         Objects.requireNonNull(responseFormat, "response format must be non-null");
         // Only try to load SSL configs if we have to (see KAFKA-14816)
-        SslContextFactory sslContextFactory = url.startsWith("https://")
+        SslContextFactory.Client sslContextFactory = url.startsWith("https://")
                 ? SSLUtils.createClientSideSslContextFactory(config)
                 : null;
         HttpClient client = httpClient(sslContextFactory);
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServer.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServer.java
index f078b24420e51..b6bb4fa3bd157 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServer.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServer.java
@@ -159,7 +159,7 @@ public final Connector createConnector(String listener, boolean isAdmin) {
         ServerConnector connector;
 
         if (PROTOCOL_HTTPS.equals(protocol)) {
-            SslContextFactory ssl;
+            SslContextFactory.Server ssl;
             if (isAdmin) {
                 ssl = SSLUtils.createServerSideSslContextFactory(config, RestServerConfig.ADMIN_LISTENERS_HTTPS_CONFIGS_PREFIX);
             } else {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/util/SSLUtils.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/util/SSLUtils.java
index efb184849f859..341e3f3550b81 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/util/SSLUtils.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/util/SSLUtils.java
@@ -40,7 +40,7 @@ public class SSLUtils {
     /**
      * Configures SSL/TLS for HTTPS Jetty Server using configs with the given prefix
      */
-    public static SslContextFactory createServerSideSslContextFactory(AbstractConfig config, String prefix) {
+    public static SslContextFactory.Server createServerSideSslContextFactory(AbstractConfig config, String prefix) {
         Map<String, Object> sslConfigValues = config.valuesWithPrefixAllOrNothing(prefix);
 
         final SslContextFactory.Server ssl = new SslContextFactory.Server();
@@ -56,14 +56,14 @@ public static SslContextFactory createServerSideSslContextFactory(AbstractConfig
     /**
      * Configures SSL/TLS for HTTPS Jetty Server
      */
-    public static SslContextFactory createServerSideSslContextFactory(AbstractConfig config) {
+    public static SslContextFactory.Server createServerSideSslContextFactory(AbstractConfig config) {
         return createServerSideSslContextFactory(config, "listeners.https.");
     }
 
     /**
      * Configures SSL/TLS for HTTPS Jetty Client
      */
-    public static SslContextFactory createClientSideSslContextFactory(AbstractConfig config) {
+    public static SslContextFactory.Client createClientSideSslContextFactory(AbstractConfig config) {
         Map<String, Object> sslConfigValues = config.valuesWithPrefixAllOrNothing("listeners.https.");
 
         final SslContextFactory.Client ssl = new SslContextFactory.Client();
@@ -147,7 +147,7 @@ protected static void configureSslContextFactoryAlgorithms(SslContextFactory ssl
     /**
      * Configures hostname verification related settings in SslContextFactory
      */
-    protected static void configureSslContextFactoryEndpointIdentification(SslContextFactory ssl, Map<String, Object> sslConfigValues) {
+    protected static void configureSslContextFactoryEndpointIdentification(SslContextFactory.Client ssl, Map<String, Object> sslConfigValues) {
         String sslEndpointIdentificationAlg = (String) sslConfigValues.get(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
         if (sslEndpointIdentificationAlg != null)
             ssl.setEndpointIdentificationAlgorithm(sslEndpointIdentificationAlg);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestForwardingIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestForwardingIntegrationTest.java
index 0bbad10a57d91..2e63a2518b345 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestForwardingIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RestForwardingIntegrationTest.java
@@ -89,7 +89,7 @@ public class RestForwardingIntegrationTest {
     @Mock
     private Herder leaderHerder;
 
-    private SslContextFactory factory;
+    private SslContextFactory.Client factory;
     private CloseableHttpClient httpClient;
     private Collection<CloseableHttpResponse> responses;
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/util/SSLUtilsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/util/SSLUtilsTest.java
index 8d486bbb9a31b..8f5c3235e7e6a 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/util/SSLUtilsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/util/SSLUtilsTest.java
@@ -26,7 +26,6 @@
 import java.util.HashMap;
 import java.util.Map;
 
-@SuppressWarnings("deprecation")
 public class SSLUtilsTest {
 
     @Test
@@ -63,7 +62,7 @@ public void testCreateServerSideSslContextFactory() {
         configMap.put("ssl.trustmanager.algorithm", "PKIX");
 
         RestServerConfig config = RestServerConfig.forPublic(null, configMap);
-        SslContextFactory ssl = SSLUtils.createServerSideSslContextFactory(config);
+        SslContextFactory.Server ssl = SSLUtils.createServerSideSslContextFactory(config);
 
         Assert.assertEquals("file:///path/to/keystore", ssl.getKeyStorePath());
         Assert.assertEquals("file:///path/to/truststore", ssl.getTrustStorePath());
@@ -101,15 +100,13 @@ public void testCreateClientSideSslContextFactory() {
         configMap.put("ssl.trustmanager.algorithm", "PKIX");
 
         RestServerConfig config = RestServerConfig.forPublic(null, configMap);
-        SslContextFactory ssl = SSLUtils.createClientSideSslContextFactory(config);
+        SslContextFactory.Client ssl = SSLUtils.createClientSideSslContextFactory(config);
 
         Assert.assertEquals("file:///path/to/keystore", ssl.getKeyStorePath());
         Assert.assertEquals("file:///path/to/truststore", ssl.getTrustStorePath());
         Assert.assertEquals("SunJSSE", ssl.getProvider());
         Assert.assertArrayEquals(new String[] {"SSL_RSA_WITH_RC4_128_SHA", "SSL_RSA_WITH_RC4_128_MD5"}, ssl.getIncludeCipherSuites());
         Assert.assertEquals("SHA1PRNG", ssl.getSecureRandomAlgorithm());
-        Assert.assertFalse(ssl.getNeedClientAuth());
-        Assert.assertFalse(ssl.getWantClientAuth());
         Assert.assertEquals("JKS", ssl.getKeyStoreType());
         Assert.assertEquals("JKS", ssl.getTrustStoreType());
         Assert.assertEquals("TLS", ssl.getProtocol());
@@ -131,7 +128,7 @@ public void testCreateServerSideSslContextFactoryDefaultValues() {
         configMap.put("ssl.secure.random.implementation", "SHA1PRNG");
 
         RestServerConfig config = RestServerConfig.forPublic(null, configMap);
-        SslContextFactory ssl = SSLUtils.createServerSideSslContextFactory(config);
+        SslContextFactory.Server ssl = SSLUtils.createServerSideSslContextFactory(config);
 
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE, ssl.getKeyStoreType());
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ssl.getTrustStoreType());
@@ -156,7 +153,7 @@ public void testCreateClientSideSslContextFactoryDefaultValues() {
         configMap.put("ssl.secure.random.implementation", "SHA1PRNG");
 
         RestServerConfig config = RestServerConfig.forPublic(null, configMap);
-        SslContextFactory ssl = SSLUtils.createClientSideSslContextFactory(config);
+        SslContextFactory.Client ssl = SSLUtils.createClientSideSslContextFactory(config);
 
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE, ssl.getKeyStoreType());
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ssl.getTrustStoreType());
@@ -164,7 +161,5 @@ public void testCreateClientSideSslContextFactoryDefaultValues() {
         Assert.assertArrayEquals(Arrays.asList(SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS.split("\\s*,\\s*")).toArray(), ssl.getIncludeProtocols());
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_KEYMANGER_ALGORITHM, ssl.getKeyManagerFactoryAlgorithm());
         Assert.assertEquals(SslConfigs.DEFAULT_SSL_TRUSTMANAGER_ALGORITHM, ssl.getTrustManagerFactoryAlgorithm());
-        Assert.assertFalse(ssl.getNeedClientAuth());
-        Assert.assertFalse(ssl.getWantClientAuth());
     }
 }

From 722967a2b71dac13b27b6e59638f853ee18b830b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?=
 <jsancio@users.noreply.github.com>
Date: Wed, 13 Mar 2024 14:36:03 -0700
Subject: [PATCH 147/258] MINOR; Make string from array (#15526)

If toString is called on an array it returns the string representing the object reference.  Use mkString instead to print the content of the array.

Reviewers: Luke Chen <showuon@gmail.com>, Justine Olshan <jolshan@confluent.io>, Lingnan Liu <liliu@confluent.io>
---
 core/src/main/scala/kafka/log/LogManager.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala
index 7c8bbec5292e9..53e0b7a43e3bb 100755
--- a/core/src/main/scala/kafka/log/LogManager.scala
+++ b/core/src/main/scala/kafka/log/LogManager.scala
@@ -1564,7 +1564,7 @@ object LogManager {
       Option(newTopicsImage.getPartition(topicId, partitionId)) match {
         case Some(partition) =>
           if (!partition.replicas.contains(brokerId)) {
-            info(s"Found stray log dir $log: the current replica assignment ${partition.replicas} " +
+            info(s"Found stray log dir $log: the current replica assignment ${partition.replicas.mkString("[", ", ", "]")} " +
               s"does not contain the local brokerId $brokerId.")
             Some(log.topicPartition)
           } else {

From d88a97adef684a0f5403c46f7fb2f8d1723eebd5 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Wed, 13 Mar 2024 21:52:25 -0700
Subject: [PATCH 148/258] MINOR: simplify consumer logic (#15519)

For static member, the `group.instance.id` cannot change.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>, Lianet Magrans <lianetmr@gmail.com>, David Jacot <david.jacot@gmail.com>
---
 .../consumer/internals/HeartbeatRequestManager.java  | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index d2a7205d5f3ff..19e4a8b7132ba 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -529,14 +529,8 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
             // MemberEpoch - always sent
             data.setMemberEpoch(membershipManager.memberEpoch());
 
-            // InstanceId - only sent if has changed since the last heartbeat
-            // Always send when leaving the group as a static member
-            membershipManager.groupInstanceId().ifPresent(groupInstanceId -> {
-                if (!groupInstanceId.equals(sentFields.instanceId) || membershipManager.memberEpoch() == ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH) {
-                    data.setInstanceId(groupInstanceId);
-                    sentFields.instanceId = groupInstanceId;
-                }
-            });
+            // InstanceId - set if present
+            membershipManager.groupInstanceId().ifPresent(data::setInstanceId);
 
             // RebalanceTimeoutMs - only sent if has changed since the last heartbeat
             if (sentFields.rebalanceTimeoutMs != rebalanceTimeoutMs) {
@@ -593,7 +587,6 @@ private List<ConsumerGroupHeartbeatRequestData.TopicPartitions> buildTopicPartit
 
         // Fields of ConsumerHeartbeatRequest sent in the most recent request
         static class SentFields {
-            private String instanceId = null;
             private int rebalanceTimeoutMs = -1;
             private TreeSet<String> subscribedTopicNames = null;
             private String serverAssignor = null;
@@ -602,7 +595,6 @@ static class SentFields {
             SentFields() {}
 
             void reset() {
-                instanceId = null;
                 rebalanceTimeoutMs = -1;
                 subscribedTopicNames = null;
                 serverAssignor = null;

From 612a1fe1bbe0b6bdc3df4648ec3427946febc2f5 Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Wed, 13 Mar 2024 21:54:06 -0700
Subject: [PATCH 149/258] MINOR: Kafka Streams docs fixes (#15517)

- add missing section to TOC
- add default value for client.id

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>, Bruno Cadonna <bruno@confluent.io>
---
 docs/streams/developer-guide/config-streams.html               | 1 +
 .../src/main/java/org/apache/kafka/streams/StreamsConfig.java  | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/docs/streams/developer-guide/config-streams.html b/docs/streams/developer-guide/config-streams.html
index 288946d33a055..40b359488022d 100644
--- a/docs/streams/developer-guide/config-streams.html
+++ b/docs/streams/developer-guide/config-streams.html
@@ -97,6 +97,7 @@
             <ul>
               <li><a class="reference internal" href="#naming" id="id17">Naming</a></li>
               <li><a class="reference internal" href="#default-values" id="id18">Default Values</a></li>
+              <li><a class="reference internal" href="#parameters-controlled-by-kafka-streams" id="id26">Parameters controlled by Kafka Streams</a></li>
               <li><a class="reference internal" href="#enable-auto-commit" id="id19">enable.auto.commit</a></li>
             </ul>
           </li>
diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
index 9500e315800e2..44c064f40fbd3 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java
@@ -885,7 +885,8 @@ public class StreamsConfig extends AbstractConfig {
                     Type.STRING,
                     "",
                     Importance.MEDIUM,
-                    CLIENT_ID_DOC)
+                    CLIENT_ID_DOC,
+                    "<code>&lt;application.id&gt-&lt;random-UUID&gt</code>")
             .define(DEFAULT_DESERIALIZATION_EXCEPTION_HANDLER_CLASS_CONFIG,
                     Type.CLASS,
                     LogAndFailExceptionHandler.class.getName(),

From e164d4d4264b6d663952c6a65f2fda46a4cc9e04 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Thu, 14 Mar 2024 00:54:28 -0700
Subject: [PATCH 150/258] KAFKA-16249; Improve reconciliation state machine
 (#15364)

This patch re-work the reconciliation state machine on the server side with the goal to fix a few issues that we have recently discovered.
* When a member acknowledges the revocation of partitions (by not reporting them in the heartbeat), the current implementation may miss it. The issue is that the current implementation re-compute the assignment of a member whenever there is a new target assignment installed. When it happens, it does not consider the reported owned partitions at all. As the member is supposed to only report its own partitions when they change, the member is stuck.
* Similarly, as the current assignment is re-computed whenever there is a new target assignment, the rebalance timeout, as it is currently implemented, becomes useless. The issue is that the rebalance timeout is reset whenever the member enters the revocation state. In other words, in the current implementation, the timer is reset when there are no target available even if the previous revocation is not completed yet.

The patch fixes these two issues by not automatically recomputing the assignment of a member when a new target assignment is available. When the member must revoke partitions, the coordinator waits. Otherwise, it recomputes the next assignment. In other words, revoking is really blocking now.

The patch also proposes to include an explicit state in the record. It makes the implementation cleaner and it also makes it more extensible in the future.

The patch also changes the record format. This is a non-backward compatible change. I think that we should do this change to cleanup the record. As KIP-848 is only in early access in 3.7 and that we clearly state that we don't plane to support upgrade from it, this is acceptable in my opinion.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../group/GroupMetadataManager.java           | 164 +++--
 .../coordinator/group/RecordHelpers.java      |   5 +-
 .../apache/kafka/coordinator/group/Utils.java |  32 +
 .../group/consumer/ConsumerGroup.java         |   2 +-
 .../group/consumer/ConsumerGroupMember.java   | 149 ++--
 .../consumer/CurrentAssignmentBuilder.java    | 365 +++-------
 .../group/consumer/MemberState.java           |  76 ++
 ...umerGroupCurrentMemberAssignmentValue.json |  14 +-
 .../kafka/coordinator/group/Assertions.java   |   6 -
 .../group/GroupMetadataManagerTest.java       | 511 +++----------
 .../GroupMetadataManagerTestContext.java      |  14 +-
 .../group/OffsetMetadataManagerTest.java      |   6 -
 .../coordinator/group/RecordHelpersTest.java  |  20 +-
 .../consumer/ConsumerGroupMemberTest.java     |  28 +-
 .../group/consumer/ConsumerGroupTest.java     |  33 +-
 .../CurrentAssignmentBuilderTest.java         | 674 +++++++++---------
 .../GroupCoordinatorMetricsShardTest.java     |   2 +-
 17 files changed, 823 insertions(+), 1278 deletions(-)
 create mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 30d4d52c6ff84..26c2f644e9612 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -58,6 +58,7 @@
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.consumer.CurrentAssignmentBuilder;
+import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.consumer.TargetAssignmentBuilder;
 import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
@@ -98,6 +99,7 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -119,6 +121,7 @@
 import static org.apache.kafka.coordinator.group.RecordHelpers.newMemberSubscriptionRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newMemberSubscriptionTombstoneRecord;
 import static org.apache.kafka.coordinator.group.RecordHelpers.newTargetAssignmentTombstoneRecord;
+import static org.apache.kafka.coordinator.group.Utils.assignmentToString;
 import static org.apache.kafka.coordinator.group.Utils.ofSentinel;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupMember.EMPTY_ASSIGNMENT;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.COMPLETING_REBALANCE;
@@ -126,6 +129,7 @@
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.EMPTY;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.PREPARING_REBALANCE;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.STABLE;
+import static org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember.hasAssignedPartitionsChanged;
 import static org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetrics.CLASSIC_GROUP_COMPLETED_REBALANCES_SENSOR_NAME;
 import static org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetrics.CONSUMER_GROUP_REBALANCES_SENSOR_NAME;
 
@@ -1058,7 +1062,6 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
             }
         }
 
-
         int groupEpoch = group.groupEpoch();
         Map<String, TopicMetadata> subscriptionMetadata = group.subscriptionMetadata();
 
@@ -1166,38 +1169,17 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
             }
         }
 
-        // 3. Reconcile the member's assignment with the target assignment. This is only required if
-        // the member is not stable or if a new target assignment has been installed.
-        boolean assignmentUpdated = false;
-        if (updatedMember.state() != ConsumerGroupMember.MemberState.STABLE || updatedMember.targetMemberEpoch() != targetAssignmentEpoch) {
-            ConsumerGroupMember prevMember = updatedMember;
-            updatedMember = new CurrentAssignmentBuilder(updatedMember)
-                .withTargetAssignment(targetAssignmentEpoch, targetAssignment)
-                .withCurrentPartitionEpoch(group::currentPartitionEpoch)
-                .withOwnedTopicPartitions(ownedTopicPartitions)
-                .build();
-
-            // Checking the reference is enough here because a new instance
-            // is created only when the state has changed.
-            if (updatedMember != prevMember) {
-                assignmentUpdated = true;
-                records.add(newCurrentAssignmentRecord(groupId, updatedMember));
-
-                log.info("[GroupId {}] Member {} transitioned from {} to {}.",
-                    groupId, memberId, member.currentAssignmentSummary(), updatedMember.currentAssignmentSummary());
-
-                if (updatedMember.state() == ConsumerGroupMember.MemberState.REVOKING) {
-                    scheduleConsumerGroupRevocationTimeout(
-                        groupId,
-                        memberId,
-                        updatedMember.rebalanceTimeoutMs(),
-                        updatedMember.memberEpoch()
-                    );
-                } else {
-                    cancelConsumerGroupRevocationTimeout(groupId, memberId);
-                }
-            }
-        }
+        // 3. Reconcile the member's assignment with the target assignment if the member is not
+        // fully reconciled yet.
+        updatedMember = maybeReconcile(
+            groupId,
+            updatedMember,
+            group::currentPartitionEpoch,
+            targetAssignmentEpoch,
+            targetAssignment,
+            ownedTopicPartitions,
+            records
+        );
 
         scheduleConsumerGroupSessionTimeout(groupId, memberId);
 
@@ -1211,13 +1193,71 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
         // 1. The member reported its owned partitions;
         // 2. The member just joined or rejoined to group (epoch equals to zero);
         // 3. The member's assignment has been updated.
-        if (ownedTopicPartitions != null || memberEpoch == 0 || assignmentUpdated) {
+        if (ownedTopicPartitions != null || memberEpoch == 0 || hasAssignedPartitionsChanged(member, updatedMember)) {
             response.setAssignment(createResponseAssignment(updatedMember));
         }
 
         return new CoordinatorResult<>(records, response);
     }
 
+    /**
+     * Reconciles the current assignment of the member towards the target assignment if needed.
+     *
+     * @param groupId               The group id.
+     * @param member                The member to reconcile.
+     * @param currentPartitionEpoch The function returning the current epoch of
+     *                              a given partition.
+     * @param targetAssignmentEpoch The target assignment epoch.
+     * @param targetAssignment      The target assignment.
+     * @param ownedTopicPartitions  The list of partitions owned by the member. This
+     *                              is reported in the ConsumerGroupHeartbeat API and
+     *                              it could be null if not provided.
+     * @param records               The list to accumulate any new records.
+     * @return The received member if no changes have been made; or a new
+     *         member containing the new assignment.
+     */
+    private ConsumerGroupMember maybeReconcile(
+        String groupId,
+        ConsumerGroupMember member,
+        BiFunction<Uuid, Integer, Integer> currentPartitionEpoch,
+        int targetAssignmentEpoch,
+        Assignment targetAssignment,
+        List<ConsumerGroupHeartbeatRequestData.TopicPartitions> ownedTopicPartitions,
+        List<Record> records
+    ) {
+        if (member.isReconciledTo(targetAssignmentEpoch)) {
+            return member;
+        }
+
+        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+            .withTargetAssignment(targetAssignmentEpoch, targetAssignment)
+            .withCurrentPartitionEpoch(currentPartitionEpoch)
+            .withOwnedTopicPartitions(ownedTopicPartitions)
+            .build();
+
+        if (!updatedMember.equals(member)) {
+            records.add(newCurrentAssignmentRecord(groupId, updatedMember));
+
+            log.info("[GroupId {}] Member {} new assignment state: epoch={}, previousEpoch={}, state={}, "
+                     + "assignedPartitions={} and revokedPartitions={}.",
+                groupId, updatedMember.memberId(), updatedMember.memberEpoch(), updatedMember.previousMemberEpoch(), updatedMember.state(),
+                assignmentToString(updatedMember.assignedPartitions()), assignmentToString(updatedMember.partitionsPendingRevocation()));
+
+            if (updatedMember.state() == MemberState.UNREVOKED_PARTITIONS) {
+                scheduleConsumerGroupRebalanceTimeout(
+                    groupId,
+                    updatedMember.memberId(),
+                    updatedMember.memberEpoch(),
+                    updatedMember.rebalanceTimeoutMs()
+                );
+            } else {
+                cancelConsumerGroupRebalanceTimeout(groupId, updatedMember.memberId());
+            }
+        }
+
+        return updatedMember;
+    }
+
     private void removeMemberAndCancelTimers(
         List<Record> records,
         String groupId,
@@ -1353,7 +1393,7 @@ private void removeMember(List<Record> records, String groupId, String memberId)
      */
     private void cancelTimers(String groupId, String memberId) {
         cancelConsumerGroupSessionTimeout(groupId, memberId);
-        cancelConsumerGroupRevocationTimeout(groupId, memberId);
+        cancelConsumerGroupRebalanceTimeout(groupId, memberId);
     }
 
     /**
@@ -1400,35 +1440,35 @@ private void cancelConsumerGroupSessionTimeout(
     }
 
     /**
-     * Schedules a revocation timeout for the member.
+     * Schedules a rebalance timeout for the member.
      *
      * @param groupId               The group id.
      * @param memberId              The member id.
-     * @param revocationTimeoutMs   The revocation timeout.
-     * @param expectedMemberEpoch   The expected member epoch.
+     * @param memberEpoch           The member epoch.
+     * @param rebalanceTimeoutMs    The rebalance timeout.
      */
-    private void scheduleConsumerGroupRevocationTimeout(
+    private void scheduleConsumerGroupRebalanceTimeout(
         String groupId,
         String memberId,
-        long revocationTimeoutMs,
-        int expectedMemberEpoch
+        int memberEpoch,
+        int rebalanceTimeoutMs
     ) {
-        String key = consumerGroupRevocationTimeoutKey(groupId, memberId);
-        timer.schedule(key, revocationTimeoutMs, TimeUnit.MILLISECONDS, true, () -> {
+        String key = consumerGroupRebalanceTimeoutKey(groupId, memberId);
+        timer.schedule(key, rebalanceTimeoutMs, TimeUnit.MILLISECONDS, true, () -> {
             try {
                 ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, false);
                 ConsumerGroupMember member = group.getOrMaybeCreateMember(memberId, false);
 
-                if (member.state() != ConsumerGroupMember.MemberState.REVOKING ||
-                    member.memberEpoch() != expectedMemberEpoch) {
-                    log.debug("[GroupId {}] Ignoring revocation timeout for {} because the member " +
-                        "state does not match the expected state.", groupId, memberId);
+                if (member.memberEpoch() == memberEpoch) {
+                    log.info("[GroupId {}] Member {} fenced from the group because " +
+                            "it failed to transition from epoch {} within {}ms.",
+                        groupId, memberId, memberEpoch, rebalanceTimeoutMs);
+                    return new CoordinatorResult<>(consumerGroupFenceMember(group, member));
+                } else {
+                    log.debug("[GroupId {}] Ignoring rebalance timeout for {} because the member " +
+                        "left the epoch {}.", groupId, memberId, memberEpoch);
                     return new CoordinatorResult<>(Collections.emptyList());
                 }
-
-                log.info("[GroupId {}] Member {} fenced from the group because " +
-                    "it failed to revoke partitions within {}ms.", groupId, memberId, revocationTimeoutMs);
-                return new CoordinatorResult<>(consumerGroupFenceMember(group, member));
             } catch (GroupIdNotFoundException ex) {
                 log.debug("[GroupId {}] Could not fence {}} because the group does not exist.",
                     groupId, memberId);
@@ -1442,16 +1482,16 @@ private void scheduleConsumerGroupRevocationTimeout(
     }
 
     /**
-     * Cancels the revocation timeout of the member.
+     * Cancels the rebalance timeout of the member.
      *
      * @param groupId       The group id.
      * @param memberId      The member id.
      */
-    private void cancelConsumerGroupRevocationTimeout(
+    private void cancelConsumerGroupRebalanceTimeout(
         String groupId,
         String memberId
     ) {
-        timer.cancel(consumerGroupRevocationTimeoutKey(groupId, memberId));
+        timer.cancel(consumerGroupRebalanceTimeoutKey(groupId, memberId));
     }
 
     /**
@@ -1744,10 +1784,8 @@ public void replay(
             ConsumerGroupMember newMember = new ConsumerGroupMember.Builder(oldMember)
                 .setMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                 .setPreviousMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
-                .setTargetMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                 .setAssignedPartitions(Collections.emptyMap())
                 .setPartitionsPendingRevocation(Collections.emptyMap())
-                .setPartitionsPendingAssignment(Collections.emptyMap())
                 .build();
             consumerGroup.updateMember(newMember);
         }
@@ -1796,12 +1834,12 @@ public void onLoaded() {
                     consumerGroup.members().forEach((memberId, member) -> {
                         log.debug("Loaded member {} in consumer group {}.", memberId, groupId);
                         scheduleConsumerGroupSessionTimeout(groupId, memberId);
-                        if (member.state() == ConsumerGroupMember.MemberState.REVOKING) {
-                            scheduleConsumerGroupRevocationTimeout(
+                        if (member.state() == MemberState.UNREVOKED_PARTITIONS) {
+                            scheduleConsumerGroupRebalanceTimeout(
                                 groupId,
-                                memberId,
-                                member.rebalanceTimeoutMs(),
-                                member.memberEpoch()
+                                member.memberId(),
+                                member.memberEpoch(),
+                                member.rebalanceTimeoutMs()
                             );
                         }
                     });
@@ -1834,8 +1872,8 @@ public static String consumerGroupSessionTimeoutKey(String groupId, String membe
         return "session-timeout-" + groupId + "-" + memberId;
     }
 
-    public static String consumerGroupRevocationTimeoutKey(String groupId, String memberId) {
-        return "revocation-timeout-" + groupId + "-" + memberId;
+    public static String consumerGroupRebalanceTimeoutKey(String groupId, String memberId) {
+        return "rebalance-timeout-" + groupId + "-" + memberId;
     }
 
     /**
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
index 47ba36c84a3b5..b162100be3745 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
@@ -346,10 +346,9 @@ public static Record newCurrentAssignmentRecord(
                 new ConsumerGroupCurrentMemberAssignmentValue()
                     .setMemberEpoch(member.memberEpoch())
                     .setPreviousMemberEpoch(member.previousMemberEpoch())
-                    .setTargetMemberEpoch(member.targetMemberEpoch())
+                    .setState(member.state().value())
                     .setAssignedPartitions(toTopicPartitions(member.assignedPartitions()))
-                    .setPartitionsPendingRevocation(toTopicPartitions(member.partitionsPendingRevocation()))
-                    .setPartitionsPendingAssignment(toTopicPartitions(member.partitionsPendingAssignment())),
+                    .setPartitionsPendingRevocation(toTopicPartitions(member.partitionsPendingRevocation())),
                 (short) 0
             )
         );
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/Utils.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/Utils.java
index 7fe92c5bd8cf4..cbbe45cc61fc7 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/Utils.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/Utils.java
@@ -16,8 +16,13 @@
  */
 package org.apache.kafka.coordinator.group;
 
+import org.apache.kafka.common.Uuid;
+
+import java.util.Iterator;
+import java.util.Map;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
+import java.util.Set;
 
 public class Utils {
     private Utils() {}
@@ -37,4 +42,31 @@ public static OptionalInt ofSentinel(int value) {
     public static OptionalLong ofSentinel(long value) {
         return value != -1 ? OptionalLong.of(value) : OptionalLong.empty();
     }
+
+    /**
+     * @return The provided assignment as a String.
+     *
+     * Example:
+     * [topicid1-0, topicid1-1, topicid2-0, topicid2-1]
+     */
+    public static String assignmentToString(
+        Map<Uuid, Set<Integer>> assignment
+    ) {
+        StringBuilder builder = new StringBuilder("[");
+        Iterator<Map.Entry<Uuid, Set<Integer>>> topicsIterator = assignment.entrySet().iterator();
+        while (topicsIterator.hasNext()) {
+            Map.Entry<Uuid, Set<Integer>> entry = topicsIterator.next();
+            Iterator<Integer> partitionsIterator = entry.getValue().iterator();
+            while (partitionsIterator.hasNext()) {
+                builder.append(entry.getKey());
+                builder.append("-");
+                builder.append(partitionsIterator.next());
+                if (partitionsIterator.hasNext() || topicsIterator.hasNext()) {
+                    builder.append(", ");
+                }
+            }
+        }
+        builder.append("]");
+        return builder.toString();
+    }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
index d4077ef99d208..68851b5176ce0 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
@@ -779,7 +779,7 @@ private void maybeUpdateGroupState() {
             newState = ASSIGNING;
         } else {
             for (ConsumerGroupMember member : members.values()) {
-                if (member.targetMemberEpoch() != targetAssignmentEpoch.get() || member.state() != ConsumerGroupMember.MemberState.STABLE) {
+                if (!member.isReconciledTo(targetAssignmentEpoch.get())) {
                     newState = RECONCILING;
                     break;
                 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
index 5159f1f27cac6..fdf919cc81cf7 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
@@ -41,6 +41,7 @@
  * by records stored in the __consumer_offsets topic.
  */
 public class ConsumerGroupMember {
+
     /**
      * A builder that facilitates the creation of a new member or the update of
      * an existing one.
@@ -52,7 +53,7 @@ public static class Builder {
         private final String memberId;
         private int memberEpoch = 0;
         private int previousMemberEpoch = -1;
-        private int targetMemberEpoch = 0;
+        private MemberState state = MemberState.STABLE;
         private String instanceId = null;
         private String rackId = null;
         private int rebalanceTimeoutMs = -1;
@@ -64,7 +65,6 @@ public static class Builder {
         private List<ClientAssignor> clientAssignors = Collections.emptyList();
         private Map<Uuid, Set<Integer>> assignedPartitions = Collections.emptyMap();
         private Map<Uuid, Set<Integer>> partitionsPendingRevocation = Collections.emptyMap();
-        private Map<Uuid, Set<Integer>> partitionsPendingAssignment = Collections.emptyMap();
 
         public Builder(String memberId) {
             this.memberId = Objects.requireNonNull(memberId);
@@ -76,7 +76,6 @@ public Builder(ConsumerGroupMember member) {
             this.memberId = member.memberId;
             this.memberEpoch = member.memberEpoch;
             this.previousMemberEpoch = member.previousMemberEpoch;
-            this.targetMemberEpoch = member.targetMemberEpoch;
             this.instanceId = member.instanceId;
             this.rackId = member.rackId;
             this.rebalanceTimeoutMs = member.rebalanceTimeoutMs;
@@ -86,23 +85,25 @@ public Builder(ConsumerGroupMember member) {
             this.subscribedTopicRegex = member.subscribedTopicRegex;
             this.serverAssignorName = member.serverAssignorName;
             this.clientAssignors = member.clientAssignors;
+            this.state = member.state;
             this.assignedPartitions = member.assignedPartitions;
             this.partitionsPendingRevocation = member.partitionsPendingRevocation;
-            this.partitionsPendingAssignment = member.partitionsPendingAssignment;
         }
 
-        public Builder setMemberEpoch(int memberEpoch) {
+        public Builder updateMemberEpoch(int memberEpoch) {
+            int currentMemberEpoch = this.memberEpoch;
             this.memberEpoch = memberEpoch;
+            this.previousMemberEpoch = currentMemberEpoch;
             return this;
         }
 
-        public Builder setPreviousMemberEpoch(int previousMemberEpoch) {
-            this.previousMemberEpoch = previousMemberEpoch;
+        public Builder setMemberEpoch(int memberEpoch) {
+            this.memberEpoch = memberEpoch;
             return this;
         }
 
-        public Builder setTargetMemberEpoch(int targetMemberEpoch) {
-            this.targetMemberEpoch = targetMemberEpoch;
+        public Builder setPreviousMemberEpoch(int previousMemberEpoch) {
+            this.previousMemberEpoch = previousMemberEpoch;
             return this;
         }
 
@@ -188,6 +189,11 @@ public Builder maybeUpdateClientAssignors(Optional<List<ClientAssignor>> clientA
             return this;
         }
 
+        public Builder setState(MemberState state) {
+            this.state = state;
+            return this;
+        }
+
         public Builder setAssignedPartitions(Map<Uuid, Set<Integer>> assignedPartitions) {
             this.assignedPartitions = assignedPartitions;
             return this;
@@ -198,11 +204,6 @@ public Builder setPartitionsPendingRevocation(Map<Uuid, Set<Integer>> partitions
             return this;
         }
 
-        public Builder setPartitionsPendingAssignment(Map<Uuid, Set<Integer>> partitionsPendingAssignment) {
-            this.partitionsPendingAssignment = partitionsPendingAssignment;
-            return this;
-        }
-
         public Builder updateWith(ConsumerGroupMemberMetadataValue record) {
             setInstanceId(record.instanceId());
             setRackId(record.rackId());
@@ -221,10 +222,9 @@ public Builder updateWith(ConsumerGroupMemberMetadataValue record) {
         public Builder updateWith(ConsumerGroupCurrentMemberAssignmentValue record) {
             setMemberEpoch(record.memberEpoch());
             setPreviousMemberEpoch(record.previousMemberEpoch());
-            setTargetMemberEpoch(record.targetMemberEpoch());
+            setState(MemberState.fromValue(record.state()));
             setAssignedPartitions(assignmentFromTopicPartitions(record.assignedPartitions()));
             setPartitionsPendingRevocation(assignmentFromTopicPartitions(record.partitionsPendingRevocation()));
-            setPartitionsPendingAssignment(assignmentFromTopicPartitions(record.partitionsPendingAssignment()));
             return this;
         }
 
@@ -237,20 +237,10 @@ private Map<Uuid, Set<Integer>> assignmentFromTopicPartitions(
         }
 
         public ConsumerGroupMember build() {
-            MemberState state;
-            if (!partitionsPendingRevocation.isEmpty()) {
-                state = MemberState.REVOKING;
-            } else if (!partitionsPendingAssignment.isEmpty()) {
-                state = MemberState.ASSIGNING;
-            } else {
-                state = MemberState.STABLE;
-            }
-
             return new ConsumerGroupMember(
                 memberId,
                 memberEpoch,
                 previousMemberEpoch,
-                targetMemberEpoch,
                 instanceId,
                 rackId,
                 rebalanceTimeoutMs,
@@ -262,33 +252,11 @@ public ConsumerGroupMember build() {
                 clientAssignors,
                 state,
                 assignedPartitions,
-                partitionsPendingRevocation,
-                partitionsPendingAssignment
+                partitionsPendingRevocation
             );
         }
     }
 
-    /**
-     * The various states that a member can be in. For their definition,
-     * refer to the documentation of {{@link CurrentAssignmentBuilder}}.
-     */
-    public enum MemberState {
-        REVOKING("revoking"),
-        ASSIGNING("assigning"),
-        STABLE("stable");
-
-        private final String name;
-
-        MemberState(String name) {
-            this.name = name;
-        }
-
-        @Override
-        public String toString() {
-            return name;
-        }
-    }
-
     /**
      * The member id.
      */
@@ -305,11 +273,9 @@ public String toString() {
     private final int previousMemberEpoch;
 
     /**
-     * The next member epoch. This corresponds to the target
-     * assignment epoch used to compute the current assigned,
-     * revoking and assigning partitions.
+     * The member state.
      */
-    private final int targetMemberEpoch;
+    private final MemberState state;
 
     /**
      * The instance id provided by the member.
@@ -356,11 +322,6 @@ public String toString() {
      */
     private final List<ClientAssignor> clientAssignors;
 
-    /**
-     * The member state.
-     */
-    private final MemberState state;
-
     /**
      * The partitions assigned to this member.
      */
@@ -371,18 +332,10 @@ public String toString() {
      */
     private final Map<Uuid, Set<Integer>> partitionsPendingRevocation;
 
-    /**
-     * The partitions waiting to be assigned to this
-     * member. They will be assigned when they are
-     * released by their previous owners.
-     */
-    private final Map<Uuid, Set<Integer>> partitionsPendingAssignment;
-
     private ConsumerGroupMember(
         String memberId,
         int memberEpoch,
         int previousMemberEpoch,
-        int targetMemberEpoch,
         String instanceId,
         String rackId,
         int rebalanceTimeoutMs,
@@ -394,13 +347,12 @@ private ConsumerGroupMember(
         List<ClientAssignor> clientAssignors,
         MemberState state,
         Map<Uuid, Set<Integer>> assignedPartitions,
-        Map<Uuid, Set<Integer>> partitionsPendingRevocation,
-        Map<Uuid, Set<Integer>> partitionsPendingAssignment
+        Map<Uuid, Set<Integer>> partitionsPendingRevocation
     ) {
         this.memberId = memberId;
         this.memberEpoch = memberEpoch;
         this.previousMemberEpoch = previousMemberEpoch;
-        this.targetMemberEpoch = targetMemberEpoch;
+        this.state = state;
         this.instanceId = instanceId;
         this.rackId = rackId;
         this.rebalanceTimeoutMs = rebalanceTimeoutMs;
@@ -410,10 +362,8 @@ private ConsumerGroupMember(
         this.subscribedTopicRegex = subscribedTopicRegex;
         this.serverAssignorName = serverAssignorName;
         this.clientAssignors = clientAssignors;
-        this.state = state;
         this.assignedPartitions = assignedPartitions;
         this.partitionsPendingRevocation = partitionsPendingRevocation;
-        this.partitionsPendingAssignment = partitionsPendingAssignment;
     }
 
     /**
@@ -437,13 +387,6 @@ public int previousMemberEpoch() {
         return previousMemberEpoch;
     }
 
-    /**
-     * @return The target member epoch.
-     */
-    public int targetMemberEpoch() {
-        return targetMemberEpoch;
-    }
-
     /**
      * @return The instance id.
      */
@@ -514,6 +457,13 @@ public MemberState state() {
         return state;
     }
 
+    /**
+     * @return True if the member is in the Stable state and at the desired epoch.
+     */
+    public boolean isReconciledTo(int targetAssignmentEpoch) {
+        return state == MemberState.STABLE && memberEpoch == targetAssignmentEpoch;
+    }
+
     /**
      * @return The set of assigned partitions.
      */
@@ -528,27 +478,6 @@ public Map<Uuid, Set<Integer>> partitionsPendingRevocation() {
         return partitionsPendingRevocation;
     }
 
-    /**
-     * @return The set of partitions awaiting assignment to the member.
-     */
-    public Map<Uuid, Set<Integer>> partitionsPendingAssignment() {
-        return partitionsPendingAssignment;
-    }
-
-    /**
-     * @return A string representation of the current assignment state.
-     */
-    public String currentAssignmentSummary() {
-        return "CurrentAssignment(memberEpoch=" + memberEpoch +
-            ", previousMemberEpoch=" + previousMemberEpoch +
-            ", targetMemberEpoch=" + targetMemberEpoch +
-            ", state=" + state +
-            ", assignedPartitions=" + assignedPartitions +
-            ", partitionsPendingRevocation=" + partitionsPendingRevocation +
-            ", partitionsPendingAssignment=" + partitionsPendingAssignment +
-            ')';
-    }
-
     /**
      * @param targetAssignment The target assignment of this member in the corresponding group.
      *
@@ -612,7 +541,7 @@ public boolean equals(Object o) {
         ConsumerGroupMember that = (ConsumerGroupMember) o;
         return memberEpoch == that.memberEpoch
             && previousMemberEpoch == that.previousMemberEpoch
-            && targetMemberEpoch == that.targetMemberEpoch
+            && state == that.state
             && rebalanceTimeoutMs == that.rebalanceTimeoutMs
             && Objects.equals(memberId, that.memberId)
             && Objects.equals(instanceId, that.instanceId)
@@ -624,8 +553,7 @@ public boolean equals(Object o) {
             && Objects.equals(serverAssignorName, that.serverAssignorName)
             && Objects.equals(clientAssignors, that.clientAssignors)
             && Objects.equals(assignedPartitions, that.assignedPartitions)
-            && Objects.equals(partitionsPendingRevocation, that.partitionsPendingRevocation)
-            && Objects.equals(partitionsPendingAssignment, that.partitionsPendingAssignment);
+            && Objects.equals(partitionsPendingRevocation, that.partitionsPendingRevocation);
     }
 
     @Override
@@ -633,7 +561,7 @@ public int hashCode() {
         int result = memberId != null ? memberId.hashCode() : 0;
         result = 31 * result + memberEpoch;
         result = 31 * result + previousMemberEpoch;
-        result = 31 * result + targetMemberEpoch;
+        result = 31 * result + Objects.hashCode(state);
         result = 31 * result + Objects.hashCode(instanceId);
         result = 31 * result + Objects.hashCode(rackId);
         result = 31 * result + rebalanceTimeoutMs;
@@ -645,7 +573,6 @@ public int hashCode() {
         result = 31 * result + Objects.hashCode(clientAssignors);
         result = 31 * result + Objects.hashCode(assignedPartitions);
         result = 31 * result + Objects.hashCode(partitionsPendingRevocation);
-        result = 31 * result + Objects.hashCode(partitionsPendingAssignment);
         return result;
     }
 
@@ -655,7 +582,7 @@ public String toString() {
             "memberId='" + memberId + '\'' +
             ", memberEpoch=" + memberEpoch +
             ", previousMemberEpoch=" + previousMemberEpoch +
-            ", targetMemberEpoch=" + targetMemberEpoch +
+            ", state='" + state + '\'' +
             ", instanceId='" + instanceId + '\'' +
             ", rackId='" + rackId + '\'' +
             ", rebalanceTimeoutMs=" + rebalanceTimeoutMs +
@@ -665,10 +592,18 @@ public String toString() {
             ", subscribedTopicRegex='" + subscribedTopicRegex + '\'' +
             ", serverAssignorName='" + serverAssignorName + '\'' +
             ", clientAssignors=" + clientAssignors +
-            ", state=" + state +
             ", assignedPartitions=" + assignedPartitions +
             ", partitionsPendingRevocation=" + partitionsPendingRevocation +
-            ", partitionsPendingAssignment=" + partitionsPendingAssignment +
             ')';
     }
+
+    /**
+     * @return True of the two provided members have different assigned partitions.
+     */
+    public static boolean hasAssignedPartitionsChanged(
+        ConsumerGroupMember member1,
+        ConsumerGroupMember member2
+    ) {
+        return !member1.assignedPartitions().equals(member2.assignedPartitions());
+    }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilder.java
index fce5b8a85bd41..f4306d95fcf5f 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilder.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilder.java
@@ -17,12 +17,12 @@
 package org.apache.kafka.coordinator.group.consumer;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
 
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -33,49 +33,6 @@
  * The CurrentAssignmentBuilder class encapsulates the reconciliation engine of the
  * consumer group protocol. Given the current state of a member and a desired or target
  * assignment state, the state machine takes the necessary steps to converge them.
- *
- * The member state has the following properties:
- * - Current Epoch:
- *   The current epoch of the member.
- *
- * - Next Epoch:
- *   The desired epoch of the member. It corresponds to the epoch of the target/desired assignment.
- *   The member transitions to this epoch when it has revoked the partitions that it does not own
- *   or if it does not have to revoke any.
- *
- * - Previous Epoch:
- *   The epoch of the member when the state was last updated.
- *
- * - Assigned Partitions:
- *   The set of partitions currently assigned to the member. This represents what the member should have.
- *
- * - Partitions Pending Revocation:
- *   The set of partitions that the member should revoke before it can transition to the next state.
- *
- * - Partitions Pending Assignment:
- *   The set of partitions that the member will eventually receive. The partitions in this set are
- *   still owned by other members in the group.
- *
- * The state machine has three states:
- * - REVOKING:
- *   This state means that the member must revoke partitions before it can transition to the next epoch
- *   and thus start receiving new partitions. This is to guarantee that offsets of revoked partitions
- *   are committed with the current epoch. The member transitions to the next state only when it has
- *   acknowledged the revocation.
- *
- * - ASSIGNING:
- *   This state means that the member waits on partitions which are still owned by other members in the
- *   group. It remains in this state until they are all freed up.
- *
- * - STABLE:
- *   This state means that the member has received all its assigned partitions.
- *
- * The reconciliation process is started or re-started whenever a new target assignment is installed;
- * the epoch of the new target assignment is different from the next epoch of the member. In this transient
- * state, the assigned partitions, the partitions pending revocation and the partitions pending assignment
- * are updated. If the partitions pending revocation is not empty, the state machine transitions to
- * REVOKING; if partitions pending assignment is not empty, it transitions to ASSIGNING; otherwise it
- * transitions to STABLE.
  */
 public class CurrentAssignmentBuilder {
     /**
@@ -170,72 +127,122 @@ public CurrentAssignmentBuilder withOwnedTopicPartitions(
      * @return A new ConsumerGroupMember or the current one.
      */
     public ConsumerGroupMember build() {
-        // A new target assignment has been installed, we need to restart
-        // the reconciliation loop from the beginning.
-        if (targetAssignmentEpoch != member.targetMemberEpoch()) {
-            return transitionToNewTargetAssignmentState();
-        }
-
         switch (member.state()) {
-            // Check if the partitions have been revoked by the member.
-            case REVOKING:
-                return maybeTransitionFromRevokingToAssigningOrStable();
+            case STABLE:
+                // When the member is in the STABLE state, we verify if a newer
+                // epoch (or target assignment) is available. If it is, we can
+                // reconcile the member towards it. Otherwise, we return.
+                if (member.memberEpoch() != targetAssignmentEpoch) {
+                    return computeNextAssignment(
+                        member.memberEpoch(),
+                        member.assignedPartitions()
+                    );
+                } else {
+                    return member;
+                }
 
-            // Check if pending partitions have been freed up.
-            case ASSIGNING:
-                return maybeTransitionFromAssigningToAssigningOrStable();
+            case UNREVOKED_PARTITIONS:
+                // When the member is in the UNREVOKED_PARTITIONS state, we wait
+                // until the member has revoked the necessary partitions. They are
+                // considered revoked when they are not anymore reported in the
+                // owned partitions set in the ConsumerGroupHeartbeat API.
 
-            // Nothing to do.
-            case STABLE:
-                return member;
+                // If the member does not provide its owned partitions. We cannot
+                // progress.
+                if (ownedTopicPartitions == null) {
+                    return member;
+                }
+
+                // If the member provides its owned partitions. We verify if it still
+                // owns any of the revoked partitions. If it does, we cannot progress.
+                for (ConsumerGroupHeartbeatRequestData.TopicPartitions topicPartitions : ownedTopicPartitions) {
+                    for (Integer partitionId : topicPartitions.partitions()) {
+                        boolean stillHasRevokedPartition = member
+                            .partitionsPendingRevocation()
+                            .getOrDefault(topicPartitions.topicId(), Collections.emptySet())
+                            .contains(partitionId);
+                        if (stillHasRevokedPartition) {
+                            return member;
+                        }
+                    }
+                }
+
+                // When the member has revoked all the pending partitions, it can
+                // transition to the next epoch (current + 1) and we can reconcile
+                // its state towards the latest target assignment.
+                return computeNextAssignment(
+                    member.memberEpoch() + 1,
+                    member.assignedPartitions()
+                );
+
+            case UNRELEASED_PARTITIONS:
+                // When the member is in the UNRELEASED_PARTITIONS, we reconcile the
+                // member towards the latest target assignment. This will assign any
+                // of the unreleased partitions when they become available.
+                return computeNextAssignment(
+                    member.memberEpoch(),
+                    member.assignedPartitions()
+                );
+
+            case UNKNOWN:
+                // We could only end up in this state if a new state is added in the
+                // future and the group coordinator is downgraded. In this case, the
+                // best option is to fence the member to force it to rejoin the group
+                // without any partitions and to reconcile it again from scratch.
+                if (ownedTopicPartitions == null || !ownedTopicPartitions.isEmpty()) {
+                    throw new FencedMemberEpochException("The consumer group member is in a unknown state. "
+                        + "The member must abandon all its partitions and rejoin.");
+                }
+
+                return computeNextAssignment(
+                    targetAssignmentEpoch,
+                    member.assignedPartitions()
+                );
         }
 
         return member;
     }
 
     /**
-     * Transitions to NewTargetAssignment state. This is a transient state where
-     * we compute the assigned partitions, the partitions pending revocation,
-     * the partitions pending assignment, and transition to the next state.
+     * Computes the next assignment.
      *
+     * @param memberEpoch               The epoch of the member to use. This may be different
+     *                                  from the epoch in {@link CurrentAssignmentBuilder#member}.
+     * @param memberAssignedPartitions  The assigned partitions of the member to use.
      * @return A new ConsumerGroupMember.
      */
-    private ConsumerGroupMember transitionToNewTargetAssignmentState() {
+    private ConsumerGroupMember computeNextAssignment(
+        int memberEpoch,
+        Map<Uuid, Set<Integer>> memberAssignedPartitions
+    ) {
+        boolean hasUnreleasedPartitions = false;
         Map<Uuid, Set<Integer>> newAssignedPartitions = new HashMap<>();
         Map<Uuid, Set<Integer>> newPartitionsPendingRevocation = new HashMap<>();
         Map<Uuid, Set<Integer>> newPartitionsPendingAssignment = new HashMap<>();
 
-        // Compute the combined set of topics.
         Set<Uuid> allTopicIds = new HashSet<>(targetAssignment.partitions().keySet());
-        allTopicIds.addAll(member.assignedPartitions().keySet());
-        allTopicIds.addAll(member.partitionsPendingRevocation().keySet());
-        allTopicIds.addAll(member.partitionsPendingAssignment().keySet());
+        allTopicIds.addAll(memberAssignedPartitions.keySet());
 
         for (Uuid topicId : allTopicIds) {
             Set<Integer> target = targetAssignment.partitions()
                 .getOrDefault(topicId, Collections.emptySet());
-            Set<Integer> currentAssignedPartitions = member.assignedPartitions()
-                .getOrDefault(topicId, Collections.emptySet());
-            Set<Integer> currentRevokingPartitions = member.partitionsPendingRevocation()
+            Set<Integer> currentAssignedPartitions = memberAssignedPartitions
                 .getOrDefault(topicId, Collections.emptySet());
 
-            // Assigned_1 = (Assigned_0 + Pending_Revocation_0) ∩ Target
-            // Assigned_0 + Pending_Revocation_0 is used here because the partitions
-            // being revoked are still owned until the revocation is acknowledged.
+            // New Assigned Partitions = Previous Assigned Partitions ∩ Target
             Set<Integer> assignedPartitions = new HashSet<>(currentAssignedPartitions);
-            assignedPartitions.addAll(currentRevokingPartitions);
             assignedPartitions.retainAll(target);
 
-            // Pending_Revocation_1 = (Assigned_0 + Pending_Revocation_0) - Assigned_1
-            // Assigned_0 + Pending_Revocation_0 is used here because the partitions
-            // being revoked are still owned until the revocation is acknowledged.
+            // Partitions Pending Revocation = Previous Assigned Partitions - New Assigned Partitions
             Set<Integer> partitionsPendingRevocation = new HashSet<>(currentAssignedPartitions);
-            partitionsPendingRevocation.addAll(currentRevokingPartitions);
             partitionsPendingRevocation.removeAll(assignedPartitions);
 
-            // Pending_Assignment_1 = Target - Assigned_1
+            // Partitions Pending Assignment = Target - New Assigned Partitions - Unreleased Partitions
             Set<Integer> partitionsPendingAssignment = new HashSet<>(target);
             partitionsPendingAssignment.removeAll(assignedPartitions);
+            hasUnreleasedPartitions = partitionsPendingAssignment.removeIf(partitionId ->
+                currentPartitionEpoch.apply(topicId, partitionId) != -1
+            ) || hasUnreleasedPartitions;
 
             if (!assignedPartitions.isEmpty()) {
                 newAssignedPartitions.put(topicId, assignedPartitions);
@@ -251,195 +258,51 @@ private ConsumerGroupMember transitionToNewTargetAssignmentState() {
         }
 
         if (!newPartitionsPendingRevocation.isEmpty()) {
-            // If the partition pending revocation set is not empty, we transition the
-            // member to revoking and keep the current epoch. The transition to the new
-            // state is done when the member is updated.
+            // If there are partitions to be revoked, the member remains in its current
+            // epoch and requests the revocation of those partitions. It transitions to
+            // the UNREVOKED_PARTITIONS state to wait until the client acknowledges the
+            // revocation of the partitions.
             return new ConsumerGroupMember.Builder(member)
+                .setState(MemberState.UNREVOKED_PARTITIONS)
+                .updateMemberEpoch(memberEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
                 .setPartitionsPendingRevocation(newPartitionsPendingRevocation)
-                .setPartitionsPendingAssignment(newPartitionsPendingAssignment)
-                .setTargetMemberEpoch(targetAssignmentEpoch)
                 .build();
-        } else {
-            if (!newPartitionsPendingAssignment.isEmpty()) {
-                // If the partitions pending assignment set is not empty, we check
-                // if some or all partitions are free to use. If they are, we move
-                // them to the partitions assigned set.
-                maybeAssignPendingPartitions(newAssignedPartitions, newPartitionsPendingAssignment);
-            }
-
-            // We transition to the target epoch. If the partitions pending assignment
-            // set is empty, the member transition to stable, otherwise to assigning.
-            // The transition to the new state is done when the member is updated.
+        } else if (!newPartitionsPendingAssignment.isEmpty()) {
+            // If there are partitions to be assigned, the member transitions to the
+            // target epoch and requests the assignment of those partitions. Note that
+            // the partitions are directly added to the assigned partitions set. The
+            // member transitions to the STABLE state or to the UNRELEASED_PARTITIONS
+            // state depending on whether there are unreleased partitions or not.
+            newPartitionsPendingAssignment.forEach((topicId, partitions) -> newAssignedPartitions
+                .computeIfAbsent(topicId, __ -> new HashSet<>())
+                .addAll(partitions));
+            MemberState newState = hasUnreleasedPartitions ? MemberState.UNRELEASED_PARTITIONS : MemberState.STABLE;
             return new ConsumerGroupMember.Builder(member)
+                .setState(newState)
+                .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
                 .setPartitionsPendingRevocation(Collections.emptyMap())
-                .setPartitionsPendingAssignment(newPartitionsPendingAssignment)
-                .setPreviousMemberEpoch(member.memberEpoch())
-                .setMemberEpoch(targetAssignmentEpoch)
-                .setTargetMemberEpoch(targetAssignmentEpoch)
                 .build();
-        }
-    }
-
-    /**
-     * Tries to transition from Revoke to Assigning or Stable. This is only
-     * possible when the member acknowledges that it only owns the partition
-     * in the assigned partitions.
-     *
-     * @return A new ConsumerGroupMember with the new state or the current one
-     *         if the member stays in the current state.
-     */
-    private ConsumerGroupMember maybeTransitionFromRevokingToAssigningOrStable() {
-        if (member.partitionsPendingRevocation().isEmpty() || matchesAssignedPartitions(ownedTopicPartitions)) {
-            Map<Uuid, Set<Integer>> newAssignedPartitions = deepCopy(member.assignedPartitions());
-            Map<Uuid, Set<Integer>> newPartitionsPendingAssignment = deepCopy(member.partitionsPendingAssignment());
-
-            if (!newPartitionsPendingAssignment.isEmpty()) {
-                // If the partitions pending assignment set is not empty, we check
-                // if some or all partitions are free to use. If they are, we move
-                // them to the assigned set.
-                maybeAssignPendingPartitions(newAssignedPartitions, newPartitionsPendingAssignment);
-            }
-
-            // We transition to the target epoch. If the partitions pending assignment
-            // set is empty, the member transition to stable, otherwise to assigning.
-            // The transition to the new state is done when the member is updated.
+        } else if (hasUnreleasedPartitions) {
+            // If there are no partitions to be revoked nor to be assigned but some
+            // partitions are not available yet, the member transitions to the target
+            // epoch, to the UNRELEASED_PARTITIONS state and waits.
             return new ConsumerGroupMember.Builder(member)
+                .setState(MemberState.UNRELEASED_PARTITIONS)
+                .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
                 .setPartitionsPendingRevocation(Collections.emptyMap())
-                .setPartitionsPendingAssignment(newPartitionsPendingAssignment)
-                .setPreviousMemberEpoch(member.memberEpoch())
-                .setMemberEpoch(targetAssignmentEpoch)
-                .setTargetMemberEpoch(targetAssignmentEpoch)
                 .build();
         } else {
-            return member;
-        }
-    }
-
-    /**
-     * Tries to transition from Assigning to Assigning or Stable. This is only
-     * possible when one or more partitions in the partitions pending assignment
-     * set have been freed up by other members in the group.
-     *
-     * @return A new ConsumerGroupMember with the new state or the current one
-     *         if the member stays in the current state.
-     */
-    private ConsumerGroupMember maybeTransitionFromAssigningToAssigningOrStable() {
-        Map<Uuid, Set<Integer>> newAssignedPartitions = deepCopy(member.assignedPartitions());
-        Map<Uuid, Set<Integer>> newPartitionsPendingAssignment = deepCopy(member.partitionsPendingAssignment());
-
-        // If any partition can transition from assigning to assigned, we update
-        // the member. Otherwise, we return the current one. The transition to the
-        // new state is done when the member is updated.
-        if (maybeAssignPendingPartitions(newAssignedPartitions, newPartitionsPendingAssignment)) {
+            // Otherwise, the member transitions to the target epoch and to the
+            // STABLE state.
             return new ConsumerGroupMember.Builder(member)
+                .setState(MemberState.STABLE)
+                .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
                 .setPartitionsPendingRevocation(Collections.emptyMap())
-                .setPartitionsPendingAssignment(newPartitionsPendingAssignment)
-                .setPreviousMemberEpoch(member.memberEpoch())
-                .setMemberEpoch(targetAssignmentEpoch)
-                .setTargetMemberEpoch(targetAssignmentEpoch)
                 .build();
-        } else {
-            return member;
-        }
-    }
-
-    /**
-     * Tries to move partitions from the partitions pending assignment set to
-     * the partitions assigned set if they are no longer owned.
-     *
-     * @param newAssignedPartitions             The assigned partitions.
-     * @param newPartitionsPendingAssignment    The partitions pending assignment.
-     * @return A boolean indicating if any partitions were moved.
-     */
-    private boolean maybeAssignPendingPartitions(
-        Map<Uuid, Set<Integer>> newAssignedPartitions,
-        Map<Uuid, Set<Integer>> newPartitionsPendingAssignment
-    ) {
-        boolean changed = false;
-
-        Iterator<Map.Entry<Uuid, Set<Integer>>> assigningSetIterator =
-            newPartitionsPendingAssignment.entrySet().iterator();
-
-        while (assigningSetIterator.hasNext()) {
-            Map.Entry<Uuid, Set<Integer>> pair = assigningSetIterator.next();
-            Uuid topicId = pair.getKey();
-            Set<Integer> assigning = pair.getValue();
-
-            Iterator<Integer> assigningIterator = assigning.iterator();
-            while (assigningIterator.hasNext()) {
-                Integer partitionId = assigningIterator.next();
-
-                // A partition can be assigned to this member iff it has been
-                // released by its previous owner. This is signaled by -1.
-                Integer partitionEpoch = currentPartitionEpoch.apply(topicId, partitionId);
-                if (partitionEpoch == -1) {
-                    assigningIterator.remove();
-                    put(newAssignedPartitions, topicId, partitionId);
-                    changed = true;
-                }
-            }
-
-            if (assigning.isEmpty()) {
-                assigningSetIterator.remove();
-            }
-        }
-
-        return changed;
-    }
-
-    /**
-     * Checks whether the owned topic partitions passed by the member to the state
-     * machine via the ConsumerGroupHeartbeat request corresponds to the assigned
-     * partitions.
-     *
-     * @param ownedTopicPartitions The topic partitions owned by the remove client.
-     * @return A boolean indicating if the owned partitions matches the Assigned set.
-     */
-    private boolean matchesAssignedPartitions(
-        List<ConsumerGroupHeartbeatRequestData.TopicPartitions> ownedTopicPartitions
-    ) {
-        if (ownedTopicPartitions == null) return false;
-        if (ownedTopicPartitions.size() != member.assignedPartitions().size()) return false;
-
-        for (ConsumerGroupHeartbeatRequestData.TopicPartitions topicPartitions : ownedTopicPartitions) {
-            Set<Integer> partitions = member.assignedPartitions().get(topicPartitions.topicId());
-            if (partitions == null) return false;
-            for (Integer partitionId : topicPartitions.partitions()) {
-                if (!partitions.contains(partitionId)) return false;
-            }
         }
-
-        return true;
-    }
-
-    /**
-     * Makes a deep copy of an assignment map.
-     *
-     * @param map The Map to copy.
-     * @return The copy.
-     */
-    private Map<Uuid, Set<Integer>> deepCopy(Map<Uuid, Set<Integer>> map) {
-        Map<Uuid, Set<Integer>> copy = new HashMap<>();
-        map.forEach((topicId, partitions) -> copy.put(topicId, new HashSet<>(partitions)));
-        return copy;
-    }
-
-    /**
-     * Puts the given TopicId and Partitions to the given map.
-     */
-    private void put(
-        Map<Uuid, Set<Integer>> map,
-        Uuid topicId,
-        Integer partitionId
-    ) {
-        map.compute(topicId, (__, partitionsOrNull) -> {
-            if (partitionsOrNull == null) partitionsOrNull = new HashSet<>();
-            partitionsOrNull.add(partitionId);
-            return partitionsOrNull;
-        });
     }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java
new file mode 100644
index 0000000000000..3f3237bfe0b12
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java
@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.consumer;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * The various states that a member can be in. For their definition,
+ * refer to the documentation of {{@link CurrentAssignmentBuilder}}.
+ */
+public enum MemberState {
+    /**
+     * The member is fully reconciled with the desired target assignment.
+     */
+    STABLE((byte) 0),
+
+    /**
+     * The member must revoke some partitions in order to be able to
+     * transition to the next epoch.
+     */
+    UNREVOKED_PARTITIONS((byte) 1),
+
+    /**
+     * The member transitioned to the last epoch but waits on some
+     * partitions which have not been revoked by their previous
+     * owners yet.
+     */
+    UNRELEASED_PARTITIONS((byte) 2),
+
+    /**
+     * The member is in an unknown state. This can only happen if a future
+     * version of the software introduces a new state unknown by this version.
+     */
+    UNKNOWN((byte) 127);
+
+    private final static Map<Byte, MemberState> VALUES_TO_ENUMS = new HashMap<>();
+
+    static {
+        for (MemberState state: MemberState.values()) {
+            VALUES_TO_ENUMS.put(state.value(), state);
+        }
+    }
+
+    private final byte value;
+
+    MemberState(byte value) {
+        this.value = value;
+    }
+
+    public byte value() {
+        return value;
+    }
+
+    public static MemberState fromValue(byte value) {
+        MemberState state = VALUES_TO_ENUMS.get(value);
+        if (state == null) {
+            return UNKNOWN;
+        }
+        return state;
+    }
+}
diff --git a/group-coordinator/src/main/resources/common/message/ConsumerGroupCurrentMemberAssignmentValue.json b/group-coordinator/src/main/resources/common/message/ConsumerGroupCurrentMemberAssignmentValue.json
index 575f8ac133b1d..b91665e69da7e 100644
--- a/group-coordinator/src/main/resources/common/message/ConsumerGroupCurrentMemberAssignmentValue.json
+++ b/group-coordinator/src/main/resources/common/message/ConsumerGroupCurrentMemberAssignmentValue.json
@@ -24,20 +24,12 @@
       "about": "The current member epoch that is expected from the member in the heartbeat request." },
     { "name": "PreviousMemberEpoch", "versions": "0+", "type": "int32",
       "about": "If the last epoch bump is lost before reaching the member, the member will retry with the previous epoch." },
-    { "name": "TargetMemberEpoch", "versions": "0+", "type": "int32",
-      "about": "The target epoch corresponding to the assignment used to compute the AssignedPartitions, the PartitionsPendingRevocation and the PartitionsPendingAssignment fields." },
+    { "name": "State", "versions": "0+", "type": "int8",
+      "about": "The member state. See ConsumerGroupMember.MemberState for the possible values." },
     { "name": "AssignedPartitions", "versions": "0+", "type": "[]TopicPartitions",
       "about": "The partitions assigned to (or owned by) this member." },
     { "name": "PartitionsPendingRevocation", "versions": "0+", "type": "[]TopicPartitions",
-      "about": "The partitions that must be revoked by this member." },
-    { "name": "PartitionsPendingAssignment", "versions": "0+", "type": "[]TopicPartitions",
-      "about": "The partitions that will be assigned to this member when they are freed up by their current owners." },
-    { "name": "Error", "versions": "0+", "type": "int8",
-      "about": "The error reported by the assignor." },
-    { "name": "MetadataVersion", "versions": "0+", "type": "int16",
-      "about": "The version of the metadata bytes." },
-    { "name": "MetadataBytes", "versions": "0+", "type": "bytes",
-      "about": "The metadata bytes." }
+      "about": "The partitions that must be revoked by this member." }
   ],
   "commonStructs": [
     { "name": "TopicPartitions", "versions": "0+", "fields": [
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
index 77d4c7d03c8d8..b5f40240ebb7f 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
@@ -144,18 +144,12 @@ private static void assertApiMessageAndVersionEquals(
 
             assertEquals(expectedValue.memberEpoch(), actualValue.memberEpoch());
             assertEquals(expectedValue.previousMemberEpoch(), actualValue.previousMemberEpoch());
-            assertEquals(expectedValue.targetMemberEpoch(), actualValue.targetMemberEpoch());
-            assertEquals(expectedValue.error(), actualValue.error());
-            assertEquals(expectedValue.metadataVersion(), actualValue.metadataVersion());
-            assertEquals(expectedValue.metadataBytes(), actualValue.metadataBytes());
 
             // We transform those to Maps before comparing them.
             assertEquals(fromTopicPartitions(expectedValue.assignedPartitions()),
                 fromTopicPartitions(actualValue.assignedPartitions()));
             assertEquals(fromTopicPartitions(expectedValue.partitionsPendingRevocation()),
                 fromTopicPartitions(actualValue.partitionsPendingRevocation()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingAssignment()),
-                fromTopicPartitions(actualValue.partitionsPendingAssignment()));
         } else if (actual.message() instanceof ConsumerGroupPartitionMetadataValue) {
             // The order of the racks stored in the PartitionMetadata of the ConsumerGroupPartitionMetadataValue
             // is not always guaranteed. Therefore, we need a special comparator.
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index 4a47fd3dc32a5..66592417828e9 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -65,6 +65,7 @@
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
+import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
 import org.apache.kafka.coordinator.group.generated.GroupMetadataValue;
 import org.apache.kafka.coordinator.group.classic.ClassicGroup;
@@ -106,7 +107,7 @@
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.appendGroupMetadataErrorToResponseError;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRevocationTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRebalanceTimeoutKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupSessionTimeoutKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.EMPTY_RESULT;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
@@ -134,7 +135,6 @@
 import static org.mockito.Mockito.when;
 
 public class GroupMetadataManagerTest {
-
     @Test
     public void testConsumerHeartbeatRequestValidation() {
         MockPartitionAssignor assignor = new MockPartitionAssignor("range");
@@ -336,9 +336,9 @@ public void testConsumerGroupMemberEpochValidation() {
             .build();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(100)
             .setPreviousMemberEpoch(99)
-            .setTargetMemberEpoch(100)
             .setRebalanceTimeoutMs(5000)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
@@ -464,9 +464,9 @@ public void testMemberJoinsEmptyConsumerGroup() {
         );
 
         ConsumerGroupMember expectedMember = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(1)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(1)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setRebalanceTimeoutMs(5000)
@@ -516,9 +516,9 @@ public void testUpdatingSubscriptionTriggersNewTargetAssignment() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Collections.singletonList("foo"))
@@ -563,9 +563,9 @@ public void testUpdatingSubscriptionTriggersNewTargetAssignment() {
         );
 
         ConsumerGroupMember expectedMember = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -618,9 +618,9 @@ public void testNewJoiningMemberTriggersNewTargetAssignment() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -631,9 +631,9 @@ public void testNewJoiningMemberTriggersNewTargetAssignment() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -692,17 +692,14 @@ public void testNewJoiningMemberTriggersNewTargetAssignment() {
         );
 
         ConsumerGroupMember expectedMember3 = new ConsumerGroupMember.Builder(memberId3)
+            .setState(MemberState.UNRELEASED_PARTITIONS)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(11)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setRebalanceTimeoutMs(5000)
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setServerAssignorName("range")
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(fooTopicId, 4, 5),
-                mkTopicAssignment(barTopicId, 2)))
             .build();
 
         List<Record> expectedRecords = Arrays.asList(
@@ -762,9 +759,9 @@ public void testLeavingMemberBumpsGroupEpoch() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -774,9 +771,9 @@ public void testLeavingMemberBumpsGroupEpoch() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     // Use zar only here to ensure that metadata needs to be recomputed.
@@ -854,10 +851,10 @@ public void testGroupEpochBumpWhenNewStaticMemberJoins() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -867,10 +864,10 @@ public void testGroupEpochBumpWhenNewStaticMemberJoins() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId2)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     // Use zar only here to ensure that metadata needs to be recomputed.
@@ -931,17 +928,14 @@ public void testGroupEpochBumpWhenNewStaticMemberJoins() {
 
         ConsumerGroupMember expectedMember3 = new ConsumerGroupMember.Builder(memberId3)
             .setMemberEpoch(11)
+            .setState(MemberState.UNRELEASED_PARTITIONS)
             .setInstanceId(memberId3)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(11)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setRebalanceTimeoutMs(5000)
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setServerAssignorName("range")
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(fooTopicId, 4, 5),
-                mkTopicAssignment(barTopicId, 2)))
             .build();
 
         List<Record> expectedRecords = Arrays.asList(
@@ -989,10 +983,10 @@ public void testStaticMemberGetsBackAssignmentUponRejoin() {
 
         MockPartitionAssignor assignor = new MockPartitionAssignor("range");
         ConsumerGroupMember member1 = new ConsumerGroupMember.Builder(memberId1)
+            .setState(MemberState.STABLE)
             .setInstanceId(memberId1)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(10)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1002,10 +996,10 @@ public void testStaticMemberGetsBackAssignmentUponRejoin() {
                 mkTopicAssignment(barTopicId, 0, 1)))
             .build();
         ConsumerGroupMember member2 = new ConsumerGroupMember.Builder(memberId2)
+            .setState(MemberState.STABLE)
             .setInstanceId(memberId2)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(10)
             .setRebalanceTimeoutMs(5000)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
@@ -1114,10 +1108,10 @@ public void testStaticMemberGetsBackAssignmentUponRejoin() {
         );
 
         ConsumerGroupMember expectedRejoinedMember = new ConsumerGroupMember.Builder(member2RejoinId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(10)
             .setInstanceId(memberId2)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(10)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setRebalanceTimeoutMs(5000)
@@ -1143,7 +1137,7 @@ public void testStaticMemberGetsBackAssignmentUponRejoin() {
         assertRecordsEquals(expectedRecordsAfterRejoin, rejoinResult.records());
         // Verify that there are no timers.
         context.assertNoSessionTimeout(groupId, memberId2);
-        context.assertNoRevocationTimeout(groupId, memberId2);
+        context.assertNoRebalanceTimeout(groupId, memberId2);
     }
 
     @Test
@@ -1160,10 +1154,10 @@ public void testNoGroupEpochBumpWhenStaticMemberTemporarilyLeaves() {
 
         MockPartitionAssignor assignor = new MockPartitionAssignor("range");
         ConsumerGroupMember member1 = new ConsumerGroupMember.Builder(memberId1)
+            .setState(MemberState.STABLE)
             .setInstanceId(memberId1)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(10)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1173,10 +1167,10 @@ public void testNoGroupEpochBumpWhenStaticMemberTemporarilyLeaves() {
                 mkTopicAssignment(barTopicId, 0, 1)))
             .build();
         ConsumerGroupMember member2 = new ConsumerGroupMember.Builder(memberId2)
+            .setState(MemberState.STABLE)
             .setInstanceId(memberId2)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(10)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             // Use zar only here to ensure that metadata needs to be recomputed.
@@ -1262,10 +1256,10 @@ public void testLeavingStaticMemberBumpsGroupEpoch() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1275,10 +1269,10 @@ public void testLeavingStaticMemberBumpsGroupEpoch() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId2)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     // Use zar only here to ensure that metadata needs to be recomputed.
@@ -1353,10 +1347,10 @@ public void testShouldThrownUnreleasedInstanceIdExceptionWhenNewMemberJoinsWithI
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1402,10 +1396,10 @@ public void testShouldThrownUnknownMemberIdExceptionWhenUnknownStaticMemberJoins
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1450,10 +1444,10 @@ public void testShouldThrowFencedInstanceIdExceptionWhenStaticMemberWithDifferen
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1490,10 +1484,10 @@ public void testConsumerGroupMemberEpochValidationForStaticMember() {
             .build();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setInstanceId(memberId)
             .setMemberEpoch(100)
             .setPreviousMemberEpoch(99)
-            .setTargetMemberEpoch(100)
             .setRebalanceTimeoutMs(5000)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
@@ -1583,10 +1577,10 @@ public void testShouldThrowUnknownMemberIdExceptionWhenUnknownStaticMemberLeaves
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1629,10 +1623,10 @@ public void testShouldThrowFencedInstanceIdExceptionWhenStaticMemberWithDifferen
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setInstanceId(memberId1)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -1680,9 +1674,9 @@ public void testReconciliationProcess() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -1693,9 +1687,9 @@ public void testReconciliationProcess() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -1737,13 +1731,14 @@ public void testReconciliationProcess() {
         CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result;
 
         // Members in the group are in Stable state.
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId1));
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId2));
+        assertEquals(MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId1));
+        assertEquals(MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId2));
         assertEquals(ConsumerGroup.ConsumerGroupState.STABLE, context.consumerGroupState(groupId));
 
         // Member 3 joins the group. This triggers the computation of a new target assignment
         // for the group. Member 3 does not get any assigned partitions yet because they are
-        // all owned by other members. However, it transitions to epoch 11 / Assigning state.
+        // all owned by other members. However, it transitions to epoch 11 and the
+        // Unreleased Partitions state.
         result = context.consumerGroupHeartbeat(
             new ConsumerGroupHeartbeatRequestData()
                 .setGroupId(groupId)
@@ -1767,21 +1762,18 @@ public void testReconciliationProcess() {
         // already covered by other tests.
         assertRecordEquals(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId3)
+                .setState(MemberState.UNRELEASED_PARTITIONS)
                 .setMemberEpoch(11)
                 .setPreviousMemberEpoch(0)
-                .setTargetMemberEpoch(11)
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 4, 5),
-                    mkTopicAssignment(barTopicId, 1)))
                 .build()),
             result.records().get(result.records().size() - 1)
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId3));
+        assertEquals(MemberState.UNRELEASED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId3));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
-        // Member 1 heartbeats. It remains at epoch 10 but transitions to Revoking state until
-        // it acknowledges the revocation of its partitions. The response contains the new
+        // Member 1 heartbeats. It remains at epoch 10 but transitions to Unrevoked Partitions
+        // state until it acknowledges the revocation of its partitions. The response contains the new
         // assignment without the partitions that must be revoked.
         result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
             .setGroupId(groupId)
@@ -1807,9 +1799,9 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
+                .setState(MemberState.UNREVOKED_PARTITIONS)
                 .setMemberEpoch(10)
-                .setPreviousMemberEpoch(9)
-                .setTargetMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(fooTopicId, 0, 1),
                     mkTopicAssignment(barTopicId, 0)))
@@ -1820,11 +1812,11 @@ public void testReconciliationProcess() {
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, context.consumerGroupMemberState(groupId, memberId1));
+        assertEquals(MemberState.UNREVOKED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId1));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
-        // Member 2 heartbeats. It remains at epoch 10 but transitions to Revoking state until
-        // it acknowledges the revocation of its partitions. The response contains the new
+        // Member 2 heartbeats. It remains at epoch 10 but transitions to Unrevoked Partitions
+        // state until it acknowledges the revocation of its partitions. The response contains the new
         // assignment without the partitions that must be revoked.
         result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
             .setGroupId(groupId)
@@ -1850,21 +1842,19 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId2)
+                .setState(MemberState.UNREVOKED_PARTITIONS)
                 .setMemberEpoch(10)
-                .setPreviousMemberEpoch(9)
-                .setTargetMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(fooTopicId, 3),
                     mkTopicAssignment(barTopicId, 2)))
                 .setPartitionsPendingRevocation(mkAssignment(
                     mkTopicAssignment(fooTopicId, 4, 5)))
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 2)))
                 .build())),
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, context.consumerGroupMemberState(groupId, memberId2));
+        assertEquals(MemberState.UNREVOKED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId2));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 3 heartbeats. The response does not contain any assignment
@@ -1882,8 +1872,16 @@ public void testReconciliationProcess() {
             result.response()
         );
 
-        assertEquals(Collections.emptyList(), result.records());
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId3));
+        assertRecordsEquals(Collections.singletonList(
+            RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId3)
+                .setState(MemberState.UNRELEASED_PARTITIONS)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(11)
+                .build())),
+            result.records()
+        );
+
+        assertEquals(MemberState.UNRELEASED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId3));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 1 acknowledges the revocation of the partitions. It does so by providing the
@@ -1921,9 +1919,9 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
+                .setState(MemberState.STABLE)
                 .setMemberEpoch(11)
                 .setPreviousMemberEpoch(10)
-                .setTargetMemberEpoch(11)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(fooTopicId, 0, 1),
                     mkTopicAssignment(barTopicId, 0)))
@@ -1931,7 +1929,7 @@ public void testReconciliationProcess() {
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId1));
+        assertEquals(MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId1));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 2 heartbeats but without acknowledging the revocation yet. This is basically a no-op.
@@ -1949,11 +1947,11 @@ public void testReconciliationProcess() {
         );
 
         assertEquals(Collections.emptyList(), result.records());
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, context.consumerGroupMemberState(groupId, memberId2));
+        assertEquals(MemberState.UNREVOKED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId2));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 3 heartbeats. It receives the partitions revoked by member 1 but remains
-        // in Assigning state because it still waits on other partitions.
+        // in Unreleased Partitions state because it still waits on other partitions.
         result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
             .setGroupId(groupId)
             .setMemberId(memberId3)
@@ -1974,18 +1972,16 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId3)
+                .setState(MemberState.UNRELEASED_PARTITIONS)
                 .setMemberEpoch(11)
                 .setPreviousMemberEpoch(11)
-                .setTargetMemberEpoch(11)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(barTopicId, 1)))
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 4, 5)))
                 .build())),
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId3));
+        assertEquals(MemberState.UNRELEASED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId3));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 3 heartbeats. Member 2 has not acknowledged the revocation of its partition so
@@ -2004,7 +2000,7 @@ public void testReconciliationProcess() {
         );
 
         assertEquals(Collections.emptyList(), result.records());
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId3));
+        assertEquals(MemberState.UNRELEASED_PARTITIONS, context.consumerGroupMemberState(groupId, memberId3));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         // Member 2 acknowledges the revocation of the partitions. It does so by providing the
@@ -2042,9 +2038,9 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId2)
+                .setState(MemberState.STABLE)
                 .setMemberEpoch(11)
                 .setPreviousMemberEpoch(10)
-                .setTargetMemberEpoch(11)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(fooTopicId, 2, 3),
                     mkTopicAssignment(barTopicId, 2)))
@@ -2052,14 +2048,19 @@ public void testReconciliationProcess() {
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId2));
+        assertEquals(MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId2));
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
-        // Member 3 heartbeats. It receives all its partitions and transitions to Stable.
+        // Member 3 heartbeats to acknowledge its current assignment. It receives all its partitions and
+        // transitions to Stable state.
         result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
             .setGroupId(groupId)
             .setMemberId(memberId3)
-            .setMemberEpoch(11));
+            .setMemberEpoch(11)
+            .setTopicPartitions(Collections.singletonList(
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(barTopicId)
+                    .setPartitions(Collections.singletonList(1)))));
 
         assertResponseEquals(
             new ConsumerGroupHeartbeatResponseData()
@@ -2079,9 +2080,9 @@ public void testReconciliationProcess() {
 
         assertRecordsEquals(Collections.singletonList(
             RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId3)
+                .setState(MemberState.STABLE)
                 .setMemberEpoch(11)
                 .setPreviousMemberEpoch(11)
-                .setTargetMemberEpoch(11)
                 .setAssignedPartitions(mkAssignment(
                     mkTopicAssignment(fooTopicId, 4, 5),
                     mkTopicAssignment(barTopicId, 1)))
@@ -2089,243 +2090,10 @@ public void testReconciliationProcess() {
             result.records()
         );
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId3));
+        assertEquals(MemberState.STABLE, context.consumerGroupMemberState(groupId, memberId3));
         assertEquals(ConsumerGroup.ConsumerGroupState.STABLE, context.consumerGroupState(groupId));
     }
 
-    @Test
-    public void testReconciliationRestartsWhenNewTargetAssignmentIsInstalled() {
-        String groupId = "fooup";
-        // Use a static member id as it makes the test easier.
-        String memberId1 = Uuid.randomUuid().toString();
-        String memberId2 = Uuid.randomUuid().toString();
-        String memberId3 = Uuid.randomUuid().toString();
-
-        Uuid fooTopicId = Uuid.randomUuid();
-        String fooTopicName = "foo";
-
-        // Create a context with one consumer group containing one member.
-        MockPartitionAssignor assignor = new MockPartitionAssignor("range");
-        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
-            .withAssignors(Collections.singletonList(assignor))
-            .withMetadataImage(new MetadataImageBuilder()
-                .addTopic(fooTopicId, fooTopicName, 6)
-                .addRacks()
-                .build())
-            .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
-                .withMember(new ConsumerGroupMember.Builder(memberId1)
-                    .setMemberEpoch(10)
-                    .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
-                    .setClientId("client")
-                    .setClientHost("localhost/127.0.0.1")
-                    .setRebalanceTimeoutMs(5000)
-                    .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
-                    .setServerAssignorName("range")
-                    .setAssignedPartitions(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0, 1, 2)))
-                    .build())
-                .withAssignment(memberId1, mkAssignment(
-                    mkTopicAssignment(fooTopicId, 0, 1, 2)))
-                .withAssignmentEpoch(10))
-            .build();
-
-        CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result;
-
-        // Prepare new assignment for the group.
-        assignor.prepareGroupAssignment(new GroupAssignment(
-            new HashMap<String, MemberAssignment>() {
-                {
-                    put(memberId1, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0, 1)
-                    )));
-                    put(memberId2, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 2)
-                    )));
-                }
-            }
-        ));
-
-        // Member 2 joins.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId2)
-                .setMemberEpoch(0)
-                .setRebalanceTimeoutMs(5000)
-                .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
-                .setServerAssignor("range")
-                .setTopicPartitions(Collections.emptyList()));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId2)
-                .setMemberEpoch(11)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()),
-            result.response()
-        );
-
-        assertRecordEquals(
-            RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId2)
-                .setMemberEpoch(11)
-                .setPreviousMemberEpoch(0)
-                .setTargetMemberEpoch(11)
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 2)))
-                .build()),
-            result.records().get(result.records().size() - 1)
-        );
-
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId2));
-
-        // Member 1 heartbeats and transitions to Revoking.
-        result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
-            .setGroupId(groupId)
-            .setMemberId(memberId1)
-            .setMemberEpoch(10));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId1)
-                .setMemberEpoch(10)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0, 1))))),
-            result.response()
-        );
-
-        assertRecordsEquals(Collections.singletonList(
-            RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
-                .setMemberEpoch(10)
-                .setPreviousMemberEpoch(9)
-                .setTargetMemberEpoch(11)
-                .setAssignedPartitions(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 0, 1)))
-                .setPartitionsPendingRevocation(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 2)))
-                .build())),
-            result.records()
-        );
-
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, context.consumerGroupMemberState(groupId, memberId1));
-
-        // Prepare new assignment for the group.
-        assignor.prepareGroupAssignment(new GroupAssignment(
-            new HashMap<String, MemberAssignment>() {
-                {
-                    put(memberId1, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0)
-                    )));
-                    put(memberId2, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 2)
-                    )));
-                    put(memberId3, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 1)
-                    )));
-                }
-            }
-        ));
-
-        // Member 3 joins.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId3)
-                .setMemberEpoch(0)
-                .setRebalanceTimeoutMs(5000)
-                .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
-                .setServerAssignor("range")
-                .setTopicPartitions(Collections.emptyList()));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId3)
-                .setMemberEpoch(12)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()),
-            result.response()
-        );
-
-        assertRecordEquals(
-            RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId3)
-                .setMemberEpoch(12)
-                .setPreviousMemberEpoch(0)
-                .setTargetMemberEpoch(12)
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 1)))
-                .build()),
-            result.records().get(result.records().size() - 1)
-        );
-
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId3));
-
-        // When member 1 heartbeats, it transitions to Revoke again but an updated state.
-        result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
-            .setGroupId(groupId)
-            .setMemberId(memberId1)
-            .setMemberEpoch(10));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId1)
-                .setMemberEpoch(10)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Collections.singletonList(0))))),
-            result.response()
-        );
-
-        assertRecordsEquals(Collections.singletonList(
-                RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
-                    .setMemberEpoch(10)
-                    .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(12)
-                    .setAssignedPartitions(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0)))
-                    .setPartitionsPendingRevocation(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 1, 2)))
-                    .build())),
-            result.records()
-        );
-
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, context.consumerGroupMemberState(groupId, memberId1));
-
-        // When member 2 heartbeats, it transitions to Assign again but with an updated state.
-        result = context.consumerGroupHeartbeat(new ConsumerGroupHeartbeatRequestData()
-            .setGroupId(groupId)
-            .setMemberId(memberId2)
-            .setMemberEpoch(11));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId2)
-                .setMemberEpoch(12)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()),
-            result.response()
-        );
-
-        assertRecordsEquals(Collections.singletonList(
-            RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId2)
-                .setMemberEpoch(12)
-                .setPreviousMemberEpoch(11)
-                .setTargetMemberEpoch(12)
-                .setPartitionsPendingAssignment(mkAssignment(
-                    mkTopicAssignment(fooTopicId, 2)))
-                .build())),
-            result.records()
-        );
-
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, context.consumerGroupMemberState(groupId, memberId2));
-    }
-
     @Test
     public void testNewMemberIsRejectedWithMaximumMembersIsReached() {
         String groupId = "fooup";
@@ -2350,9 +2118,9 @@ public void testNewMemberIsRejectedWithMaximumMembersIsReached() {
             .withConsumerGroupMaxSize(2)
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId1)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -2363,9 +2131,9 @@ public void testNewMemberIsRejectedWithMaximumMembersIsReached() {
                         mkTopicAssignment(barTopicId, 0, 1)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder(memberId2)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -2412,6 +2180,7 @@ public void testConsumerGroupStates() {
         assertEquals(ConsumerGroup.ConsumerGroupState.EMPTY, context.consumerGroupState(groupId));
 
         context.replay(RecordHelpers.newMemberSubscriptionRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
+            .setState(MemberState.STABLE)
             .setSubscribedTopicNames(Collections.singletonList(fooTopicName))
             .build()));
         context.replay(RecordHelpers.newGroupEpochRecord(groupId, 11));
@@ -2425,19 +2194,18 @@ public void testConsumerGroupStates() {
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         context.replay(RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
+            .setState(MemberState.UNREVOKED_PARTITIONS)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
-            .setAssignedPartitions(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2)))
-            .setPartitionsPendingAssignment(mkAssignment(mkTopicAssignment(fooTopicId, 3)))
+            .setAssignedPartitions(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2, 3)))
             .build()));
 
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, context.consumerGroupState(groupId));
 
         context.replay(RecordHelpers.newCurrentAssignmentRecord(groupId, new ConsumerGroupMember.Builder(memberId1)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2, 3)))
             .build()));
 
@@ -2502,9 +2270,9 @@ public void testSubscriptionMetadataRefreshedAfterGroupIsLoaded() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(10)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -2560,9 +2328,9 @@ public void testSubscriptionMetadataRefreshedAfterGroupIsLoaded() {
         );
 
         ConsumerGroupMember expectedMember = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -2613,9 +2381,9 @@ public void testSubscriptionMetadataRefreshedAgainAfterWriteFailure() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10)
                 .withMember(new ConsumerGroupMember.Builder(memberId)
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(10)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setRebalanceTimeoutMs(5000)
@@ -2689,9 +2457,9 @@ public void testSubscriptionMetadataRefreshedAgainAfterWriteFailure() {
         );
 
         ConsumerGroupMember expectedMember = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setClientId("client")
             .setClientHost("localhost/127.0.0.1")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
@@ -2998,7 +2766,7 @@ public void testSessionTimeoutLifecycle() {
 
         // Verify that there are no timers.
         context.assertNoSessionTimeout(groupId, memberId);
-        context.assertNoRevocationTimeout(groupId, memberId);
+        context.assertNoRebalanceTimeout(groupId, memberId);
     }
 
     @Test
@@ -3062,7 +2830,7 @@ public void testSessionTimeoutExpiration() {
 
         // Verify that there are no timers.
         context.assertNoSessionTimeout(groupId, memberId);
-        context.assertNoRevocationTimeout(groupId, memberId);
+        context.assertNoRebalanceTimeout(groupId, memberId);
     }
 
     @Test
@@ -3143,16 +2911,15 @@ public void testSessionTimeoutExpirationStaticMember() {
 
         // Verify that there are no timers.
         context.assertNoSessionTimeout(groupId, memberId);
-        context.assertNoRevocationTimeout(groupId, memberId);
+        context.assertNoRebalanceTimeout(groupId, memberId);
     }
 
     @Test
-    public void testRevocationTimeoutLifecycle() {
+    public void testRebalanceTimeoutLifecycle() {
         String groupId = "fooup";
         // Use a static member id as it makes the test easier.
         String memberId1 = Uuid.randomUuid().toString();
         String memberId2 = Uuid.randomUuid().toString();
-        String memberId3 = Uuid.randomUuid().toString();
 
         Uuid fooTopicId = Uuid.randomUuid();
         String fooTopicName = "foo";
@@ -3243,7 +3010,7 @@ public void testRevocationTimeoutLifecycle() {
             context.sleep(result.response().heartbeatIntervalMs())
         );
 
-        // Member 1 heartbeats and transitions to revoking. The revocation timeout
+        // Member 1 heartbeats and transitions to unrevoked partitions. The rebalance timeout
         // is scheduled.
         result = context.consumerGroupHeartbeat(
             new ConsumerGroupHeartbeatRequestData()
@@ -3266,82 +3033,10 @@ public void testRevocationTimeoutLifecycle() {
             result.response()
         );
 
-        // Verify that there is a revocation timeout.
-        context.assertRevocationTimeout(groupId, memberId1, 12000);
-
-        assertEquals(
-            Collections.emptyList(),
-            context.sleep(result.response().heartbeatIntervalMs())
-        );
-
-        // Prepare next assignment.
-        assignor.prepareGroupAssignment(new GroupAssignment(
-            new HashMap<String, MemberAssignment>() {
-                {
-                    put(memberId1, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0)
-                    )));
-                    put(memberId2, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 2)
-                    )));
-                    put(memberId3, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 1)
-                    )));
-                }
-            }
-        ));
-
-        // Member 3 joins the group.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId3)
-                .setMemberEpoch(0)
-                .setRebalanceTimeoutMs(90000)
-                .setSubscribedTopicNames(Collections.singletonList("foo"))
-                .setTopicPartitions(Collections.emptyList()));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId3)
-                .setMemberEpoch(3)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()),
-            result.response()
-        );
-
-        assertEquals(
-            Collections.emptyList(),
-            context.sleep(result.response().heartbeatIntervalMs())
-        );
-
-        // Member 1 heartbeats and re-transitions to revoking. The revocation timeout
-        // is re-scheduled.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId1)
-                .setMemberEpoch(1)
-                .setRebalanceTimeoutMs(90000)
-                .setSubscribedTopicNames(Collections.singletonList("foo")));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId1)
-                .setMemberEpoch(1)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Collections.singletonList(0))))),
-            result.response()
-        );
-
         // Verify that there is a revocation timeout. Keep a reference
         // to the timeout for later.
         ScheduledTimeout<Void, Record> scheduledTimeout =
-            context.assertRevocationTimeout(groupId, memberId1, 90000);
+            context.assertRebalanceTimeout(groupId, memberId1, 12000);
 
         assertEquals(
             Collections.emptyList(),
@@ -3356,23 +3051,23 @@ public void testRevocationTimeoutLifecycle() {
                 .setMemberEpoch(1)
                 .setTopicPartitions(Collections.singletonList(new ConsumerGroupHeartbeatRequestData.TopicPartitions()
                     .setTopicId(fooTopicId)
-                    .setPartitions(Collections.singletonList(0)))));
+                    .setPartitions(Arrays.asList(0, 1)))));
 
         assertResponseEquals(
             new ConsumerGroupHeartbeatResponseData()
                 .setMemberId(memberId1)
-                .setMemberEpoch(3)
+                .setMemberEpoch(2)
                 .setHeartbeatIntervalMs(5000)
                 .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
                     .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                             .setTopicId(fooTopicId)
-                            .setPartitions(Collections.singletonList(0))))),
+                            .setPartitions(Arrays.asList(0, 1))))),
             result.response()
         );
 
         // Verify that there is not revocation timeout.
-        context.assertNoRevocationTimeout(groupId, memberId1);
+        context.assertNoRebalanceTimeout(groupId, memberId1);
 
         // Execute the scheduled revocation timeout captured earlier to simulate a
         // stale timeout. This should be a no-op.
@@ -3380,7 +3075,7 @@ public void testRevocationTimeoutLifecycle() {
     }
 
     @Test
-    public void testRevocationTimeoutExpiration() {
+    public void testRebalanceTimeoutExpiration() {
         String groupId = "fooup";
         // Use a static member id as it makes the test easier.
         String memberId1 = Uuid.randomUuid().toString();
@@ -3502,7 +3197,7 @@ public void testRevocationTimeoutExpiration() {
         // Verify the expired timeout.
         assertEquals(
             Collections.singletonList(new ExpiredTimeout<Void, Record>(
-                consumerGroupRevocationTimeoutKey(groupId, memberId1),
+                consumerGroupRebalanceTimeoutKey(groupId, memberId1),
                 new CoordinatorResult<>(
                     Arrays.asList(
                         RecordHelpers.newCurrentAssignmentTombstoneRecord(groupId, memberId1),
@@ -3517,7 +3212,7 @@ public void testRevocationTimeoutExpiration() {
 
         // Verify that there are no timers.
         context.assertNoSessionTimeout(groupId, memberId1);
-        context.assertNoRevocationTimeout(groupId, memberId1);
+        context.assertNoRebalanceTimeout(groupId, memberId1);
     }
 
     @Test
@@ -3535,9 +3230,9 @@ public void testOnLoaded() {
                 .build())
             .withConsumerGroup(new ConsumerGroupBuilder("foo", 10)
                 .withMember(new ConsumerGroupMember.Builder("foo-1")
+                    .setState(MemberState.UNREVOKED_PARTITIONS)
                     .setMemberEpoch(9)
                     .setPreviousMemberEpoch(9)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Collections.singletonList("foo"))
@@ -3548,19 +3243,15 @@ public void testOnLoaded() {
                         mkTopicAssignment(fooTopicId, 3, 4, 5)))
                     .build())
                 .withMember(new ConsumerGroupMember.Builder("foo-2")
+                    .setState(MemberState.STABLE)
                     .setMemberEpoch(10)
                     .setPreviousMemberEpoch(10)
-                    .setTargetMemberEpoch(10)
                     .setClientId("client")
                     .setClientHost("localhost/127.0.0.1")
                     .setSubscribedTopicNames(Collections.singletonList("foo"))
                     .setServerAssignorName("range")
-                    .setPartitionsPendingAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 3, 4, 5)))
                     .build())
                 .withAssignment("foo-1", mkAssignment(
-                    mkTopicAssignment(fooTopicId, 0, 1, 2)))
-                .withAssignment("foo-2", mkAssignment(
                     mkTopicAssignment(fooTopicId, 3, 4, 5)))
                 .withAssignmentEpoch(10))
             .build();
@@ -3574,7 +3265,7 @@ public void testOnLoaded() {
         assertNotNull(context.timer.timeout(consumerGroupSessionTimeoutKey("foo", "foo-2")));
 
         // foo-1 should also have a revocation timeout in place.
-        assertNotNull(context.timer.timeout(consumerGroupRevocationTimeoutKey("foo", "foo-1")));
+        assertNotNull(context.timer.timeout(consumerGroupRebalanceTimeoutKey("foo", "foo-1")));
     }
 
     @Test
@@ -9452,8 +9143,8 @@ public void testConsumerGroupMaybeDelete() {
 
         records = new ArrayList<>();
         group.updateMember(new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
index ad70107506ccb..a02359ffed17c 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -47,7 +47,7 @@
 import org.apache.kafka.coordinator.group.classic.ClassicGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
-import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
+import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
@@ -85,7 +85,7 @@
 import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.EMPTY_RESULT;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRevocationTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRebalanceTimeoutKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupSessionTimeoutKey;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.COMPLETING_REBALANCE;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.DEAD;
@@ -487,7 +487,7 @@ public ConsumerGroup.ConsumerGroupState consumerGroupState(
             .state();
     }
 
-    public ConsumerGroupMember.MemberState consumerGroupMemberState(
+    public MemberState consumerGroupMemberState(
         String groupId,
         String memberId
     ) {
@@ -556,24 +556,24 @@ public void assertNoSessionTimeout(
         assertNull(timeout);
     }
 
-    public MockCoordinatorTimer.ScheduledTimeout<Void, Record> assertRevocationTimeout(
+    public MockCoordinatorTimer.ScheduledTimeout<Void, Record> assertRebalanceTimeout(
         String groupId,
         String memberId,
         long delayMs
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
+            timer.timeout(consumerGroupRebalanceTimeoutKey(groupId, memberId));
         assertNotNull(timeout);
         assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
         return timeout;
     }
 
-    public void assertNoRevocationTimeout(
+    public void assertNoRebalanceTimeout(
         String groupId,
         String memberId
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupRevocationTimeoutKey(groupId, memberId));
+            timer.timeout(consumerGroupRebalanceTimeoutKey(groupId, memberId));
         assertNull(timeout);
     }
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
index 87f7b471035fc..d9eacb46c3ea3 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
@@ -1116,7 +1116,6 @@ public void testConsumerGroupOffsetCommitWithStaleMemberEpoch() {
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
@@ -1162,7 +1161,6 @@ public void testConsumerGroupOffsetCommitWithVersionSmallerThanVersion9(short ve
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
@@ -1256,7 +1254,6 @@ public void testConsumerGroupOffsetCommit() {
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
@@ -1327,7 +1324,6 @@ public void testConsumerGroupOffsetCommitWithOffsetMetadataTooLarge() {
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
@@ -1405,7 +1401,6 @@ public void testConsumerGroupTransactionalOffsetCommit() {
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
@@ -1525,7 +1520,6 @@ public void testConsumerGroupTransactionalOffsetCommitWithStaleMemberEpoch() {
         // Add member.
         group.updateMember(new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setPreviousMemberEpoch(10)
             .build()
         );
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
index 2a15fb7decd9c..50bfcbd52356c 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
@@ -24,6 +24,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.coordinator.group.consumer.ClientAssignor;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
+import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
 import org.apache.kafka.coordinator.group.consumer.VersionedMetadata;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
@@ -397,11 +398,6 @@ public void testNewCurrentAssignmentRecord() {
             mkSortedTopicAssignment(topicId2, 24, 25, 26)
         );
 
-        Map<Uuid, Set<Integer>> assigning = mkSortedAssignment(
-            mkSortedTopicAssignment(topicId1, 17, 18, 19),
-            mkSortedTopicAssignment(topicId2, 27, 28, 29)
-        );
-
         Record expectedRecord = new Record(
             new ApiMessageAndVersion(
                 new ConsumerGroupCurrentMemberAssignmentKey()
@@ -410,9 +406,9 @@ public void testNewCurrentAssignmentRecord() {
                 (short) 8),
             new ApiMessageAndVersion(
                 new ConsumerGroupCurrentMemberAssignmentValue()
+                    .setState(MemberState.UNREVOKED_PARTITIONS.value())
                     .setMemberEpoch(22)
                     .setPreviousMemberEpoch(21)
-                    .setTargetMemberEpoch(23)
                     .setAssignedPartitions(Arrays.asList(
                         new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                             .setTopicId(topicId1)
@@ -426,25 +422,17 @@ public void testNewCurrentAssignmentRecord() {
                             .setPartitions(Arrays.asList(14, 15, 16)),
                         new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                             .setTopicId(topicId2)
-                            .setPartitions(Arrays.asList(24, 25, 26))))
-                    .setPartitionsPendingAssignment(Arrays.asList(
-                        new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
-                            .setTopicId(topicId1)
-                            .setPartitions(Arrays.asList(17, 18, 19)),
-                        new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
-                            .setTopicId(topicId2)
-                            .setPartitions(Arrays.asList(27, 28, 29)))),
+                            .setPartitions(Arrays.asList(24, 25, 26)))),
                 (short) 0));
 
         assertEquals(expectedRecord, newCurrentAssignmentRecord(
             "group-id",
             new ConsumerGroupMember.Builder("member-id")
+                .setState(MemberState.UNREVOKED_PARTITIONS)
                 .setMemberEpoch(22)
                 .setPreviousMemberEpoch(21)
-                .setTargetMemberEpoch(23)
                 .setAssignedPartitions(assigned)
                 .setPartitionsPendingRevocation(revoking)
-                .setPartitionsPendingAssignment(assigning)
                 .build()
         ));
     }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
index a5daf3f92b372..3baa6f77af4ee 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
@@ -51,7 +51,6 @@ public void testNewMember() {
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member-id")
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(11)
             .setInstanceId("instance-id")
             .setRackId("rack-id")
             .setRebalanceTimeoutMs(5000)
@@ -73,14 +72,11 @@ public void testNewMember() {
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(topicId2, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId3, 7, 8, 9)))
             .build();
 
         assertEquals("member-id", member.memberId());
         assertEquals(10, member.memberEpoch());
         assertEquals(9, member.previousMemberEpoch());
-        assertEquals(11, member.targetMemberEpoch());
         assertEquals("instance-id", member.instanceId());
         assertEquals("rack-id", member.rackId());
         assertEquals("client-id", member.clientId());
@@ -102,7 +98,6 @@ public void testNewMember() {
             member.clientAssignors());
         assertEquals(mkAssignment(mkTopicAssignment(topicId1, 1, 2, 3)), member.assignedPartitions());
         assertEquals(mkAssignment(mkTopicAssignment(topicId2, 4, 5, 6)), member.partitionsPendingRevocation());
-        assertEquals(mkAssignment(mkTopicAssignment(topicId3, 7, 8, 9)), member.partitionsPendingAssignment());
     }
 
     @Test
@@ -114,7 +109,6 @@ public void testEquals() {
         ConsumerGroupMember member1 = new ConsumerGroupMember.Builder("member-id")
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(11)
             .setInstanceId("instance-id")
             .setRackId("rack-id")
             .setRebalanceTimeoutMs(5000)
@@ -136,14 +130,11 @@ public void testEquals() {
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(topicId2, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId3, 7, 8, 9)))
             .build();
 
         ConsumerGroupMember member2 = new ConsumerGroupMember.Builder("member-id")
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(11)
             .setInstanceId("instance-id")
             .setRackId("rack-id")
             .setRebalanceTimeoutMs(5000)
@@ -165,8 +156,6 @@ public void testEquals() {
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(topicId2, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId3, 7, 8, 9)))
             .build();
 
         assertEquals(member1, member2);
@@ -181,7 +170,6 @@ public void testUpdateMember() {
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member-id")
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(11)
             .setInstanceId("instance-id")
             .setRackId("rack-id")
             .setRebalanceTimeoutMs(5000)
@@ -203,8 +191,6 @@ public void testUpdateMember() {
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(topicId2, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId3, 7, 8, 9)))
             .build();
 
         // This is a no-op.
@@ -291,16 +277,12 @@ public void testUpdateWithConsumerGroupCurrentMemberAssignmentValue() {
         ConsumerGroupCurrentMemberAssignmentValue record = new ConsumerGroupCurrentMemberAssignmentValue()
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(9)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                 .setTopicId(topicId1)
                 .setPartitions(Arrays.asList(0, 1, 2))))
             .setPartitionsPendingRevocation(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                 .setTopicId(topicId2)
-                .setPartitions(Arrays.asList(3, 4, 5))))
-            .setPartitionsPendingAssignment(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
-                .setTopicId(topicId3)
-                .setPartitions(Arrays.asList(6, 7, 8))));
+                .setPartitions(Arrays.asList(3, 4, 5))));
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member-id")
             .updateWith(record)
@@ -308,10 +290,8 @@ public void testUpdateWithConsumerGroupCurrentMemberAssignmentValue() {
 
         assertEquals(10, member.memberEpoch());
         assertEquals(9, member.previousMemberEpoch());
-        assertEquals(11, member.targetMemberEpoch());
         assertEquals(mkAssignment(mkTopicAssignment(topicId1, 0, 1, 2)), member.assignedPartitions());
         assertEquals(mkAssignment(mkTopicAssignment(topicId2, 3, 4, 5)), member.partitionsPendingRevocation());
-        assertEquals(mkAssignment(mkTopicAssignment(topicId3, 6, 7, 8)), member.partitionsPendingAssignment());
     }
 
     @Test
@@ -331,16 +311,12 @@ public void testAsConsumerGroupDescribeMember() {
         ConsumerGroupCurrentMemberAssignmentValue record = new ConsumerGroupCurrentMemberAssignmentValue()
             .setMemberEpoch(epoch)
             .setPreviousMemberEpoch(epoch - 1)
-            .setTargetMemberEpoch(epoch + 1)
             .setAssignedPartitions(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                 .setTopicId(topicId1)
                 .setPartitions(assignedPartitions)))
             .setPartitionsPendingRevocation(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                 .setTopicId(topicId2)
-                .setPartitions(Arrays.asList(3, 4, 5))))
-            .setPartitionsPendingAssignment(Collections.singletonList(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
-                .setTopicId(topicId3)
-                .setPartitions(Arrays.asList(6, 7, 8))));
+                .setPartitions(Arrays.asList(3, 4, 5))));
         String memberId = Uuid.randomUuid().toString();
         String clientId = "clientId";
         String instanceId = "instanceId";
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
index 7d90122c2dfb4..1526152706751 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
@@ -177,8 +177,6 @@ public void testUpdatingMemberUpdatesPartitionEpoch() {
                 mkTopicAssignment(fooTopicId, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(barTopicId, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(zarTopicId, 7, 8, 9)))
             .build();
 
         consumerGroup.updateMember(member);
@@ -199,8 +197,6 @@ public void testUpdatingMemberUpdatesPartitionEpoch() {
                 mkTopicAssignment(barTopicId, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(zarTopicId, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(fooTopicId, 7, 8, 9)))
             .build();
 
         consumerGroup.updateMember(member);
@@ -337,8 +333,6 @@ public void testDeletingMemberRemovesPartitionEpoch() {
                 mkTopicAssignment(fooTopicId, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
                 mkTopicAssignment(barTopicId, 4, 5, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(zarTopicId, 7, 8, 9)))
             .build();
 
         consumerGroup.updateMember(member);
@@ -373,27 +367,27 @@ public void testGroupState() {
         assertEquals(ConsumerGroup.ConsumerGroupState.EMPTY, consumerGroup.state());
 
         ConsumerGroupMember member1 = new ConsumerGroupMember.Builder("member1")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(1)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(1)
             .build();
 
         consumerGroup.updateMember(member1);
         consumerGroup.setGroupEpoch(1);
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member1.state());
+        assertEquals(MemberState.STABLE, member1.state());
         assertEquals(ConsumerGroup.ConsumerGroupState.ASSIGNING, consumerGroup.state());
 
         ConsumerGroupMember member2 = new ConsumerGroupMember.Builder("member2")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(1)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(1)
             .build();
 
         consumerGroup.updateMember(member2);
         consumerGroup.setGroupEpoch(2);
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member2.state());
+        assertEquals(MemberState.STABLE, member2.state());
         assertEquals(ConsumerGroup.ConsumerGroupState.ASSIGNING, consumerGroup.state());
 
         consumerGroup.setTargetAssignmentEpoch(2);
@@ -401,42 +395,37 @@ public void testGroupState() {
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, consumerGroup.state());
 
         member1 = new ConsumerGroupMember.Builder(member1)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(2)
             .setPreviousMemberEpoch(1)
-            .setTargetMemberEpoch(2)
             .build();
 
         consumerGroup.updateMember(member1);
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member1.state());
+        assertEquals(MemberState.STABLE, member1.state());
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, consumerGroup.state());
 
         // Member 2 is not stable so the group stays in reconciling state.
         member2 = new ConsumerGroupMember.Builder(member2)
+            .setState(MemberState.UNREVOKED_PARTITIONS)
             .setMemberEpoch(2)
             .setPreviousMemberEpoch(1)
-            .setTargetMemberEpoch(2)
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(fooTopicId, 0)))
             .build();
 
         consumerGroup.updateMember(member2);
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, member2.state());
+        assertEquals(MemberState.UNREVOKED_PARTITIONS, member2.state());
         assertEquals(ConsumerGroup.ConsumerGroupState.RECONCILING, consumerGroup.state());
 
         member2 = new ConsumerGroupMember.Builder(member2)
+            .setState(MemberState.STABLE)
             .setMemberEpoch(2)
             .setPreviousMemberEpoch(1)
-            .setTargetMemberEpoch(2)
-            .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(fooTopicId, 0)))
-            .setPartitionsPendingAssignment(Collections.emptyMap())
             .build();
 
         consumerGroup.updateMember(member2);
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member2.state());
+        assertEquals(MemberState.STABLE, member2.state());
         assertEquals(ConsumerGroup.ConsumerGroupState.STABLE, consumerGroup.state());
 
         consumerGroup.removeMember("member1");
@@ -882,7 +871,6 @@ public void testValidateDeleteGroup() {
         ConsumerGroupMember member1 = new ConsumerGroupMember.Builder("member1")
             .setMemberEpoch(1)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(1)
             .build();
         consumerGroup.updateMember(member1);
 
@@ -1014,7 +1002,6 @@ public void testStateTransitionMetrics() {
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(1)
             .setPreviousMemberEpoch(0)
-            .setTargetMemberEpoch(1)
             .build();
 
         consumerGroup.updateMember(member);
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java
index 037a6ccbcd463..c1db6f228900b 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/CurrentAssignmentBuilderTest.java
@@ -17,542 +17,522 @@
 package org.apache.kafka.coordinator.group.consumer;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.FencedMemberEpochException;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.Arguments;
-import org.junit.jupiter.params.provider.MethodSource;
 
-import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.stream.Stream;
 
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class CurrentAssignmentBuilderTest {
 
     @Test
-    public void testTransitionFromNewTargetToRevoke() {
+    public void testStableToStable() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3),
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 1, 2, 3),
+                mkTopicAssignment(topicId2, 4, 5, 6))))
             .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(10, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3),
-            mkTopicAssignment(topicId2, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2),
-            mkTopicAssignment(topicId2, 4, 5)
-        ), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 4, 5),
-            mkTopicAssignment(topicId2, 7, 8)
-        ), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 1, 2, 3),
+                    mkTopicAssignment(topicId2, 4, 5, 6)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromNewTargetToAssigning() {
+    public void testStableToStableWithNewPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3),
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3, 4, 5),
-            mkTopicAssignment(topicId2, 4, 5, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 4, 5, 6, 7))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 4, 5),
-            mkTopicAssignment(topicId2, 7, 8)
-        ), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 1, 2, 3, 4),
+                    mkTopicAssignment(topicId2, 4, 5, 6, 7)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromNewTargetToStable() {
+    public void testStableToUnrevokedPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.STABLE)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(10)
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3),
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member.state());
+        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 5, 6, 7))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
+            .build();
 
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ));
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.UNREVOKED_PARTITIONS)
+                .setMemberEpoch(10)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 2, 3),
+                    mkTopicAssignment(topicId2, 5, 6)))
+                .setPartitionsPendingRevocation(mkAssignment(
+                    mkTopicAssignment(topicId1, 1),
+                    mkTopicAssignment(topicId2, 4)))
+                .build(),
+            updatedMember
+        );
+    }
+
+    @Test
+    public void testStableToUnreleasedPartitions() {
+        Uuid topicId1 = Uuid.randomUuid();
+        Uuid topicId2 = Uuid.randomUuid();
+
+        ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.STABLE)
+            .setMemberEpoch(10)
+            .setPreviousMemberEpoch(10)
+            .setAssignedPartitions(mkAssignment(
+                mkTopicAssignment(topicId1, 1, 2, 3),
+                mkTopicAssignment(topicId2, 4, 5, 6)))
+            .build();
 
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 4, 5, 6, 7))))
             .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingAssignment());
-    }
-
-    private static Stream<Arguments> ownedTopicPartitionsArguments() {
-        return Stream.of(
-            // Field not set in the heartbeat request.
-            null,
-            // Owned partitions does not match the assigned partitions.
-            Collections.emptyList()
-        ).map(Arguments::of);
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.UNRELEASED_PARTITIONS)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 1, 2, 3),
+                    mkTopicAssignment(topicId2, 4, 5, 6)))
+                .build(),
+            updatedMember
+        );
     }
 
-    @ParameterizedTest
-    @MethodSource("ownedTopicPartitionsArguments")
-    public void testTransitionFromRevokeToRevoke(
-        List<ConsumerGroupHeartbeatRequestData.TopicPartitions> ownedTopicPartitions
-    ) {
+    @Test
+    public void testUnrevokedPartitionsToStable() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.UNREVOKED_PARTITIONS)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .setPartitionsPendingRevocation(mkAssignment(
-                mkTopicAssignment(topicId1, 1, 2),
-                mkTopicAssignment(topicId2, 4, 5)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 1),
+                mkTopicAssignment(topicId2, 4)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6))))
             .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
-            .withOwnedTopicPartitions(ownedTopicPartitions)
+            .withOwnedTopicPartitions(Arrays.asList(
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId1)
+                    .setPartitions(Arrays.asList(2, 3)),
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId2)
+                    .setPartitions(Arrays.asList(5, 6))))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(10, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3),
-            mkTopicAssignment(topicId2, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2),
-            mkTopicAssignment(topicId2, 4, 5)
-        ), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 4, 5),
-            mkTopicAssignment(topicId2, 7, 8)
-        ), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 2, 3),
+                    mkTopicAssignment(topicId2, 5, 6)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromRevokeToAssigning() {
+    public void testRemainsInUnrevokedPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.UNREVOKED_PARTITIONS)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .setPartitionsPendingRevocation(mkAssignment(
-                mkTopicAssignment(topicId1, 1, 2),
-                mkTopicAssignment(topicId2, 4, 5)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 1),
+                mkTopicAssignment(topicId2, 4)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
-            .withOwnedTopicPartitions(requestFromAssignment(mkAssignment(
+        CurrentAssignmentBuilder currentAssignmentBuilder = new CurrentAssignmentBuilder(member)
+            .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6))))
-            .build();
-
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3),
-            mkTopicAssignment(topicId2, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 4, 5),
-            mkTopicAssignment(topicId2, 7, 8)
-        ), updatedMember.partitionsPendingAssignment());
+            .withCurrentPartitionEpoch((topicId, partitionId) -> -1);
+
+        assertEquals(
+            member,
+            currentAssignmentBuilder
+                .withOwnedTopicPartitions(null)
+                .build()
+        );
+
+        assertEquals(
+            member,
+            currentAssignmentBuilder
+                .withOwnedTopicPartitions(Arrays.asList(
+                    new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                        .setTopicId(topicId1)
+                        .setPartitions(Arrays.asList(1, 2, 3)),
+                    new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                        .setTopicId(topicId2)
+                        .setPartitions(Arrays.asList(5, 6))))
+                .build()
+        );
+
+        assertEquals(
+            member,
+            currentAssignmentBuilder
+                .withOwnedTopicPartitions(Arrays.asList(
+                    new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                        .setTopicId(topicId1)
+                        .setPartitions(Arrays.asList(2, 3)),
+                    new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                        .setTopicId(topicId2)
+                        .setPartitions(Arrays.asList(4, 5, 6))))
+                .build()
+        );
     }
 
     @Test
-    public void testTransitionFromRevokeToStable() {
+    public void testUnrevokedPartitionsToUnrevokedPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.UNREVOKED_PARTITIONS)
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .setPartitionsPendingRevocation(mkAssignment(
-                mkTopicAssignment(topicId1, 1, 2),
-                mkTopicAssignment(topicId2, 4, 5)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 1),
+                mkTopicAssignment(topicId2, 4)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
-            .withOwnedTopicPartitions(requestFromAssignment(mkAssignment(
+            .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
+            .withOwnedTopicPartitions(Arrays.asList(
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId1)
+                    .setPartitions(Arrays.asList(2, 3)),
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId2)
+                    .setPartitions(Arrays.asList(5, 6))))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.UNREVOKED_PARTITIONS)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(10)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 3),
+                    mkTopicAssignment(topicId2, 6)))
+                .setPartitionsPendingRevocation(mkAssignment(
+                    mkTopicAssignment(topicId1, 2),
+                    mkTopicAssignment(topicId2, 5)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromRevokeToStableWhenPartitionsPendingRevocationAreReassignedBeforeBeingRevoked() {
+    public void testUnrevokedPartitionsToUnreleasedPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
-            .setMemberEpoch(10)
+            .setState(MemberState.UNREVOKED_PARTITIONS)
+            .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
-            .setPartitionsPendingRevocation(mkAssignment(
-                mkTopicAssignment(topicId1, 1, 2),
-                mkTopicAssignment(topicId2, 4, 5)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, member.state());
-
-        // A new target assignment is computed (epoch 12) before the partitions
-        // pending revocation are revoked by the member and those partitions
-        // have been reassigned to the member. In this case, the member
-        // can keep them a jump to epoch 12.
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(12, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 5, 6, 7))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
+            .withOwnedTopicPartitions(Arrays.asList(
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId1)
+                    .setPartitions(Arrays.asList(2, 3)),
+                new ConsumerGroupHeartbeatRequestData.TopicPartitions()
+                    .setTopicId(topicId2)
+                    .setPartitions(Arrays.asList(5, 6))))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(12, updatedMember.memberEpoch());
-        assertEquals(12, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.UNRELEASED_PARTITIONS)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(11)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 2, 3),
+                    mkTopicAssignment(topicId2, 5, 6)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromAssigningToAssigning() {
+    public void testUnreleasedPartitionsToStable() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
-            .setMemberEpoch(10)
+            .setState(MemberState.UNRELEASED_PARTITIONS)
+            .setMemberEpoch(11)
             .setPreviousMemberEpoch(11)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> {
-                if (topicId.equals(topicId1))
-                    return -1;
-                else
-                    return 10;
-            })
+            .withTargetAssignment(12, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId2, 7, 8)
-        ), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(12)
+                .setPreviousMemberEpoch(11)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 2, 3),
+                    mkTopicAssignment(topicId2, 5, 6)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromAssigningToStable() {
+    public void testUnreleasedPartitionsToStableWithNewPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
-            .setMemberEpoch(10)
+            .setState(MemberState.UNRELEASED_PARTITIONS)
+            .setMemberEpoch(11)
             .setPreviousMemberEpoch(11)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3),
-                mkTopicAssignment(topicId2, 6)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.ASSIGNING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
-
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 5, 6, 7))))
             .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(11)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 2, 3, 4),
+                    mkTopicAssignment(topicId2, 5, 6, 7)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testTransitionFromStableToStable() {
+    public void testUnreleasedPartitionsToUnreleasedPartitions() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.UNRELEASED_PARTITIONS)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(11)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
-                mkTopicAssignment(topicId1, 3, 4, 5),
-                mkTopicAssignment(topicId2, 6, 7, 8)))
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
+            .build();
+
+        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+            .withTargetAssignment(11, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3, 4),
+                mkTopicAssignment(topicId2, 5, 6, 7))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, member.state());
+        assertEquals(member, updatedMember);
+    }
 
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ));
+    @Test
+    public void testUnreleasedPartitionsToUnrevokedPartitions() {
+        Uuid topicId1 = Uuid.randomUuid();
+        Uuid topicId2 = Uuid.randomUuid();
+
+        ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
+            .setState(MemberState.UNRELEASED_PARTITIONS)
+            .setMemberEpoch(11)
+            .setPreviousMemberEpoch(11)
+            .setAssignedPartitions(mkAssignment(
+                mkTopicAssignment(topicId1, 2, 3),
+                mkTopicAssignment(topicId2, 5, 6)))
+            .build();
 
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(11, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
+            .withTargetAssignment(12, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 3),
+                mkTopicAssignment(topicId2, 6))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.STABLE, updatedMember.state());
-        assertEquals(11, updatedMember.previousMemberEpoch());
-        assertEquals(11, updatedMember.memberEpoch());
-        assertEquals(11, updatedMember.targetMemberEpoch());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 3, 4, 5),
-            mkTopicAssignment(topicId2, 6, 7, 8)
-        ), updatedMember.assignedPartitions());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingRevocation());
-        assertEquals(Collections.emptyMap(), updatedMember.partitionsPendingAssignment());
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.UNREVOKED_PARTITIONS)
+                .setMemberEpoch(11)
+                .setPreviousMemberEpoch(11)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 3),
+                    mkTopicAssignment(topicId2, 6)))
+                .setPartitionsPendingRevocation(mkAssignment(
+                    mkTopicAssignment(topicId1, 2),
+                    mkTopicAssignment(topicId2, 5)))
+                .build(),
+            updatedMember
+        );
     }
 
     @Test
-    public void testNewTargetRestartReconciliation() {
+    public void testUnknownState() {
         Uuid topicId1 = Uuid.randomUuid();
         Uuid topicId2 = Uuid.randomUuid();
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member")
-            .setMemberEpoch(10)
-            .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
+            .setState(MemberState.UNKNOWN)
+            .setMemberEpoch(11)
+            .setPreviousMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6)))
             .setPartitionsPendingRevocation(mkAssignment(
-                mkTopicAssignment(topicId1, 1, 2),
-                mkTopicAssignment(topicId2, 4, 5)))
-            .setPartitionsPendingAssignment(mkAssignment(
-                mkTopicAssignment(topicId1, 4, 5),
-                mkTopicAssignment(topicId2, 7, 8)))
+                mkTopicAssignment(topicId1, 2),
+                mkTopicAssignment(topicId2, 5)))
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, member.state());
-
-        Assignment targetAssignment = new Assignment(mkAssignment(
-            mkTopicAssignment(topicId1, 6, 7, 8),
-            mkTopicAssignment(topicId2, 9, 10, 11)
-        ));
+        // When the member is in an unknown state, the member is first to force
+        // a reset of the client side member state.
+        assertThrows(FencedMemberEpochException.class, () -> new CurrentAssignmentBuilder(member)
+            .withTargetAssignment(12, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 3),
+                mkTopicAssignment(topicId2, 6))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 10)
+            .build());
 
+        // Then the member rejoins with no owned partitions.
         ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(12, targetAssignment)
-            .withCurrentPartitionEpoch((topicId, partitionId) -> -1)
+            .withTargetAssignment(12, new Assignment(mkAssignment(
+                mkTopicAssignment(topicId1, 3),
+                mkTopicAssignment(topicId2, 6))))
+            .withCurrentPartitionEpoch((topicId, partitionId) -> 11)
+            .withOwnedTopicPartitions(Collections.emptyList())
             .build();
 
-        assertEquals(ConsumerGroupMember.MemberState.REVOKING, updatedMember.state());
-        assertEquals(10, updatedMember.previousMemberEpoch());
-        assertEquals(10, updatedMember.memberEpoch());
-        assertEquals(12, updatedMember.targetMemberEpoch());
-        assertEquals(Collections.emptyMap(), updatedMember.assignedPartitions());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 1, 2, 3),
-            mkTopicAssignment(topicId2, 4, 5, 6)
-        ), updatedMember.partitionsPendingRevocation());
-        assertEquals(mkAssignment(
-            mkTopicAssignment(topicId1, 6, 7, 8),
-            mkTopicAssignment(topicId2, 9, 10, 11)
-        ), updatedMember.partitionsPendingAssignment());
-    }
-
-    private static List<ConsumerGroupHeartbeatRequestData.TopicPartitions> requestFromAssignment(
-        Map<Uuid, Set<Integer>> assignment
-    ) {
-        List<ConsumerGroupHeartbeatRequestData.TopicPartitions> topicPartitions = new ArrayList<>();
-
-        assignment.forEach((topicId, partitions) -> {
-            ConsumerGroupHeartbeatRequestData.TopicPartitions topic = new ConsumerGroupHeartbeatRequestData.TopicPartitions()
-                .setTopicId(topicId)
-                .setPartitions(new ArrayList<>(partitions));
-            topicPartitions.add(topic);
-        });
-
-        return topicPartitions;
+        assertEquals(
+            new ConsumerGroupMember.Builder("member")
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(12)
+                .setPreviousMemberEpoch(11)
+                .setAssignedPartitions(mkAssignment(
+                    mkTopicAssignment(topicId1, 3),
+                    mkTopicAssignment(topicId2, 6)))
+                .build(),
+            updatedMember
+        );
     }
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/metrics/GroupCoordinatorMetricsShardTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/metrics/GroupCoordinatorMetricsShardTest.java
index b16ff6c6535ac..5696f2bca6951 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/metrics/GroupCoordinatorMetricsShardTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/metrics/GroupCoordinatorMetricsShardTest.java
@@ -214,7 +214,7 @@ public void testConsumerGroupStateTransitionMetrics() {
 
         // Set member2 to ASSIGNING state.
         new ConsumerGroupMember.Builder(member2)
-            .setPartitionsPendingAssignment(Collections.singletonMap(Uuid.ZERO_UUID, Collections.singleton(0)))
+            .setPartitionsPendingRevocation(Collections.singletonMap(Uuid.ZERO_UUID, Collections.singleton(0)))
             .build();
 
         snapshotRegistry.getOrCreateSnapshot(4000);

From 2c49a765735152f72d6a5120b6c6403fc03af43a Mon Sep 17 00:00:00 2001
From: Andras Katona <41361962+akatona84@users.noreply.github.com>
Date: Thu, 14 Mar 2024 10:36:46 +0100
Subject: [PATCH 151/258] KAFKA-13922: Adjustments for jacoco, coverage
 reporting (#11982)

Jacoco and scoverage reporting hasn't been working for a while. This commit fixes report generation. After this PR only subproject level reports are generated as Jenkins and Sonar only cares about that.
This PR doesn't change Kafka's Jenkinsfile.

Reviewers: Viktor Somogyi-Vass <viktorsomogyi@gmail.com>
---
 build.gradle               | 42 +++++++++-----------------------------
 gradle/dependencies.gradle |  2 +-
 2 files changed, 11 insertions(+), 33 deletions(-)

diff --git a/build.gradle b/build.gradle
index 786c731dccd8d..ffac25d7bd9ad 100644
--- a/build.gradle
+++ b/build.gradle
@@ -43,7 +43,7 @@ plugins {
   // includes spotbugs version 4.7.4, in which case CVE-2022-42920 can
   // be dropped from gradle/resources/dependencycheck-suppressions.xml
   id "com.github.spotbugs" version '5.1.3' apply false
-  id 'org.scoverage' version '7.0.1' apply false
+  id 'org.scoverage' version '8.0.3' apply false
   // Updating the shadow plugin version to 8.1.1 causes issue with signing and publishing the shadowed
   // artifacts - see https://github.com/johnrengelman/shadow/issues/901
   id 'com.github.johnrengelman.shadow' version '8.1.0' apply false
@@ -750,7 +750,7 @@ subprojects {
 
   if (userEnableTestCoverage) {
     def coverageGen = it.path == ':core' ? 'reportScoverage' : 'jacocoTestReport'
-    task reportCoverage(dependsOn: [coverageGen])
+    tasks.register('reportCoverage').configure { dependsOn(coverageGen) }
   }
 
   dependencyCheck {
@@ -800,33 +800,8 @@ def checkstyleConfigProperties(configFileName) {
   [importControlFile: "$configFileName"]
 }
 
-// Aggregates all jacoco results into the root project directory
 if (userEnableTestCoverage) {
-  task jacocoRootReport(type: org.gradle.testing.jacoco.tasks.JacocoReport) {
-    def javaProjects = subprojects.findAll { it.path != ':core' }
-
-    description = 'Generates an aggregate report from all subprojects'
-    dependsOn(javaProjects.test)
-
-    additionalSourceDirs.from = javaProjects.sourceSets.main.allSource.srcDirs
-    sourceDirectories.from = javaProjects.sourceSets.main.allSource.srcDirs
-    classDirectories.from = javaProjects.sourceSets.main.output
-    executionData.from = javaProjects.jacocoTestReport.executionData
-
-    reports {
-      html.required = true
-      xml.required = true
-    }
-    // workaround to ignore projects that don't have any tests at all
-    onlyIf = { true }
-    doFirst {
-      executionData = files(executionData.findAll { it.exists() })
-    }
-  }
-}
-
-if (userEnableTestCoverage) {
-  task reportCoverage(dependsOn: ['jacocoRootReport', 'core:reportCoverage'])
+  tasks.register('reportCoverage').configure { dependsOn(subprojects.reportCoverage) }
 }
 
 def connectPkgs = [
@@ -1003,6 +978,9 @@ project(':core') {
   if (userEnableTestCoverage) {
     scoverage {
       scoverageVersion = versions.scoverage
+      if (versions.baseScala == '2.13') {
+        scoverageScalaVersion = '2.13.9' // there's no newer 2.13 artifact, org.scoverage:scalac-scoverage-plugin_2.13.9:2.0.11 is the latest as of now
+      }
       reportDir = file("${rootProject.buildDir}/scoverage")
       highlighting = false
       minimumRate = 0.0
@@ -1294,7 +1272,7 @@ project(':metadata') {
     }
     test {
       java {
-        srcDirs = ["src/generated/java", "src/test/java"]
+        srcDirs = ["src/test/java"]
       }
     }
   }
@@ -1661,7 +1639,7 @@ project(':raft') {
     }
     test {
       java {
-        srcDirs = ["src/generated/java", "src/test/java"]
+        srcDirs = ["src/test/java"]
       }
     }
   }
@@ -1904,7 +1882,7 @@ project(':storage') {
     }
     test {
       java {
-        srcDirs = ["src/generated/java", "src/test/java"]
+        srcDirs = ["src/test/java"]
       }
     }
   }
@@ -2218,7 +2196,7 @@ project(':streams') {
     }
     test {
       java {
-        srcDirs = ["src/generated/java", "src/test/java"]
+        srcDirs = ["src/test/java"]
       }
     }
   }
diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle
index 4f31bcd25c1d4..92cd8e1b943e2 100644
--- a/gradle/dependencies.gradle
+++ b/gradle/dependencies.gradle
@@ -159,7 +159,7 @@ versions += [
   // https://github.com/scalameta/scalafmt/releases/tag/v3.1.0.
   scalafmt: "3.7.14",
   scalaJava8Compat : "1.0.2",
-  scoverage: "1.9.3",
+  scoverage: "2.0.11",
   slf4j: "1.7.36",
   snappy: "1.1.10.5",
   spotbugs: "4.8.0",

From 37212bb242f937c0f0f8936350777686528291a4 Mon Sep 17 00:00:00 2001
From: David Mao <47232755+splett2@users.noreply.github.com>
Date: Thu, 14 Mar 2024 10:53:26 -0700
Subject: [PATCH 152/258] MINOR: AddPartitionsToTxnManager performance
 optimizations (#15454)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Justine Olshan <jolshan@confluent.io>
---
 .../server/AddPartitionsToTxnManager.scala    |  41 ++----
 .../AddPartitionsToTxnManagerTest.scala       | 137 ++++--------------
 2 files changed, 38 insertions(+), 140 deletions(-)

diff --git a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
index 59f75d4b11faf..8604c97ecc303 100644
--- a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
+++ b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
@@ -32,7 +32,7 @@ import org.apache.kafka.server.util.{InterBrokerSendThread, RequestAndCompletion
 
 import java.util
 import java.util.concurrent.TimeUnit
-import scala.collection.{Set, Seq, mutable}
+import scala.collection.{Seq, mutable}
 import scala.jdk.CollectionConverters._
 
 object AddPartitionsToTxnManager {
@@ -66,6 +66,7 @@ class AddPartitionsToTxnManager(
 
   this.logIdent = logPrefix
 
+  private val interBrokerListenerName = config.interBrokerListenerName
   private val inflightNodes = mutable.HashSet[Node]()
   private val nodesToTransactions = mutable.Map[Node, TransactionDataAndCallbacks]()
 
@@ -80,10 +81,9 @@ class AddPartitionsToTxnManager(
     topicPartitions: Seq[TopicPartition],
     callback: AddPartitionsToTxnManager.AppendCallback
   ): Unit = {
-    val (error, node) = getTransactionCoordinator(partitionFor(transactionalId))
-
-    if (error != Errors.NONE) {
-      callback(topicPartitions.map(tp => tp -> error).toMap)
+    val coordinatorNode = getTransactionCoordinator(partitionFor(transactionalId))
+    if (coordinatorNode.isEmpty) {
+      callback(topicPartitions.map(tp => tp -> Errors.COORDINATOR_NOT_AVAILABLE).toMap)
     } else {
       val topicCollection = new AddPartitionsToTxnTopicCollection()
       topicPartitions.groupBy(_.topic).forKeyValue { (topic, tps) =>
@@ -99,7 +99,7 @@ class AddPartitionsToTxnManager(
         .setVerifyOnly(true)
         .setTopics(topicCollection)
 
-      addTxnData(node, transactionData, callback)
+      addTxnData(coordinatorNode.get, transactionData, callback)
     }
   }
 
@@ -146,31 +146,10 @@ class AddPartitionsToTxnManager(
     }
   }
 
-  private def getTransactionCoordinator(partition: Int): (Errors, Node) = {
-    val listenerName = config.interBrokerListenerName
-
-    val topicMetadata = metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), listenerName)
-
-    if (topicMetadata.headOption.isEmpty) {
-      // If topic is not created, then the transaction is definitely not started.
-      (Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
-    } else {
-      if (topicMetadata.head.errorCode != Errors.NONE.code) {
-        (Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
-      } else {
-        val coordinatorEndpoint = topicMetadata.head.partitions.asScala
-          .find(_.partitionIndex == partition)
-          .filter(_.leaderId != MetadataResponse.NO_LEADER_ID)
-          .flatMap(metadata => metadataCache.getAliveBrokerNode(metadata.leaderId, listenerName))
-
-        coordinatorEndpoint match {
-          case Some(endpoint) =>
-            (Errors.NONE, endpoint)
-          case _ =>
-            (Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
-        }
-      }
-    }
+  private def getTransactionCoordinator(partition: Int): Option[Node] = {
+   metadataCache.getPartitionInfo(Topic.TRANSACTION_STATE_TOPIC_NAME, partition)
+      .filter(_.leader != MetadataResponse.NO_LEADER_ID)
+      .flatMap(metadata => metadataCache.getAliveBrokerNode(metadata.leader, interBrokerListenerName))
   }
 
   private def topicPartitionsToError(transactionData: AddPartitionsToTxnTransaction, error: Errors): Map[TopicPartition, Errors] = {
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
index 9e34322ec9614..e849af43bd474 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
@@ -24,11 +24,12 @@ import org.apache.kafka.clients.{ClientResponse, NetworkClient}
 import org.apache.kafka.common.errors.{AuthenticationException, SaslAuthenticationException, UnsupportedVersionException}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.{AddPartitionsToTxnTopic, AddPartitionsToTxnTopicCollection, AddPartitionsToTxnTransaction, AddPartitionsToTxnTransactionCollection}
-import org.apache.kafka.common.message.{AddPartitionsToTxnResponseData, MetadataResponseData}
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData
 import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnResultCollection
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState
 import org.apache.kafka.common.{Node, TopicPartition}
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.{AbstractResponse, AddPartitionsToTxnRequest, AddPartitionsToTxnResponse}
+import org.apache.kafka.common.requests.{AbstractResponse, AddPartitionsToTxnRequest, AddPartitionsToTxnResponse, MetadataResponse}
 import org.apache.kafka.common.utils.MockTime
 import org.apache.kafka.server.metrics.KafkaMetricsGroup
 import org.apache.kafka.server.util.RequestAndCompletionHandler
@@ -98,23 +99,8 @@ class AddPartitionsToTxnManagerTest {
     when(partitionFor.apply(transactionalId1)).thenReturn(0)
     when(partitionFor.apply(transactionalId2)).thenReturn(1)
     when(partitionFor.apply(transactionalId3)).thenReturn(0)
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(0),
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(1)
-              .setLeaderId(1)
-          ).asJava)
-      ))
-    when(metadataCache.getAliveBrokerNode(0, config.interBrokerListenerName))
-      .thenReturn(Some(node0))
-    when(metadataCache.getAliveBrokerNode(1, config.interBrokerListenerName))
-      .thenReturn(Some(node1))
+    mockTransactionStateMetadata(0, 0, Some(node0))
+    mockTransactionStateMetadata(1, 1, Some(node1))
 
     val transaction1Errors = mutable.Map[TopicPartition, Errors]()
     val transaction2Errors = mutable.Map[TopicPartition, Errors]()
@@ -167,28 +153,9 @@ class AddPartitionsToTxnManagerTest {
     when(partitionFor.apply(transactionalId1)).thenReturn(0)
     when(partitionFor.apply(transactionalId2)).thenReturn(1)
     when(partitionFor.apply(transactionalId3)).thenReturn(2)
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(0),
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(1)
-              .setLeaderId(1),
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(2)
-              .setLeaderId(2)
-          ).asJava)
-      ))
-    when(metadataCache.getAliveBrokerNode(0, config.interBrokerListenerName))
-      .thenReturn(Some(node0))
-    when(metadataCache.getAliveBrokerNode(1, config.interBrokerListenerName))
-      .thenReturn(Some(node1))
-    when(metadataCache.getAliveBrokerNode(2, config.interBrokerListenerName))
-      .thenReturn(Some(node2))
+    mockTransactionStateMetadata(0, 0, Some(node0))
+    mockTransactionStateMetadata(1, 1, Some(node1))
+    mockTransactionStateMetadata(2, 2, Some(node2))
 
     val transactionErrors = mutable.Map[TopicPartition, Errors]()
 
@@ -247,51 +214,16 @@ class AddPartitionsToTxnManagerTest {
     }
 
     // The transaction state topic does not exist.
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq())
-    checkError()
-
-    // The metadata of the transaction state topic returns an error.
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setErrorCode(Errors.BROKER_NOT_AVAILABLE.code)
-      ))
-    checkError()
-
-    // The partition does not exist.
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-      ))
+    when(metadataCache.getPartitionInfo(Topic.TRANSACTION_STATE_TOPIC_NAME, 0))
+      .thenReturn(Option.empty)
     checkError()
 
     // The partition has no leader.
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(-1)
-          ).asJava)
-      ))
+    mockTransactionStateMetadata(0, -1, Option.empty)
     checkError()
 
     // The leader is not available.
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(0)
-          ).asJava)
-      ))
+    mockTransactionStateMetadata(0, 0, Option.empty)
     checkError()
   }
 
@@ -299,18 +231,7 @@ class AddPartitionsToTxnManagerTest {
   def testAddPartitionsToTxnHandlerErrorHandling(): Unit = {
     when(partitionFor.apply(transactionalId1)).thenReturn(0)
     when(partitionFor.apply(transactionalId2)).thenReturn(0)
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(0)
-          ).asJava)
-      ))
-    when(metadataCache.getAliveBrokerNode(0, config.interBrokerListenerName))
-      .thenReturn(Some(node0))
+    mockTransactionStateMetadata(0, 0, Some(node0))
 
     val transaction1Errors = mutable.Map[TopicPartition, Errors]()
     val transaction2Errors = mutable.Map[TopicPartition, Errors]()
@@ -380,23 +301,8 @@ class AddPartitionsToTxnManagerTest {
 
     when(partitionFor.apply(transactionalId1)).thenReturn(0)
     when(partitionFor.apply(transactionalId2)).thenReturn(1)
-    when(metadataCache.getTopicMetadata(Set(Topic.TRANSACTION_STATE_TOPIC_NAME), config.interBrokerListenerName))
-      .thenReturn(Seq(
-        new MetadataResponseData.MetadataResponseTopic()
-          .setName(Topic.TRANSACTION_STATE_TOPIC_NAME)
-          .setPartitions(List(
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(0)
-              .setLeaderId(0),
-            new MetadataResponseData.MetadataResponsePartition()
-              .setPartitionIndex(1)
-              .setLeaderId(1)
-          ).asJava)
-      ))
-    when(metadataCache.getAliveBrokerNode(0, config.interBrokerListenerName))
-      .thenReturn(Some(node0))
-    when(metadataCache.getAliveBrokerNode(1, config.interBrokerListenerName))
-      .thenReturn(Some(node1))
+    mockTransactionStateMetadata(0, 0, Some(node0))
+    mockTransactionStateMetadata(1, 1, Some(node1))
 
     // Update max verification time when we see a higher verification time.
     when(mockVerificationTime.update(anyLong())).thenAnswer { invocation =>
@@ -458,6 +364,19 @@ class AddPartitionsToTxnManagerTest {
     }
   }
 
+  private def mockTransactionStateMetadata(partitionIndex: Int, leaderId: Int, leaderNode: Option[Node]): Unit = {
+    when(metadataCache.getPartitionInfo(Topic.TRANSACTION_STATE_TOPIC_NAME, partitionIndex))
+      .thenReturn(Some(
+        new UpdateMetadataPartitionState()
+          .setTopicName(Topic.TRANSACTION_STATE_TOPIC_NAME)
+          .setPartitionIndex(partitionIndex)
+          .setLeader(leaderId)))
+    if (leaderId != MetadataResponse.NO_LEADER_ID) {
+      when(metadataCache.getAliveBrokerNode(leaderId, config.interBrokerListenerName))
+        .thenReturn(leaderNode)
+    }
+  }
+
   private def clientResponse(
     response: AbstractResponse,
     authException: AuthenticationException = null,

From 178761eb36a3d0e82526a9694378fc2c9bdc7dd6 Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Thu, 14 Mar 2024 15:50:57 -0400
Subject: [PATCH 153/258] KAFKA-14683 Cleanup WorkerSinkTaskTest (#15506)

1) Rename WorkerSinkTaskMockitoTest back to WorkerSinkTaskTest
2) Tidy up the code a bit
3) rewrite "fail" by "assertThrow"

Reviewers: Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle                                  |  3 +-
 checkstyle/suppressions.xml                   |  2 +-
 ...ckitoTest.java => WorkerSinkTaskTest.java} | 77 ++++++-------------
 3 files changed, 26 insertions(+), 56 deletions(-)
 rename connect/runtime/src/test/java/org/apache/kafka/connect/runtime/{WorkerSinkTaskMockitoTest.java => WorkerSinkTaskTest.java} (97%)

diff --git a/build.gradle b/build.gradle
index ffac25d7bd9ad..fb5cbf5869d6d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -420,8 +420,7 @@ subprojects {
   if (JavaVersion.current().isCompatibleWith(JavaVersion.VERSION_16)) {
     testsToExclude.addAll([
       // connect tests
-      "**/KafkaConfigBackingStoreTest.*",
-      "**/WorkerSinkTaskTest.*"
+      "**/KafkaConfigBackingStoreTest.*"
     ])
   }
 
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 27c252f83906c..5ab49f7f64c3f 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -96,7 +96,7 @@
               files="(AbstractFetch|ClientTelemetryReporter|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/>
 
     <suppress checks="JavaNCSS"
-              files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|WorkerSinkTaskMockitoTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest|KafkaRaftClientTest).java"/>
+              files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest|KafkaRaftClientTest).java"/>
 
     <suppress checks="NPathComplexity"
               files="(ConsumerCoordinator|BufferPool|MetricName|Node|ConfigDef|RecordBatch|SslFactory|SslTransportLayer|MetadataResponse|KerberosLogin|Selector|Sender|Serdes|TokenInformation|Agent|Values|PluginUtils|MiniTrogdorCluster|TasksRequest|KafkaProducer|AbstractStickyAssignor|KafkaRaftClient|Authorizer|FetchSessionHandler|RecordAccumulator).java"/>
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
similarity index 97%
rename from connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
rename to connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
index 83bf4fe0c7b0a..21b2b10c16e75 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskMockitoTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
@@ -22,9 +22,8 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyMap;
@@ -111,7 +110,7 @@
 import org.mockito.stubbing.OngoingStubbing;
 
 @RunWith(MockitoJUnitRunner.StrictStubs.class)
-public class WorkerSinkTaskMockitoTest {
+public class WorkerSinkTaskTest {
     // These are fixed to keep this code simpler. In this example we assume byte[] raw values
     // with mix of integer/string in Connect
     private static final String TOPIC = "test";
@@ -316,9 +315,7 @@ public void testPause() {
         // And unpause
         verify(statusListener).onResume(taskId);
         verify(consumer, times(2)).wakeup();
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            verify(consumer).resume(singleton(tp));
-        });
+        INITIAL_ASSIGNMENT.forEach(tp -> verify(consumer).resume(singleton(tp)));
         verify(sinkTask, times(4)).put(anyList());
     }
 
@@ -419,9 +416,7 @@ public void testPollRedelivery() {
         time.sleep(30000L);
 
         verify(sinkTask, times(3)).put(anyList());
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            verify(consumer).resume(Collections.singleton(tp));
-        });
+        INITIAL_ASSIGNMENT.forEach(tp -> verify(consumer).resume(Collections.singleton(tp)));
 
         assertSinkMetricValue("sink-record-read-total", 1.0);
         assertSinkMetricValue("sink-record-send-total", 1.0);
@@ -528,9 +523,7 @@ public void testPollRedeliveryWithConsumerRebalance() {
         verify(sinkTask).close(INITIAL_ASSIGNMENT);
 
         // All partitions are resumed, as all previously paused-for-redelivery partitions were revoked
-        newAssignment.forEach(tp -> {
-            verify(consumer).resume(Collections.singleton(tp));
-        });
+        newAssignment.forEach(tp -> verify(consumer).resume(Collections.singleton(tp)));
     }
 
     @Test
@@ -553,12 +546,8 @@ public void testErrorInRebalancePartitionLoss() {
         workerTask.iteration();
         verifyPollInitialAssignment();
 
-        try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
-        }
+        RuntimeException thrownException = assertThrows(RuntimeException.class, () -> workerTask.iteration());
+        assertEquals(exception, thrownException);
     }
 
     @Test
@@ -580,12 +569,9 @@ public void testErrorInRebalancePartitionRevocation() {
 
         workerTask.iteration();
         verifyPollInitialAssignment();
-        try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
-        }
+
+        RuntimeException thrownException = assertThrows(RuntimeException.class, () -> workerTask.iteration());
+        assertEquals(exception, thrownException);
     }
 
     @Test
@@ -608,11 +594,10 @@ public void testErrorInRebalancePartitionAssignment() {
         verifyPollInitialAssignment();
 
         expectRebalanceAssignmentError(exception);
+
         try {
-            workerTask.iteration();
-            fail("Poll should have raised the rebalance exception");
-        } catch (RuntimeException e) {
-            assertEquals(exception, e);
+            RuntimeException thrownException = assertThrows(RuntimeException.class, () -> workerTask.iteration());
+            assertEquals(exception, thrownException);
         } finally {
             verify(sinkTask).close(INITIAL_ASSIGNMENT);
         }
@@ -798,7 +783,7 @@ public void testWakeupInCommitSyncCausesRetry() {
         doThrow(new WakeupException())
                 // and succeed the second time
                 .doNothing()
-                .when(consumer).commitSync(eq(offsets));
+                .when(consumer).commitSync(offsets);
 
         workerTask.iteration(); // first record delivered
 
@@ -808,9 +793,7 @@ public void testWakeupInCommitSyncCausesRetry() {
         verify(sinkTask).close(INITIAL_ASSIGNMENT);
         verify(sinkTask, times(2)).open(INITIAL_ASSIGNMENT);
 
-        INITIAL_ASSIGNMENT.forEach(tp -> {
-            verify(consumer).resume(Collections.singleton(tp));
-        });
+        INITIAL_ASSIGNMENT.forEach(tp -> verify(consumer).resume(Collections.singleton(tp)));
 
         verify(statusListener).onResume(taskId);
 
@@ -862,7 +845,7 @@ public void testWakeupNotThrownDuringShutdown() {
         doThrow(new WakeupException())
                 // and succeed the second time
                 .doNothing()
-                .when(consumer).commitSync(eq(offsets));
+                .when(consumer).commitSync(offsets);
 
         workerTask.execute();
 
@@ -1091,7 +1074,7 @@ public void testIgnoredCommit() {
     // Test that the commitTimeoutMs timestamp is correctly computed and checked in WorkerSinkTask.iteration()
     // when there is a long running commit in process. See KAFKA-4942 for more information.
     @Test
-    public void testLongRunningCommitWithoutTimeout() throws InterruptedException {
+    public void testLongRunningCommitWithoutTimeout() {
         createTask(initialState);
 
         workerTask.initialize(TASK_CONFIG);
@@ -1172,12 +1155,8 @@ public void testSinkTasksHandleCloseErrors() {
         // Throw another exception while closing the task's assignment
         doThrow(closeException).when(sinkTask).close(any(Collection.class));
 
-        try {
-            workerTask.execute();
-            fail("workerTask.execute should have thrown an exception");
-        } catch (RuntimeException e) {
-            assertSame("Exception from close should propagate as-is", closeException, e);
-        }
+        RuntimeException thrownException = assertThrows(RuntimeException.class, () -> workerTask.execute());
+        assertEquals(closeException, thrownException);
 
         verify(consumer).wakeup();
     }
@@ -1215,17 +1194,13 @@ public void testSuppressCloseErrors() {
 
         workerTask.initialize(TASK_CONFIG);
         workerTask.initializeAndStart();
-        try {
-            workerTask.execute();
-            fail("workerTask.execute should have thrown an exception");
-        } catch (ConnectException e) {
-            assertSame("Exception from put should be the cause", putException, e.getCause());
-            assertTrue("Exception from close should be suppressed", e.getSuppressed().length > 0);
-            assertSame(closeException, e.getSuppressed()[0]);
-        }
+
+        RuntimeException thrownException = assertThrows(ConnectException.class, () -> workerTask.execute());
+        assertEquals("Exception from put should be the cause", putException, thrownException.getCause());
+        assertTrue("Exception from close should be suppressed", thrownException.getSuppressed().length > 0);
+        assertEquals(closeException, thrownException.getSuppressed()[0]);
     }
 
-    @SuppressWarnings("unchecked")
     @Test
     public void testTaskCancelPreventsFinalOffsetCommit() {
         createTask(initialState);
@@ -1287,10 +1262,6 @@ public void testCommitWithOutOfOrderCallback() {
         expectTaskGetTopic();
         expectConversionAndTransformation(null, new RecordHeaders());
 
-        final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
-        workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
-        workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
-
         final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
         workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
         workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));

From 834efa66062687f520a9e936a5ee2798310d8bcf Mon Sep 17 00:00:00 2001
From: Luke Chen <showuon@gmail.com>
Date: Fri, 15 Mar 2024 06:09:45 +0800
Subject: [PATCH 154/258] KAFKA-16342 fix getOffsetByMaxTimestamp for
 compressed records (#15474)

Fix getOffsetByMaxTimestamp for compressed records.

This PR adds:

1) For inPlaceAssignment case, compute the correct offset for maxTimestamp when traversing the batch records, and set to ValidationResult in the end, instead of setting to last offset always.

2) For not inPlaceAssignment, set the offsetOfMaxTimestamp for the log create time, like non-compressed, and inPlaceAssignment cases, instead of setting to last offset always.

3) Add tests to verify the fix.

Reviewers: Jun Rao <junrao@apache.org>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../common/record/MemoryRecordsBuilder.java   |  29 +--
 .../record/MemoryRecordsBuilderTest.java      |  23 +-
 .../common/record/MemoryRecordsTest.java      |   5 +-
 .../admin/ListOffsetsIntegrationTest.scala    | 202 +++++++++++++++---
 .../kafka/server/QuorumTestHarness.scala      |   2 +-
 .../unit/kafka/log/LogValidatorTest.scala     |  29 +--
 .../storage/internals/log/LogValidator.java   |  13 +-
 .../kafka/tools/GetOffsetShellTest.java       |  11 +-
 8 files changed, 233 insertions(+), 81 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
index 143a515bcd00a..3663c5ea7e339 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
@@ -242,34 +242,23 @@ public MemoryRecords build() {
 
     /**
      * Get the max timestamp and its offset. The details of the offset returned are a bit subtle.
+     * Note: The semantic for the offset of max timestamp is the first offset with the max timestamp if there are multi-records having same timestamp.
      *
-     * If the log append time is used, the offset will be the last offset unless no compression is used and
-     * the message format version is 0 or 1, in which case, it will be the first offset.
+     * If the log append time is used, the offset will be the first offset of the record.
      *
-     * If create time is used, the offset will be the last offset unless no compression is used and the message
-     * format version is 0 or 1, in which case, it will be the offset of the record with the max timestamp.
+     * If create time is used, the offset will always be the offset of the record with the max timestamp.
+     *
+     * If it's NO_TIMESTAMP (i.e. MAGIC_VALUE_V0), we'll return offset -1 since no timestamp info in records.
      *
      * @return The max timestamp and its offset
      */
     public RecordsInfo info() {
         if (timestampType == TimestampType.LOG_APPEND_TIME) {
-            long shallowOffsetOfMaxTimestamp;
-            // Use the last offset when dealing with record batches
-            if (compressionType != CompressionType.NONE || magic >= RecordBatch.MAGIC_VALUE_V2)
-                shallowOffsetOfMaxTimestamp = lastOffset;
-            else
-                shallowOffsetOfMaxTimestamp = baseOffset;
-            return new RecordsInfo(logAppendTime, shallowOffsetOfMaxTimestamp);
-        } else if (maxTimestamp == RecordBatch.NO_TIMESTAMP) {
-            return new RecordsInfo(RecordBatch.NO_TIMESTAMP, lastOffset);
+            return new RecordsInfo(logAppendTime, baseOffset);
         } else {
-            long shallowOffsetOfMaxTimestamp;
-            // Use the last offset when dealing with record batches
-            if (compressionType != CompressionType.NONE || magic >= RecordBatch.MAGIC_VALUE_V2)
-                shallowOffsetOfMaxTimestamp = lastOffset;
-            else
-                shallowOffsetOfMaxTimestamp = offsetOfMaxTimestamp;
-            return new RecordsInfo(maxTimestamp, shallowOffsetOfMaxTimestamp);
+            // For create time, we always use offsetOfMaxTimestamp for the correct time -> offset mapping
+            // If it's MAGIC_VALUE_V0, the value will be the default value: [-1, -1]
+            return new RecordsInfo(maxTimestamp, offsetOfMaxTimestamp);
         }
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
index 33a419e64ee9d..81aa162feee1f 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
@@ -377,11 +377,8 @@ public void buildUsingLogAppendTime(Args args) {
 
         MemoryRecordsBuilder.RecordsInfo info = builder.info();
         assertEquals(logAppendTime, info.maxTimestamp);
-
-        if (args.compressionType == CompressionType.NONE && magic <= MAGIC_VALUE_V1)
-            assertEquals(0L, info.shallowOffsetOfMaxTimestamp);
-        else
-            assertEquals(2L, info.shallowOffsetOfMaxTimestamp);
+        // When logAppendTime is used, the first offset of the batch will be the offset of maxTimestamp
+        assertEquals(0L, info.shallowOffsetOfMaxTimestamp);
 
         for (RecordBatch batch : records.batches()) {
             if (magic == MAGIC_VALUE_V0) {
@@ -415,10 +412,11 @@ public void buildUsingCreateTime(Args args) {
             assertEquals(2L, info.maxTimestamp);
         }
 
-        if (args.compressionType == CompressionType.NONE && magic == MAGIC_VALUE_V1)
-            assertEquals(1L, info.shallowOffsetOfMaxTimestamp);
+        if (magic == MAGIC_VALUE_V0)
+            // in MAGIC_VALUE_V0's case, we don't have timestamp info in records, so always return -1.
+            assertEquals(-1L, info.shallowOffsetOfMaxTimestamp);
         else
-            assertEquals(2L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(1L, info.shallowOffsetOfMaxTimestamp);
 
         int i = 0;
         long[] expectedTimestamps = new long[] {0L, 2L, 1L};
@@ -495,12 +493,13 @@ public void writePastLimit(Args args) {
         MemoryRecords records = builder.build();
 
         MemoryRecordsBuilder.RecordsInfo info = builder.info();
-        if (magic == MAGIC_VALUE_V0)
+        if (magic == MAGIC_VALUE_V0) {
             assertEquals(-1, info.maxTimestamp);
-        else
+            assertEquals(-1L, info.shallowOffsetOfMaxTimestamp);
+        } else {
             assertEquals(2L, info.maxTimestamp);
-
-        assertEquals(2L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(2L, info.shallowOffsetOfMaxTimestamp);
+        }
 
         long i = 0L;
         for (RecordBatch batch : records.batches()) {
diff --git a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
index 3f0195bf5d149..50821af841c42 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
@@ -893,10 +893,7 @@ public void testFilterTo(Args args) {
         assertEquals(filtered.limit(), result.bytesRetained());
         if (magic > RecordBatch.MAGIC_VALUE_V0) {
             assertEquals(20L, result.maxTimestamp());
-            if (compression == CompressionType.NONE && magic < RecordBatch.MAGIC_VALUE_V2)
-                assertEquals(4L, result.shallowOffsetOfMaxTimestamp());
-            else
-                assertEquals(5L, result.shallowOffsetOfMaxTimestamp());
+            assertEquals(4L, result.shallowOffsetOfMaxTimestamp());
         }
 
         MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
diff --git a/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala b/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
index 333bd980a603b..e5e22e9dff97d 100644
--- a/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
@@ -19,82 +19,236 @@ package kafka.admin
 
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
+import kafka.utils.TestUtils.{createProducer, plaintextBootstrapServers}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin._
 import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.config.TopicConfig
+import org.apache.kafka.common.utils.{MockTime, Time, Utils}
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo}
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 
+import java.util.Properties
 import scala.collection.{Map, Seq}
 import scala.jdk.CollectionConverters._
 
 class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
 
   val topicName = "foo"
+  val topicNameWithCustomConfigs = "foo2"
   var adminClient: Admin = _
+  var setOldMessageFormat: Boolean = false
+  val mockTime: Time = new MockTime(1)
 
   @BeforeEach
   override def setUp(testInfo: TestInfo): Unit = {
     super.setUp(testInfo)
-    createTopic(topicName, 1, 1.toShort)
-    produceMessages()
+    createTopicWithConfig(topicName, new Properties())
     adminClient = Admin.create(Map[String, Object](
       AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG -> bootstrapServers()
     ).asJava)
   }
 
+  override def brokerTime(brokerId: Int): Time = mockTime
+
   @AfterEach
   override def tearDown(): Unit = {
+    setOldMessageFormat = false
     Utils.closeQuietly(adminClient, "ListOffsetsAdminClient")
     super.tearDown()
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
-  def testEarliestOffset(quorum: String): Unit = {
-    val earliestOffset = runFetchOffsets(adminClient, OffsetSpec.earliest())
-    assertEquals(0, earliestOffset.offset())
+  def testThreeCompressedRecordsInOneBatch(quorum: String): Unit = {
+    produceMessagesInOneBatch("gzip")
+    verifyListOffsets()
+
+    // test LogAppendTime case
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInOneBatch("gzip", topicNameWithCustomConfigs)
+    // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
+    // So in this one batch test, it'll be the first offset 0
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 0)
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
-  def testLatestOffset(quorum: String): Unit = {
-    val latestOffset = runFetchOffsets(adminClient, OffsetSpec.latest())
-    assertEquals(3, latestOffset.offset())
+  def testThreeRecordsInSeparateBatch(quorum: String): Unit = {
+    produceMessagesInSeparateBatch()
+    verifyListOffsets()
+  }
+
+  // The message conversion test only run in ZK mode because KRaft mode doesn't support "inter.broker.protocol.version" < 3.0
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk"))
+  def testThreeRecordsInOneBatchWithMessageConversion(quorum: String): Unit = {
+    createOldMessageFormatBrokers()
+    produceMessagesInOneBatch()
+    verifyListOffsets()
+
+    // test LogAppendTime case
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInOneBatch(topic = topicNameWithCustomConfigs)
+    // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
+    // So in this one batch test, it'll be the first offset 0
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 0)
+  }
+
+  // The message conversion test only run in ZK mode because KRaft mode doesn't support "inter.broker.protocol.version" < 3.0
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk"))
+  def testThreeRecordsInSeparateBatchWithMessageConversion(quorum: String): Unit = {
+    createOldMessageFormatBrokers()
+    produceMessagesInSeparateBatch()
+    verifyListOffsets()
+
+    // test LogAppendTime case
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInSeparateBatch(topic = topicNameWithCustomConfigs)
+    // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
+    // So in this separate batch test, it'll be the last offset 2
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 2)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def testThreeRecordsInOneBatchHavingDifferentCompressionTypeWithServer(quorum: String): Unit = {
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.COMPRESSION_TYPE_CONFIG, "lz4")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInOneBatch(topic = topicNameWithCustomConfigs)
+    verifyListOffsets(topic = topicNameWithCustomConfigs)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def testThreeRecordsInSeparateBatchHavingDifferentCompressionTypeWithServer(quorum: String): Unit = {
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.COMPRESSION_TYPE_CONFIG, "lz4")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInSeparateBatch(topic = topicNameWithCustomConfigs)
+    verifyListOffsets(topic = topicNameWithCustomConfigs)
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
-  def testMaxTimestampOffset(quorum: String): Unit = {
-    val maxTimestampOffset = runFetchOffsets(adminClient, OffsetSpec.maxTimestamp())
-    assertEquals(1, maxTimestampOffset.offset())
+  def testThreeCompressedRecordsInSeparateBatch(quorum: String): Unit = {
+    produceMessagesInSeparateBatch("gzip")
+    verifyListOffsets()
+
+    // test LogAppendTime case
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    produceMessagesInSeparateBatch("gzip", topicNameWithCustomConfigs)
+    // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
+    // So in this separate batch test, it'll be the last offset 2
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 2)
+  }
+
+  private def createOldMessageFormatBrokers(): Unit = {
+    setOldMessageFormat = true
+    recreateBrokers(reconfigure = true, startup = true)
+    Utils.closeQuietly(adminClient, "ListOffsetsAdminClient")
+    adminClient = Admin.create(Map[String, Object](
+      AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG -> bootstrapServers()
+    ).asJava)
+  }
+
+  private def createTopicWithConfig(topic: String, props: Properties): Unit = {
+    createTopic(topic, 1, 1.toShort, topicConfig = props)
+  }
+
+  private def verifyListOffsets(topic: String = topicName, expectedMaxTimestampOffset: Int = 1): Unit = {
+    val earliestOffset = runFetchOffsets(adminClient, OffsetSpec.earliest(), topic)
+    assertEquals(0, earliestOffset.offset())
+
+    val latestOffset = runFetchOffsets(adminClient, OffsetSpec.latest(), topic)
+    assertEquals(3, latestOffset.offset())
+
+    val maxTimestampOffset = runFetchOffsets(adminClient, OffsetSpec.maxTimestamp(), topic)
+    assertEquals(expectedMaxTimestampOffset, maxTimestampOffset.offset())
   }
 
   private def runFetchOffsets(adminClient: Admin,
-                              offsetSpec: OffsetSpec): ListOffsetsResult.ListOffsetsResultInfo = {
-    val tp = new TopicPartition(topicName, 0)
+                              offsetSpec: OffsetSpec,
+                              topic: String): ListOffsetsResult.ListOffsetsResultInfo = {
+    val tp = new TopicPartition(topic, 0)
     adminClient.listOffsets(Map(
       tp -> offsetSpec
     ).asJava, new ListOffsetsOptions()).all().get().get(tp)
   }
 
-  def produceMessages(): Unit = {
+  private def produceMessagesInOneBatch(compressionType: String = "none", topic: String = topicName): Unit = {
     val records = Seq(
-      new ProducerRecord[Array[Byte], Array[Byte]](topicName, 0, 100L,
-        null, new Array[Byte](10000)),
-      new ProducerRecord[Array[Byte], Array[Byte]](topicName, 0, 999L,
-        null, new Array[Byte](10000)),
-      new ProducerRecord[Array[Byte], Array[Byte]](topicName, 0, 200L,
-        null, new Array[Byte](10000)),
+      new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 100L,
+        null, new Array[Byte](10)),
+      new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 999L,
+        null, new Array[Byte](10)),
+      new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 200L,
+        null, new Array[Byte](10)),
     )
-    TestUtils.produceMessages(brokers, records, -1)
+    // create a producer with large linger.ms and enough batch.size (default is enough for three 10 bytes records),
+    // so that we can confirm all records will be accumulated in producer until we flush them into one batch.
+    val producer = createProducer(
+      plaintextBootstrapServers(brokers),
+      deliveryTimeoutMs = Int.MaxValue,
+      lingerMs = Int.MaxValue,
+      compressionType = compressionType)
+
+    try {
+      val futures = records.map(producer.send)
+      producer.flush()
+      futures.foreach(_.get)
+    } finally {
+      producer.close()
+    }
   }
 
-  def generateConfigs: Seq[KafkaConfig] =
-    TestUtils.createBrokerConfigs(1, zkConnectOrNull).map(KafkaConfig.fromProps)
+  private def produceMessagesInSeparateBatch(compressionType: String = "none", topic: String = topicName): Unit = {
+    val records = Seq(new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 100L,
+        null, new Array[Byte](10)))
+    val records2 = Seq(new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 999L,
+      null, new Array[Byte](10)))
+    val records3 = Seq(new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, 200L,
+      null, new Array[Byte](10)))
+
+    val producer = createProducer(
+      plaintextBootstrapServers(brokers),
+      compressionType = compressionType)
+    try {
+      val futures = records.map(producer.send)
+      futures.foreach(_.get)
+      // advance the server time after each record sent to make sure the time changed when appendTime is used
+      mockTime.sleep(100)
+      val futures2 = records2.map(producer.send)
+      futures2.foreach(_.get)
+      mockTime.sleep(100)
+      val futures3 = records3.map(producer.send)
+      futures3.foreach(_.get)
+    } finally {
+      producer.close()
+    }
+  }
+
+  def generateConfigs: Seq[KafkaConfig] = {
+    TestUtils.createBrokerConfigs(1, zkConnectOrNull).map{ props =>
+      if (setOldMessageFormat) {
+        props.setProperty("log.message.format.version", "0.10.0")
+        props.setProperty("inter.broker.protocol.version", "0.10.0")
+      }
+      props
+    }.map(KafkaConfig.fromProps)
+  }
 }
 
diff --git a/core/src/test/scala/integration/kafka/server/QuorumTestHarness.scala b/core/src/test/scala/integration/kafka/server/QuorumTestHarness.scala
index 45cd92057f0cf..aabf3f9345368 100755
--- a/core/src/test/scala/integration/kafka/server/QuorumTestHarness.scala
+++ b/core/src/test/scala/integration/kafka/server/QuorumTestHarness.scala
@@ -124,7 +124,7 @@ class KRaftQuorumImplementation(
       util.EnumSet.of(REQUIRE_AT_LEAST_ONE_VALID, REQUIRE_METADATA_LOG_DIR))
     val sharedServer = new SharedServer(config,
       metaPropertiesEnsemble,
-      Time.SYSTEM,
+      time,
       new Metrics(),
       controllerQuorumVotersFuture,
       faultHandlerFactory)
diff --git a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
index 0b5d17cbcc234..ac6152b9b153c 100644
--- a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
@@ -219,8 +219,8 @@ class LogValidatorTest {
       "MessageSet should still valid")
     assertEquals(now, validatedResults.maxTimestampMs,
       s"Max timestamp should be $now")
-    assertEquals(records.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestampMs,
-      s"The offset of max timestamp should be ${records.records.asScala.size - 1}")
+    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+      s"The offset of max timestamp should be 0 if logAppendTime is used")
     assertTrue(validatedResults.messageSizeMaybeChanged,
       "Message size may have been changed")
 
@@ -271,8 +271,8 @@ class LogValidatorTest {
       "MessageSet should still valid")
     assertEquals(now, validatedResults.maxTimestampMs,
       s"Max timestamp should be $now")
-    assertEquals(records.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestampMs,
-      s"The offset of max timestamp should be ${records.records.asScala.size - 1}")
+    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+      s"The offset of max timestamp should be 0 if logAppendTime is used")
     assertFalse(validatedResults.messageSizeMaybeChanged,
       "Message size should not have been changed")
 
@@ -341,6 +341,7 @@ class LogValidatorTest {
 
   private def checkNonCompressed(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
+    // set the timestamp of seq(1) (i.e. offset 1) as the max timestamp
     val timestampSeq = Seq(now - 1, now + 1, now)
 
     val (producerId, producerEpoch, baseSequence, isTransactional, partitionLeaderEpoch) =
@@ -431,6 +432,7 @@ class LogValidatorTest {
 
   private def checkRecompression(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
+    // set the timestamp of seq(1) (i.e. offset 1) as the max timestamp
     val timestampSeq = Seq(now - 1, now + 1, now)
 
     val (producerId, producerEpoch, baseSequence, isTransactional, partitionLeaderEpoch) =
@@ -484,8 +486,8 @@ class LogValidatorTest {
     }
     assertEquals(now + 1, validatingResults.maxTimestampMs,
       s"Max timestamp should be ${now + 1}")
-    assertEquals(2, validatingResults.shallowOffsetOfMaxTimestampMs,
-      "Offset of max timestamp should be 2")
+    assertEquals(1, validatingResults.shallowOffsetOfMaxTimestampMs,
+      "Offset of max timestamp should be 1")
     assertTrue(validatingResults.messageSizeMaybeChanged,
       "Message size should have been changed")
 
@@ -536,8 +538,8 @@ class LogValidatorTest {
     }
     assertEquals(validatedResults.maxTimestampMs, RecordBatch.NO_TIMESTAMP,
       s"Max timestamp should be ${RecordBatch.NO_TIMESTAMP}")
-    assertEquals(validatedRecords.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestampMs,
-      s"Offset of max timestamp should be ${validatedRecords.records.asScala.size - 1}")
+    assertEquals(-1, validatedResults.shallowOffsetOfMaxTimestampMs,
+      s"Offset of max timestamp should be -1")
     assertTrue(validatedResults.messageSizeMaybeChanged, "Message size should have been changed")
 
     verifyRecordValidationStats(validatedResults.recordValidationStats, numConvertedRecords = 3, records,
@@ -583,8 +585,8 @@ class LogValidatorTest {
       assertEquals(RecordBatch.NO_SEQUENCE, batch.baseSequence)
     }
     assertEquals(timestamp, validatedResults.maxTimestampMs)
-    assertEquals(validatedRecords.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestampMs,
-      s"Offset of max timestamp should be ${validatedRecords.records.asScala.size - 1}")
+    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+      s"Offset of max timestamp should be 0 when multiple records having the same max timestamp.")
     assertTrue(validatedResults.messageSizeMaybeChanged, "Message size should have been changed")
 
     verifyRecordValidationStats(validatedResults.recordValidationStats, numConvertedRecords = 3, records,
@@ -598,6 +600,7 @@ class LogValidatorTest {
 
   private def checkCompressed(magic: Byte): Unit = {
     val now = System.currentTimeMillis()
+    // set the timestamp of seq(1) (i.e. offset 1) as the max timestamp
     val timestampSeq = Seq(now - 1, now + 1, now)
 
     val (producerId, producerEpoch, baseSequence, isTransactional, partitionLeaderEpoch) =
@@ -654,11 +657,9 @@ class LogValidatorTest {
     }
     assertEquals(now + 1, validatedResults.maxTimestampMs, s"Max timestamp should be ${now + 1}")
 
-    // All versions have an outer batch when compressed, so the shallow offset
-    // of max timestamp is always the offset of the last record in the batch.
-    val expectedShallowOffsetOfMaxTimestamp = recordList.size - 1
+    val expectedShallowOffsetOfMaxTimestamp = 1
     assertEquals(expectedShallowOffsetOfMaxTimestamp, validatedResults.shallowOffsetOfMaxTimestampMs,
-      s"Offset of max timestamp should be ${validatedRecords.records.asScala.size - 1}")
+      s"Offset of max timestamp should be 1")
     assertFalse(validatedResults.messageSizeMaybeChanged, "Message size should not have been changed")
 
     verifyRecordValidationStats(validatedResults.recordValidationStats, numConvertedRecords = 0, records,
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
index 378247fa27030..0cf9cd1c60f73 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
@@ -330,6 +330,8 @@ public ValidationResult validateMessagesAndAssignOffsetsCompressed(LongRef offse
         long maxTimestamp = RecordBatch.NO_TIMESTAMP;
         LongRef expectedInnerOffset = PrimitiveRef.ofLong(0);
         List<Record> validatedRecords = new ArrayList<>();
+        long offsetOfMaxTimestamp = -1;
+        long initialOffset = offsetCounter.value;
 
         int uncompressedSizeInBytes = 0;
 
@@ -379,8 +381,11 @@ public ValidationResult validateMessagesAndAssignOffsetsCompressed(LongRef offse
                             && batch.magic() > RecordBatch.MAGIC_VALUE_V0
                             && toMagic > RecordBatch.MAGIC_VALUE_V0) {
 
-                        if (record.timestamp() > maxTimestamp)
+                        if (record.timestamp() > maxTimestamp) {
                             maxTimestamp = record.timestamp();
+                            // The offset is only increased when it is a valid record
+                            offsetOfMaxTimestamp = initialOffset + validatedRecords.size();
+                        }
 
                         // Some older clients do not implement the V1 internal offsets correctly.
                         // Historically the broker handled this by rewriting the batches rather
@@ -417,8 +422,10 @@ public ValidationResult validateMessagesAndAssignOffsetsCompressed(LongRef offse
             long lastOffset = offsetCounter.value - 1;
             firstBatch.setLastOffset(lastOffset);
 
-            if (timestampType == TimestampType.LOG_APPEND_TIME)
+            if (timestampType == TimestampType.LOG_APPEND_TIME) {
                 maxTimestamp = now;
+                offsetOfMaxTimestamp = initialOffset;
+            }
 
             if (toMagic >= RecordBatch.MAGIC_VALUE_V1)
                 firstBatch.setMaxTimestamp(timestampType, maxTimestamp);
@@ -431,7 +438,7 @@ public ValidationResult validateMessagesAndAssignOffsetsCompressed(LongRef offse
                 now,
                 records,
                 maxTimestamp,
-                lastOffset,
+                offsetOfMaxTimestamp,
                 false,
                 recordValidationStats);
         }
diff --git a/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java b/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
index 417bbe711681a..ce9c1718ffb67 100644
--- a/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
@@ -30,6 +30,7 @@
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.common.utils.Exit;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.extension.ExtendWith;
 
@@ -52,20 +53,24 @@ public class GetOffsetShellTest {
     private final int topicCount = 4;
     private final int offsetTopicPartitionCount = 4;
     private final ClusterInstance cluster;
+    private final String topicName = "topic";
 
     public GetOffsetShellTest(ClusterInstance cluster) {
         this.cluster = cluster;
     }
 
     private String getTopicName(int i) {
-        return "topic" + i;
+        return topicName + i;
     }
 
-    public void setUp() {
+    @BeforeEach
+    public void before() {
         cluster.config().serverProperties().put("auto.create.topics.enable", false);
         cluster.config().serverProperties().put("offsets.topic.replication.factor", "1");
         cluster.config().serverProperties().put("offsets.topic.num.partitions", String.valueOf(offsetTopicPartitionCount));
+    }
 
+    private void setUp() {
         try (Admin admin = Admin.create(cluster.config().adminClientProperties())) {
             List<NewTopic> topics = new ArrayList<>();
 
@@ -333,7 +338,7 @@ private void assertExitCodeIsOne(String... args) {
     }
 
     private List<Row> expectedOffsetsWithInternal() {
-        List<Row> consOffsets = IntStream.range(0, offsetTopicPartitionCount + 1)
+        List<Row> consOffsets = IntStream.range(0, offsetTopicPartitionCount)
                 .mapToObj(i -> new Row("__consumer_offsets", i, 0L))
                 .collect(Collectors.toList());
 

From e4c53d093eb70c29395e98b732749964572085a8 Mon Sep 17 00:00:00 2001
From: Kamal Chandraprakash <kchandraprakash@uber.com>
Date: Fri, 15 Mar 2024 08:03:40 +0530
Subject: [PATCH 155/258] KAFKA-15206: Fix the flaky
 RemoteIndexCacheTest.testClose test (#15523)

It is possible that due to resource constraint, ShutdownableThread#run might be called later than the ShutdownableThread#close method.

Reviewers: Luke Chen <showuon@gmail.com>, Divij Vaidya <diviv@amazon.com>
---
 .../kafka/log/remote/RemoteIndexCacheTest.scala |  2 ++
 .../kafka/utils/ShutdownableThreadTest.scala    | 17 ++++++++++++++++-
 .../kafka/server/util/ShutdownableThread.java   |  4 ++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/unit/kafka/log/remote/RemoteIndexCacheTest.scala b/core/src/test/scala/unit/kafka/log/remote/RemoteIndexCacheTest.scala
index 48b01503bb61b..ff65868674bbd 100644
--- a/core/src/test/scala/unit/kafka/log/remote/RemoteIndexCacheTest.scala
+++ b/core/src/test/scala/unit/kafka/log/remote/RemoteIndexCacheTest.scala
@@ -340,6 +340,8 @@ class RemoteIndexCacheTest {
     val spyEntry = generateSpyCacheEntry()
     cache.internalCache.put(rlsMetadata.remoteLogSegmentId().id(), spyEntry)
 
+    TestUtils.waitUntilTrue(() => cache.cleanerThread().isStarted, "Cleaner thread should be started")
+
     // close the cache
     cache.close()
 
diff --git a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala b/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
index 7b086d11509e3..da998903d2746 100644
--- a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
@@ -20,7 +20,7 @@ import java.util.concurrent.{CountDownLatch, TimeUnit}
 import org.apache.kafka.common.internals.FatalExitError
 import org.apache.kafka.server.util.ShutdownableThread
 import org.junit.jupiter.api.{AfterEach, Test}
-import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue}
+import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertTrue}
 
 class ShutdownableThreadTest {
 
@@ -51,4 +51,19 @@ class ShutdownableThreadTest {
     assertEquals(1, statusCodeOption.get)
   }
 
+  @Test
+  def testIsThreadStarted(): Unit = {
+    val latch = new CountDownLatch(1)
+    val thread = new ShutdownableThread("shutdownable-thread-test") {
+      override def doWork(): Unit = {
+        latch.countDown()
+      }
+    }
+    assertFalse(thread.isStarted)
+    thread.start()
+    latch.await()
+    assertTrue(thread.isStarted)
+
+    thread.shutdown()
+  }
 }
diff --git a/server-common/src/main/java/org/apache/kafka/server/util/ShutdownableThread.java b/server-common/src/main/java/org/apache/kafka/server/util/ShutdownableThread.java
index 06c751e0bb2c2..4268af92f84f5 100644
--- a/server-common/src/main/java/org/apache/kafka/server/util/ShutdownableThread.java
+++ b/server-common/src/main/java/org/apache/kafka/server/util/ShutdownableThread.java
@@ -68,6 +68,10 @@ public boolean isShutdownComplete() {
         return shutdownComplete.getCount() == 0;
     }
 
+    public boolean isStarted() {
+        return isStarted;
+    }
+
     /**
      * @return true if there has been an unexpected error and the thread shut down
      */

From 96bfac42168c5b195bd2e0551f21697ca73ed9f7 Mon Sep 17 00:00:00 2001
From: "A. Sophie Blee-Goldman" <ableegoldman@gmail.com>
Date: Thu, 14 Mar 2024 23:08:39 -0700
Subject: [PATCH 156/258] MINOR: Update javadocs and exception string in
 "deprecated" ProcessorRecordContext#hashcode (#15508)

This PR updates the javadocs for the "deprecated" hashCode() method of ProcessorRecordContext, as well as the UnsupportedOperationException thrown in its implementation, to actually explain why the class is mutable and therefore unsafe for use in hash collections. They now point out the mutable field in the class (namely the Headers)

Reviewers: Matthias Sax <mjsax@apache.org>, Bruno Cadonna <cadonna@apache.org>
---
 .../streams/processor/internals/ProcessorRecordContext.java | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
index b68b24036fff1..839baaad87528 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
@@ -193,12 +193,14 @@ public boolean equals(final Object o) {
     }
 
     /**
-     * Equality is implemented in support of tests, *not* for use in Hash collections, since this class is mutable.
+     * Equality is implemented in support of tests, *not* for use in Hash collections, since this class is mutable
+     * due to the {@link Headers} field it contains.
      */
     @Deprecated
     @Override
     public int hashCode() {
-        throw new UnsupportedOperationException("ProcessorRecordContext is unsafe for use in Hash collections");
+        throw new UnsupportedOperationException("ProcessorRecordContext is unsafe for use in Hash collections "
+                                                    + "due to the mutable Headers field");
     }
 
     @Override

From 7b9f31e35c0304bec191c5fbe2b880386e43f895 Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Fri, 15 Mar 2024 06:56:36 -0400
Subject: [PATCH 157/258] KAFKA-16358: Sort transformations by name in
 documentation; add missing transformations to documentation; add hyperlinks
 (#15499)

Reviewers: Yash Mayya <yash.mayya@gmail.com>
---
 .../connect/tools/TransformationDoc.java      | 24 +++++++--------
 docs/connect.html                             | 29 ++++++++++---------
 2 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/TransformationDoc.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/TransformationDoc.java
index 543703a13ac07..2c7250eb588c3 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/TransformationDoc.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/TransformationDoc.java
@@ -53,22 +53,22 @@ private DocInfo(String transformationName, String overview, ConfigDef configDef)
     }
 
     private static final List<DocInfo> TRANSFORMATIONS = Arrays.asList(
+            new DocInfo(Cast.class.getName(), Cast.OVERVIEW_DOC, Cast.CONFIG_DEF),
+            new DocInfo(DropHeaders.class.getName(), DropHeaders.OVERVIEW_DOC, DropHeaders.CONFIG_DEF),
+            new DocInfo(ExtractField.class.getName(), ExtractField.OVERVIEW_DOC, ExtractField.CONFIG_DEF),
+            new DocInfo(Filter.class.getName(), Filter.OVERVIEW_DOC, Filter.CONFIG_DEF),
+            new DocInfo(Flatten.class.getName(), Flatten.OVERVIEW_DOC, Flatten.CONFIG_DEF),
+            new DocInfo(HeaderFrom.class.getName(), HeaderFrom.OVERVIEW_DOC, HeaderFrom.CONFIG_DEF),
+            new DocInfo(HoistField.class.getName(), HoistField.OVERVIEW_DOC, HoistField.CONFIG_DEF),
             new DocInfo(InsertField.class.getName(), InsertField.OVERVIEW_DOC, InsertField.CONFIG_DEF),
-            new DocInfo(ReplaceField.class.getName(), ReplaceField.OVERVIEW_DOC, ReplaceField.CONFIG_DEF),
+            new DocInfo(InsertHeader.class.getName(), InsertHeader.OVERVIEW_DOC, InsertHeader.CONFIG_DEF),
             new DocInfo(MaskField.class.getName(), MaskField.OVERVIEW_DOC, MaskField.CONFIG_DEF),
-            new DocInfo(ValueToKey.class.getName(), ValueToKey.OVERVIEW_DOC, ValueToKey.CONFIG_DEF),
-            new DocInfo(HoistField.class.getName(), HoistField.OVERVIEW_DOC, HoistField.CONFIG_DEF),
-            new DocInfo(ExtractField.class.getName(), ExtractField.OVERVIEW_DOC, ExtractField.CONFIG_DEF),
-            new DocInfo(SetSchemaMetadata.class.getName(), SetSchemaMetadata.OVERVIEW_DOC, SetSchemaMetadata.CONFIG_DEF),
-            new DocInfo(TimestampRouter.class.getName(), TimestampRouter.OVERVIEW_DOC, TimestampRouter.CONFIG_DEF),
             new DocInfo(RegexRouter.class.getName(), RegexRouter.OVERVIEW_DOC, RegexRouter.CONFIG_DEF),
-            new DocInfo(Flatten.class.getName(), Flatten.OVERVIEW_DOC, Flatten.CONFIG_DEF),
-            new DocInfo(Cast.class.getName(), Cast.OVERVIEW_DOC, Cast.CONFIG_DEF),
+            new DocInfo(ReplaceField.class.getName(), ReplaceField.OVERVIEW_DOC, ReplaceField.CONFIG_DEF),
+            new DocInfo(SetSchemaMetadata.class.getName(), SetSchemaMetadata.OVERVIEW_DOC, SetSchemaMetadata.CONFIG_DEF),
             new DocInfo(TimestampConverter.class.getName(), TimestampConverter.OVERVIEW_DOC, TimestampConverter.CONFIG_DEF),
-            new DocInfo(Filter.class.getName(), Filter.OVERVIEW_DOC, Filter.CONFIG_DEF),
-            new DocInfo(InsertHeader.class.getName(), InsertHeader.OVERVIEW_DOC, InsertHeader.CONFIG_DEF),
-            new DocInfo(DropHeaders.class.getName(), DropHeaders.OVERVIEW_DOC, DropHeaders.CONFIG_DEF),
-            new DocInfo(HeaderFrom.class.getName(), HeaderFrom.OVERVIEW_DOC, HeaderFrom.CONFIG_DEF)
+            new DocInfo(TimestampRouter.class.getName(), TimestampRouter.OVERVIEW_DOC, TimestampRouter.CONFIG_DEF),
+            new DocInfo(ValueToKey.class.getName(), ValueToKey.OVERVIEW_DOC, ValueToKey.CONFIG_DEF)
     );
 
     private static void printTransformationHtml(PrintStream out, DocInfo docInfo) {
diff --git a/docs/connect.html b/docs/connect.html
index 954c7c647d8f8..c06879813fb98 100644
--- a/docs/connect.html
+++ b/docs/connect.html
@@ -162,19 +162,22 @@ <h5><a id="connect_included_transformation" href="#connect_included_transformati
     <p>Several widely-applicable data and routing transformations are included with Kafka Connect:</p>
 
     <ul>
-        <li>InsertField - Add a field using either static data or record metadata</li>
-        <li>ReplaceField - Filter or rename fields</li>
-        <li>MaskField - Replace field with valid null value for the type (0, empty string, etc) or custom replacement (non-empty string or numeric value only)</li>
-        <li>ValueToKey - Replace the record key with a new key formed from a subset of fields in the record value</li>
-        <li>HoistField - Wrap the entire event as a single field inside a Struct or a Map</li>
-        <li>ExtractField - Extract a specific field from Struct and Map and include only this field in results</li>
-        <li>SetSchemaMetadata - modify the schema name or version</li>
-        <li>TimestampRouter - Modify the topic of a record based on original topic and timestamp. Useful when using a sink that needs to write to different tables or indexes based on timestamps</li>
-        <li>RegexRouter - modify the topic of a record based on original topic, replacement string and a regular expression</li>
-        <li>Filter - Removes messages from all further processing. This is used with a <a href="#connect_predicates">predicate</a> to selectively filter certain messages.</li>
-        <li>InsertHeader - Add a header using static data</li>
-        <li>HeadersFrom - Copy or move fields in the key or value to the record headers</li>
-        <li>DropHeaders - Remove headers by name</li>
+        <li><a href="#org.apache.kafka.connect.transforms.Cast">Cast</a> - Cast fields or the entire key or value to a specific type</li>
+        <li><a href="#org.apache.kafka.connect.transforms.DropHeaders">DropHeaders</a> - Remove headers by name</li>
+        <li><a href="#org.apache.kafka.connect.transforms.ExtractField">ExtractField</a> - Extract a specific field from Struct and Map and include only this field in results</li>
+        <li><a href="#org.apache.kafka.connect.transforms.Filter">Filter</a> - Removes messages from all further processing. This is used with a <a href="#connect_predicates">predicate</a> to selectively filter certain messages</li>
+        <li><a href="#org.apache.kafka.connect.transforms.Flatten">Flatten</a> - Flatten a nested data structure</li>
+        <li><a href="#org.apache.kafka.connect.transforms.HeaderFrom">HeaderFrom</a> - Copy or move fields in the key or value to the record headers</li>
+        <li><a href="#org.apache.kafka.connect.transforms.HoistField">HoistField</a> - Wrap the entire event as a single field inside a Struct or a Map</li>
+        <li><a href="#org.apache.kafka.connect.transforms.InsertField">InsertField</a> - Add a field using either static data or record metadata</li>
+        <li><a href="#org.apache.kafka.connect.transforms.InsertHeader">InsertHeader</a> - Add a header using static data</li>
+        <li><a href="#org.apache.kafka.connect.transforms.MaskField">MaskField</a> - Replace field with valid null value for the type (0, empty string, etc) or custom replacement (non-empty string or numeric value only)</li>
+        <li><a href="#org.apache.kafka.connect.transforms.RegexRouter">RegexRouter</a> - modify the topic of a record based on original topic, replacement string and a regular expression</li>
+        <li><a href="#org.apache.kafka.connect.transforms.ReplaceField">ReplaceField</a> - Filter or rename fields</li>
+        <li><a href="#org.apache.kafka.connect.transforms.SetSchemaMetadata">SetSchemaMetadata</a> - modify the schema name or version</li>
+        <li><a href="#org.apache.kafka.connect.transforms.TimestampConverter">TimestampConverter</a> - Convert timestamps between different formats</li>
+        <li><a href="#org.apache.kafka.connect.transforms.TimestampRouter">TimestampRouter</a> - Modify the topic of a record based on original topic and timestamp. Useful when using a sink that needs to write to different tables or indexes based on timestamps</li>
+        <li><a href="#org.apache.kafka.connect.transforms.ValueToKey">ValueToKey</a> - Replace the record key with a new key formed from a subset of fields in the record value</li>
     </ul>
 
     <p>Details on how to configure each transformation are listed below:</p>

From fa190cf18e6101b8e0212b6f79bb585ccfeb0437 Mon Sep 17 00:00:00 2001
From: Dongnuo Lyu <139248811+dongnuo123@users.noreply.github.com>
Date: Fri, 15 Mar 2024 08:24:59 -0400
Subject: [PATCH 158/258] MINOR: Only enable replay methods to modify timeline
 data structure (#15528)

The patch prevents the main method (the method generating records) from modifying the timeline data structure `groups`  by calling `getOrMaybeCreateConsumerGroup` in kip-848 new group coordinator. Only replay methods are able to add the newly created group to `groups`.

Reviewers: David Jacot <djacot@confluent.io>
---
 .../group/GroupMetadataManager.java           | 55 ++++++++++++++++---
 .../group/GroupMetadataManagerTest.java       | 36 ++++++------
 .../group/OffsetMetadataManagerTest.java      | 42 +++++++-------
 3 files changed, 85 insertions(+), 48 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 26c2f644e9612..48f0618c55d34 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -583,7 +583,8 @@ public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(
     }
 
     /**
-     * Gets or maybe creates a consumer group.
+     * Gets or maybe creates a consumer group without updating the groups map.
+     * The group will be materialized during the replay.
      *
      * @param groupId           The group id.
      * @param createIfNotExists A boolean indicating whether the group should be
@@ -605,6 +606,42 @@ ConsumerGroup getOrMaybeCreateConsumerGroup(
             throw new GroupIdNotFoundException(String.format("Consumer group %s not found.", groupId));
         }
 
+        if (group == null) {
+            return new ConsumerGroup(snapshotRegistry, groupId, metrics);
+        } else {
+            if (group.type() == CONSUMER) {
+                return (ConsumerGroup) group;
+            } else {
+                // We don't support upgrading/downgrading between protocols at the moment so
+                // we throw an exception if a group exists with the wrong type.
+                throw new GroupIdNotFoundException(String.format("Group %s is not a consumer group.", groupId));
+            }
+        }
+    }
+
+    /**
+     * The method should be called on the replay path.
+     * Gets or maybe creates a consumer group and updates the groups map if a new group is created.
+     *
+     * @param groupId           The group id.
+     * @param createIfNotExists A boolean indicating whether the group should be
+     *                          created if it does not exist.
+     *
+     * @return A ConsumerGroup.
+     * @throws IllegalStateException if the group does not exist and createIfNotExists is false or
+     *                               if the group is not a consumer group.
+     * Package private for testing.
+     */
+    ConsumerGroup getOrMaybeCreatePersistedConsumerGroup(
+        String groupId,
+        boolean createIfNotExists
+    ) throws GroupIdNotFoundException {
+        Group group = groups.get(groupId);
+
+        if (group == null && !createIfNotExists) {
+            throw new IllegalStateException(String.format("Consumer group %s not found.", groupId));
+        }
+
         if (group == null) {
             ConsumerGroup consumerGroup = new ConsumerGroup(snapshotRegistry, groupId, metrics);
             groups.put(groupId, consumerGroup);
@@ -616,7 +653,7 @@ ConsumerGroup getOrMaybeCreateConsumerGroup(
             } else {
                 // We don't support upgrading/downgrading between protocols at the moment so
                 // we throw an exception if a group exists with the wrong type.
-                throw new GroupIdNotFoundException(String.format("Group %s is not a consumer group.", groupId));
+                throw new IllegalStateException(String.format("Group %s is not a consumer group.", groupId));
             }
         }
     }
@@ -1551,7 +1588,7 @@ public void replay(
         String groupId = key.groupId();
         String memberId = key.memberId();
 
-        ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, value != null);
+        ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, value != null);
         Set<String> oldSubscribedTopicNames = new HashSet<>(consumerGroup.subscribedTopicNames());
 
         if (value != null) {
@@ -1663,10 +1700,10 @@ public void replay(
         String groupId = key.groupId();
 
         if (value != null) {
-            ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, true);
+            ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, true);
             consumerGroup.setGroupEpoch(value.epoch());
         } else {
-            ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, false);
+            ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, false);
             if (!consumerGroup.members().isEmpty()) {
                 throw new IllegalStateException("Received a tombstone record to delete group " + groupId
                     + " but the group still has " + consumerGroup.members().size() + " members.");
@@ -1698,7 +1735,7 @@ public void replay(
         ConsumerGroupPartitionMetadataValue value
     ) {
         String groupId = key.groupId();
-        ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, false);
+        ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, false);
 
         if (value != null) {
             Map<String, TopicMetadata> subscriptionMetadata = new HashMap<>();
@@ -1724,7 +1761,7 @@ public void replay(
     ) {
         String groupId = key.groupId();
         String memberId = key.memberId();
-        ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, false);
+        ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, false);
 
         if (value != null) {
             consumerGroup.updateTargetAssignment(memberId, Assignment.fromRecord(value));
@@ -1746,7 +1783,7 @@ public void replay(
         ConsumerGroupTargetAssignmentMetadataValue value
     ) {
         String groupId = key.groupId();
-        ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, false);
+        ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, false);
 
         if (value != null) {
             consumerGroup.setTargetAssignmentEpoch(value.assignmentEpoch());
@@ -1772,7 +1809,7 @@ public void replay(
     ) {
         String groupId = key.groupId();
         String memberId = key.memberId();
-        ConsumerGroup consumerGroup = getOrMaybeCreateConsumerGroup(groupId, false);
+        ConsumerGroup consumerGroup = getOrMaybeCreatePersistedConsumerGroup(groupId, false);
         ConsumerGroupMember oldMember = consumerGroup.getOrMaybeCreateMember(memberId, false);
 
         if (value != null) {
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index 66592417828e9..e9304407cdd20 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -9112,43 +9112,43 @@ public void testClassicGroupMaybeDelete() {
 
     @Test
     public void testConsumerGroupDelete() {
+        String groupId = "group-id";
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10))
             .build();
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group-id", true);
 
         List<Record> expectedRecords = Arrays.asList(
-            RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id"),
-            RecordHelpers.newGroupSubscriptionMetadataTombstoneRecord("group-id"),
-            RecordHelpers.newGroupEpochTombstoneRecord("group-id")
+            RecordHelpers.newTargetAssignmentEpochTombstoneRecord(groupId),
+            RecordHelpers.newGroupSubscriptionMetadataTombstoneRecord(groupId),
+            RecordHelpers.newGroupEpochTombstoneRecord(groupId)
         );
         List<Record> records = new ArrayList<>();
-        context.groupMetadataManager.deleteGroup("group-id", records);
+        context.groupMetadataManager.deleteGroup(groupId, records);
         assertEquals(expectedRecords, records);
     }
 
     @Test
     public void testConsumerGroupMaybeDelete() {
+        String groupId = "group-id";
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10))
             .build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group-id", true);
 
         List<Record> expectedRecords = Arrays.asList(
-            RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id"),
-            RecordHelpers.newGroupSubscriptionMetadataTombstoneRecord("group-id"),
-            RecordHelpers.newGroupEpochTombstoneRecord("group-id")
+            RecordHelpers.newTargetAssignmentEpochTombstoneRecord(groupId),
+            RecordHelpers.newGroupSubscriptionMetadataTombstoneRecord(groupId),
+            RecordHelpers.newGroupEpochTombstoneRecord(groupId)
         );
         List<Record> records = new ArrayList<>();
-        context.groupMetadataManager.maybeDeleteGroup("group-id", records);
+        context.groupMetadataManager.maybeDeleteGroup(groupId, records);
         assertEquals(expectedRecords, records);
 
         records = new ArrayList<>();
-        group.updateMember(new ConsumerGroupMember.Builder("member")
-            .setState(MemberState.STABLE)
+        context.replay(RecordHelpers.newMemberSubscriptionRecord(groupId, new ConsumerGroupMember.Builder("member")
             .setMemberEpoch(10)
             .setPreviousMemberEpoch(10)
-            .build()
-        );
-        context.groupMetadataManager.maybeDeleteGroup("group-id", records);
+            .build()));
+        context.groupMetadataManager.maybeDeleteGroup(groupId, records);
         assertEquals(Collections.emptyList(), records);
     }
 
@@ -9273,7 +9273,7 @@ public void testOnConsumerGroupStateTransition() {
         verify(context.metrics, times(1)).onConsumerGroupStateTransition(ConsumerGroup.ConsumerGroupState.EMPTY, null);
 
         // Replaying a tombstone for a group that has already been removed should not decrement metric.
-        tombstones.forEach(tombstone -> assertThrows(GroupIdNotFoundException.class, () -> context.replay(tombstone)));
+        tombstones.forEach(tombstone -> assertThrows(IllegalStateException.class, () -> context.replay(tombstone)));
         verify(context.metrics, times(1)).onConsumerGroupStateTransition(ConsumerGroup.ConsumerGroupState.EMPTY, null);
     }
 
@@ -9290,8 +9290,8 @@ public void testOnConsumerGroupStateTransitionOnLoading() {
         context.replay(RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id"));
         context.replay(RecordHelpers.newGroupEpochTombstoneRecord("group-id"));
         IntStream.range(0, 3).forEach(__ -> {
-            assertThrows(GroupIdNotFoundException.class, () -> context.replay(RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id")));
-            assertThrows(GroupIdNotFoundException.class, () -> context.replay(RecordHelpers.newGroupEpochTombstoneRecord("group-id")));
+            assertThrows(IllegalStateException.class, () -> context.replay(RecordHelpers.newTargetAssignmentEpochTombstoneRecord("group-id")));
+            assertThrows(IllegalStateException.class, () -> context.replay(RecordHelpers.newGroupEpochTombstoneRecord("group-id")));
         });
 
         verify(context.metrics, times(1)).onConsumerGroupStateTransition(null, ConsumerGroup.ConsumerGroupState.EMPTY);
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
index d9eacb46c3ea3..4b3d76bbd1c11 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/OffsetMetadataManagerTest.java
@@ -194,7 +194,7 @@ public Group getOrMaybeCreateGroup(
                         true
                     );
                 case CONSUMER:
-                    return groupMetadataManager.getOrMaybeCreateConsumerGroup(
+                    return groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
                         groupId,
                         true
                     );
@@ -1079,7 +1079,7 @@ public void testConsumerGroupOffsetCommitWithUnknownMemberId() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1108,7 +1108,7 @@ public void testConsumerGroupOffsetCommitWithStaleMemberEpoch() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1153,7 +1153,7 @@ public void testConsumerGroupOffsetCommitWithVersionSmallerThanVersion9(short ve
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1190,7 +1190,7 @@ public void testConsumerGroupOffsetCommitFromAdminClient() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1246,7 +1246,7 @@ public void testConsumerGroupOffsetCommit() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1316,7 +1316,7 @@ public void testConsumerGroupOffsetCommitWithOffsetMetadataTooLarge() {
             .build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1393,7 +1393,7 @@ public void testConsumerGroupTransactionalOffsetCommit() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1483,7 +1483,7 @@ public void testConsumerGroupTransactionalOffsetCommitWithUnknownMemberId() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1512,7 +1512,7 @@ public void testConsumerGroupTransactionalOffsetCommitWithStaleMemberEpoch() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
         // Create an empty group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -1775,7 +1775,7 @@ public void testFetchOffsetsWithUnknownGroup() {
     public void testFetchOffsetsAtDifferentCommittedOffset() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
 
         assertEquals(0, context.lastWrittenOffset);
         context.commitOffset("group", "foo", 0, 100L, 1);
@@ -1916,7 +1916,7 @@ public void testFetchOffsetsAtDifferentCommittedOffset() {
     public void testFetchOffsetsWithPendingTransactionalOffsets() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
 
         context.commitOffset("group", "foo", 0, 100L, 1);
         context.commitOffset("group", "foo", 1, 110L, 1);
@@ -2021,7 +2021,7 @@ public void testFetchAllOffsetsWithUnknownGroup() {
     public void testFetchAllOffsetsAtDifferentCommittedOffset() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
 
         assertEquals(0, context.lastWrittenOffset);
         context.commitOffset("group", "foo", 0, 100L, 1);
@@ -2108,7 +2108,7 @@ public void testFetchAllOffsetsAtDifferentCommittedOffset() {
     public void testFetchAllOffsetsWithPendingTransactionalOffsets() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
 
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
 
         context.commitOffset("group", "foo", 0, 100L, 1);
         context.commitOffset("group", "foo", 1, 110L, 1);
@@ -2182,7 +2182,7 @@ public void testFetchAllOffsetsWithPendingTransactionalOffsets() {
     public void testConsumerGroupOffsetFetchWithMemberIdAndEpoch() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
         // Create consumer group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
         // Create member.
         group.getOrMaybeCreateMember("member", true);
         // Commit offset.
@@ -2217,7 +2217,7 @@ public void testConsumerGroupOffsetFetchWithMemberIdAndEpoch() {
     public void testConsumerGroupOffsetFetchFromAdminClient() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
         // Create consumer group.
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
         // Create member.
         group.getOrMaybeCreateMember("member", true);
         // Commit offset.
@@ -2251,7 +2251,7 @@ public void testConsumerGroupOffsetFetchFromAdminClient() {
     @Test
     public void testConsumerGroupOffsetFetchWithUnknownMemberId() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
 
         // Fetch offsets case.
         List<OffsetFetchRequestData.OffsetFetchRequestTopics> topics = Collections.singletonList(
@@ -2276,7 +2276,7 @@ public void testConsumerGroupOffsetFetchWithUnknownMemberId() {
     @Test
     public void testConsumerGroupOffsetFetchWithStaleMemberEpoch() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group", true);
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup("group", true);
         group.getOrMaybeCreateMember("member", true);
 
         // Fetch offsets case.
@@ -2340,7 +2340,7 @@ public void testGenericGroupOffsetDeleteWithPendingTransactionalOffsets() {
     @Test
     public void testConsumerGroupOffsetDelete() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -2352,7 +2352,7 @@ public void testConsumerGroupOffsetDelete() {
     @Test
     public void testConsumerGroupOffsetDeleteWithErrors() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );
@@ -2382,7 +2382,7 @@ public void testConsumerGroupOffsetDeleteWithErrors() {
     @Test
     public void testConsumerGroupOffsetDeleteWithPendingTransactionalOffsets() {
         OffsetMetadataManagerTestContext context = new OffsetMetadataManagerTestContext.Builder().build();
-        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(
+        ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreatePersistedConsumerGroup(
             "foo",
             true
         );

From af0ec247ccf7b1d97ca684109854bf1881fd11bb Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 15 Mar 2024 06:17:22 -0700
Subject: [PATCH 159/258] =?UTF-8?q?KAFKA-16231:=20Update=20consumer=5Ftest?=
 =?UTF-8?q?.py=20to=20support=20KIP-848=E2=80=99s=20group=20protocol=20con?=
 =?UTF-8?q?fig=20(#15330)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the setup_consumer method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../services/kafka/consumer_group.py          | 42 +++++++++++
 .../kafkatest/services/verifiable_consumer.py |  2 +
 tests/kafkatest/tests/client/consumer_test.py | 73 ++++++++++++-------
 .../tests/verifiable_consumer_test.py         |  3 +-
 4 files changed, 93 insertions(+), 27 deletions(-)
 create mode 100644 tests/kafkatest/services/kafka/consumer_group.py

diff --git a/tests/kafkatest/services/kafka/consumer_group.py b/tests/kafkatest/services/kafka/consumer_group.py
new file mode 100644
index 0000000000000..e94bd6382d265
--- /dev/null
+++ b/tests/kafkatest/services/kafka/consumer_group.py
@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# These are the group protocols we support. Most tests that use the new group coordinator will
+# (eventually) be upgraded to test both of these consumer groups.
+classic_group_protocol = 'classic'
+consumer_group_protocol = 'consumer'
+all_group_protocols = [classic_group_protocol, consumer_group_protocol]
+
+# These are the remote assignors used by the new group coordinator.
+range_remote_assignor = 'range'
+uniform_remote_assignor = 'uniform'
+all_remote_assignors = [range_remote_assignor, uniform_remote_assignor]
+
+
+def is_consumer_group_protocol_enabled(group_protocol):
+    """Check if the KIP-848 consumer group protocol is enabled."""
+    return group_protocol is not None and group_protocol.lower() == consumer_group_protocol
+
+
+def maybe_set_group_protocol(group_protocol, config=None):
+    """Maybe include the KIP-848 group.protocol configuration if it's not None."""
+    if config is None:
+        config = {}
+
+    if group_protocol is not None:
+        config["group.protocol"] = group_protocol
+
+    return config
diff --git a/tests/kafkatest/services/verifiable_consumer.py b/tests/kafkatest/services/verifiable_consumer.py
index e1155c16aaef4..7ef5f75e22bef 100644
--- a/tests/kafkatest/services/verifiable_consumer.py
+++ b/tests/kafkatest/services/verifiable_consumer.py
@@ -177,6 +177,8 @@ def __init__(self, context, num_nodes, kafka, topic, group_id,
         self.log_level = log_level
         self.kafka = kafka
         self.topic = topic
+        self.group_protocol = group_protocol
+        self.group_remote_assignor = group_remote_assignor
         self.group_id = group_id
         self.reset_policy = reset_policy
         self.static_membership = static_membership
diff --git a/tests/kafkatest/tests/client/consumer_test.py b/tests/kafkatest/tests/client/consumer_test.py
index 6dc1fb897b06d..d4d6af9f2aad0 100644
--- a/tests/kafkatest/tests/client/consumer_test.py
+++ b/tests/kafkatest/tests/client/consumer_test.py
@@ -18,7 +18,7 @@
 from ducktape.mark.resource import cluster
 
 from kafkatest.tests.verifiable_consumer_test import VerifiableConsumerTest
-from kafkatest.services.kafka import TopicPartition, quorum
+from kafkatest.services.kafka import TopicPartition, quorum, consumer_group
 
 import signal
 
@@ -76,14 +76,15 @@ def setup_consumer(self, topic, **kwargs):
 
     @cluster(num_nodes=7)
     @matrix(
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify correct consumer behavior when the brokers are consecutively restarted.
 
@@ -108,7 +109,7 @@ def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordina
         # broker, and that coordinator will fail the consumer and trigger a group rebalance if its session times out.
         # This test is asserting that no rebalances occur, so we increase the session timeout for this to be the case.
         self.session_timeout_sec = 30
-        consumer = self.setup_consumer(self.TOPIC)
+        consumer = self.setup_consumer(self.TOPIC, group_protocol=group_protocol)
 
         producer.start()
         self.await_produced_messages(producer)
@@ -136,16 +137,17 @@ def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordina
     @matrix(
         clean_shutdown=[True],
         bounce_mode=["all", "rolling"],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         clean_shutdown=[True],
         bounce_mode=["all", "rolling"],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]   
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_consumer_bounce(self, clean_shutdown, bounce_mode, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_consumer_bounce(self, clean_shutdown, bounce_mode, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify correct consumer behavior when the consumers in the group are consecutively restarted.
 
@@ -160,7 +162,7 @@ def test_consumer_bounce(self, clean_shutdown, bounce_mode, metadata_quorum=quor
         partition = TopicPartition(self.TOPIC, 0)
 
         producer = self.setup_producer(self.TOPIC)
-        consumer = self.setup_consumer(self.TOPIC)
+        consumer = self.setup_consumer(self.TOPIC, group_protocol=group_protocol)
 
         producer.start()
         self.await_produced_messages(producer)
@@ -371,19 +373,20 @@ def test_fencing_static_consumer(self, num_conflict_consumers, fencing_stage, me
     @matrix(
         clean_shutdown=[True],
         enable_autocommit=[True, False],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         clean_shutdown=[True],
         enable_autocommit=[True, False],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_consumer_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_consumer_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         partition = TopicPartition(self.TOPIC, 0)
 
-        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit)
+        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit, group_protocol=group_protocol)
         producer = self.setup_producer(self.TOPIC)
 
         consumer.start()
@@ -436,12 +439,19 @@ def test_consumer_failure(self, clean_shutdown, enable_autocommit, metadata_quor
         clean_shutdown=[True, False],
         enable_autocommit=[True, False],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
     )
-    def test_broker_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    @matrix(
+        clean_shutdown=[True, False],
+        enable_autocommit=[True, False],
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
+    )
+    def test_broker_failure(self, clean_shutdown, enable_autocommit, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         partition = TopicPartition(self.TOPIC, 0)
 
-        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit)
+        consumer = self.setup_consumer(self.TOPIC, enable_autocommit=enable_autocommit, group_protocol=group_protocol)
         producer = self.setup_producer(self.TOPIC)
 
         producer.start()
@@ -475,14 +485,15 @@ def test_broker_failure(self, clean_shutdown, enable_autocommit, metadata_quorum
 
     @cluster(num_nodes=7)
     @matrix(
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_group_consumption(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_group_consumption(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verifies correct group rebalance behavior as consumers are started and stopped.
         In particular, this test verifies that the partition is readable after every
@@ -494,7 +505,7 @@ def test_group_consumption(self, metadata_quorum=quorum.zk, use_new_coordinator=
         - Start the consumers one by one, verifying consumption after each rebalance
         - Shutdown the consumers one by one, verifying consumption after each rebalance
         """
-        consumer = self.setup_consumer(self.TOPIC)
+        consumer = self.setup_consumer(self.TOPIC, group_protocol=group_protocol)
         producer = self.setup_producer(self.TOPIC)
 
         partition = TopicPartition(self.TOPIC, 0)
@@ -535,18 +546,25 @@ def __init__(self, test_context):
     @matrix(
         assignment_strategy=["org.apache.kafka.clients.consumer.RangeAssignor",
                              "org.apache.kafka.clients.consumer.RoundRobinAssignor",
-                             "org.apache.kafka.clients.consumer.StickyAssignor"], 
-        metadata_quorum=[quorum.zk],
+                             "org.apache.kafka.clients.consumer.StickyAssignor"],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         assignment_strategy=["org.apache.kafka.clients.consumer.RangeAssignor",
                              "org.apache.kafka.clients.consumer.RoundRobinAssignor",
-                             "org.apache.kafka.clients.consumer.StickyAssignor"], 
+                             "org.apache.kafka.clients.consumer.StickyAssignor"],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=[consumer_group.classic_group_protocol],
+    )
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=[consumer_group.consumer_group_protocol],
+        group_remote_assignor=consumer_group.all_remote_assignors
     )
-    def test_valid_assignment(self, assignment_strategy, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_valid_assignment(self, assignment_strategy=None, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None, group_remote_assignor=None):
         """
         Verify assignment strategy correctness: each partition is assigned to exactly
         one consumer instance.
@@ -556,7 +574,10 @@ def test_valid_assignment(self, assignment_strategy, metadata_quorum=quorum.zk,
         - Start the consumers one by one
         - Validate assignment after every expected rebalance
         """
-        consumer = self.setup_consumer(self.TOPIC, assignment_strategy=assignment_strategy)
+        consumer = self.setup_consumer(self.TOPIC,
+                                       assignment_strategy=assignment_strategy,
+                                       group_protocol=group_protocol,
+                                       group_remote_assignor=group_remote_assignor)
         for num_started, node in enumerate(consumer.nodes, 1):
             consumer.start_node(node)
             self.await_members(consumer, num_started)
diff --git a/tests/kafkatest/tests/verifiable_consumer_test.py b/tests/kafkatest/tests/verifiable_consumer_test.py
index e416690bfc402..38bb6cf3bd90f 100644
--- a/tests/kafkatest/tests/verifiable_consumer_test.py
+++ b/tests/kafkatest/tests/verifiable_consumer_test.py
@@ -54,10 +54,11 @@ def min_cluster_size(self):
         return super(VerifiableConsumerTest, self).min_cluster_size() + self.num_consumers + self.num_producers
 
     def setup_consumer(self, topic, static_membership=False, enable_autocommit=False,
-                       assignment_strategy="org.apache.kafka.clients.consumer.RangeAssignor", **kwargs):
+                       assignment_strategy="org.apache.kafka.clients.consumer.RangeAssignor", group_remote_assignor="range", **kwargs):
         return VerifiableConsumer(self.test_context, self.num_consumers, self.kafka,
                                   topic, self.group_id, static_membership=static_membership, session_timeout_sec=self.session_timeout_sec,
                                   assignment_strategy=assignment_strategy, enable_autocommit=enable_autocommit,
+                                  group_remote_assignor=group_remote_assignor,
                                   log_level="TRACE", **kwargs)
 
     def setup_producer(self, topic, max_messages=-1, throughput=500):

From e878654e95c3f82863a7b199fc3af56994f9841c Mon Sep 17 00:00:00 2001
From: Chris Holland <41524756+ChrisAHolland@users.noreply.github.com>
Date: Fri, 15 Mar 2024 06:18:24 -0700
Subject: [PATCH 160/258] MINOR: Cleanup BoundedList to Make Constructors More
 Safe (#15507)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/server/mutable/BoundedList.java     |   8 +-
 .../kafka/server/mutable/BoundedListTest.java | 120 +++++++++++-------
 2 files changed, 81 insertions(+), 47 deletions(-)

diff --git a/server-common/src/main/java/org/apache/kafka/server/mutable/BoundedList.java b/server-common/src/main/java/org/apache/kafka/server/mutable/BoundedList.java
index 6c5d1ba0d5a11..d31de006b7878 100644
--- a/server-common/src/main/java/org/apache/kafka/server/mutable/BoundedList.java
+++ b/server-common/src/main/java/org/apache/kafka/server/mutable/BoundedList.java
@@ -40,18 +40,22 @@ public static <E> BoundedList<E> newArrayBacked(int maxLength) {
     }
 
     public static <E> BoundedList<E> newArrayBacked(int maxLength, int initialCapacity) {
+        if (initialCapacity <= 0) {
+            throw new IllegalArgumentException("Invalid non-positive initialCapacity of " + initialCapacity);
+        }
         return new BoundedList<>(maxLength, new ArrayList<>(initialCapacity));
     }
 
-    public BoundedList(int maxLength, List<E> underlying) {
+    private BoundedList(int maxLength, List<E> underlying) {
         if (maxLength <= 0) {
             throw new IllegalArgumentException("Invalid non-positive maxLength of " + maxLength);
         }
-        this.maxLength = maxLength;
+
         if (underlying.size() > maxLength) {
             throw new BoundedListTooLongException("Cannot wrap list, because it is longer than " +
                 "the maximum length " + maxLength);
         }
+        this.maxLength = maxLength;
         this.underlying = underlying;
     }
 
diff --git a/server-common/src/test/java/org/apache/kafka/server/mutable/BoundedListTest.java b/server-common/src/test/java/org/apache/kafka/server/mutable/BoundedListTest.java
index a5a71f264bad3..df2608430f4ae 100644
--- a/server-common/src/test/java/org/apache/kafka/server/mutable/BoundedListTest.java
+++ b/server-common/src/test/java/org/apache/kafka/server/mutable/BoundedListTest.java
@@ -20,9 +20,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -33,33 +31,46 @@
 
 @Timeout(120)
 public class BoundedListTest {
+
     @Test
     public void testMaxLengthMustNotBeZero() {
         assertEquals("Invalid non-positive maxLength of 0",
             assertThrows(IllegalArgumentException.class,
-                () -> new BoundedList<>(0, new ArrayList<Integer>())).
-                    getMessage());
+                () -> BoundedList.newArrayBacked(0)).getMessage());
+
+        assertEquals("Invalid non-positive maxLength of 0",
+            assertThrows(IllegalArgumentException.class,
+                () -> BoundedList.newArrayBacked(0, 100)).getMessage());
     }
 
     @Test
     public void testMaxLengthMustNotBeNegative() {
         assertEquals("Invalid non-positive maxLength of -123",
             assertThrows(IllegalArgumentException.class,
-                () -> new BoundedList<>(-123, new ArrayList<Integer>())).
-                    getMessage());
+                () -> BoundedList.newArrayBacked(-123)).getMessage());
+
+        assertEquals("Invalid non-positive maxLength of -123",
+            assertThrows(IllegalArgumentException.class,
+                () -> BoundedList.newArrayBacked(-123, 100)).getMessage());
+    }
+
+    @Test
+    public void testInitialCapacityMustNotBeZero() {
+        assertEquals("Invalid non-positive initialCapacity of 0",
+            assertThrows(IllegalArgumentException.class,
+                () -> BoundedList.newArrayBacked(100, 0)).getMessage());
     }
 
     @Test
-    public void testOwnedListMustNotBeTooLong() {
-        assertEquals("Cannot wrap list, because it is longer than the maximum length 1",
-            assertThrows(BoundedListTooLongException.class,
-                () -> new BoundedList<>(1, new ArrayList<>(Arrays.asList(1, 2)))).
-                    getMessage());
+    public void testInitialCapacityMustNotBeNegative() {
+        assertEquals("Invalid non-positive initialCapacity of -123",
+            assertThrows(IllegalArgumentException.class,
+                () -> BoundedList.newArrayBacked(100, -123)).getMessage());
     }
 
     @Test
     public void testAddingToBoundedList() {
-        BoundedList<Integer> list = new BoundedList<>(2, new ArrayList<>(3));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(2);
         assertEquals(0, list.size());
         assertTrue(list.isEmpty());
         assertTrue(list.add(456));
@@ -70,42 +81,42 @@ public void testAddingToBoundedList() {
         assertEquals("Cannot add another element to the list because it would exceed the " +
             "maximum length of 2",
                 assertThrows(BoundedListTooLongException.class,
-                    () -> list.add(912)).
-                        getMessage());
+                    () -> list.add(912)).getMessage());
         assertEquals("Cannot add another element to the list because it would exceed the " +
             "maximum length of 2",
                 assertThrows(BoundedListTooLongException.class,
-                    () -> list.add(0, 912)).
-                        getMessage());
-    }
-
-    private static <E> void testHashCodeAndEquals(List<E> a) {
-        assertEquals(a, new BoundedList<>(123, a));
-        assertEquals(a.hashCode(), new BoundedList<>(123, a).hashCode());
-    }
-
-    @Test
-    public void testHashCodeAndEqualsForEmptyList() {
-        testHashCodeAndEquals(Collections.emptyList());
+                    () -> list.add(0, 912)).getMessage());
     }
 
     @Test
     public void testHashCodeAndEqualsForNonEmptyList() {
-        testHashCodeAndEquals(Arrays.asList(1, 2, 3, 4, 5, 6, 7));
+        BoundedList<Integer> boundedList = BoundedList.newArrayBacked(7);
+        List<Integer> otherList = Arrays.asList(1, 2, 3, 4, 5, 6, 7);
+        boundedList.addAll(otherList);
+
+        assertEquals(otherList, boundedList);
+        assertEquals(otherList.hashCode(), boundedList.hashCode());
     }
 
     @Test
     public void testSet() {
-        ArrayList<Integer> underlying = new ArrayList<>(Arrays.asList(1, 2, 3));
-        BoundedList<Integer> list = new BoundedList<>(3, underlying);
-        list.set(1, 200);
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(200);
+        list.add(3);
         assertEquals(Arrays.asList(1, 200, 3), list);
+        list.set(0, 100);
+        list.set(1, 200);
+        list.set(2, 300);
+        assertEquals(Arrays.asList(100, 200, 300), list);
     }
 
     @Test
     public void testRemove() {
-        ArrayList<String> underlying = new ArrayList<>(Arrays.asList("a", "a", "c"));
-        BoundedList<String> list = new BoundedList<>(3, underlying);
+        BoundedList<String> list = BoundedList.newArrayBacked(3);
+        list.add("a");
+        list.add("a");
+        list.add("c");
         assertEquals(0, list.indexOf("a"));
         assertEquals(1, list.lastIndexOf("a"));
         list.remove("a");
@@ -116,8 +127,10 @@ public void testRemove() {
 
     @Test
     public void testClear() {
-        ArrayList<String> underlying = new ArrayList<>(Arrays.asList("a", "b", "c"));
-        BoundedList<String> list = new BoundedList<>(3, underlying);
+        BoundedList<String> list = BoundedList.newArrayBacked(3);
+        list.add("a");
+        list.add("a");
+        list.add("c");
         list.clear();
         assertEquals(Arrays.asList(), list);
         assertTrue(list.isEmpty());
@@ -125,38 +138,49 @@ public void testClear() {
 
     @Test
     public void testGet() {
-        BoundedList<Integer> list = new BoundedList<>(3, Arrays.asList(1, 2, 3));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(2);
+        list.add(3);
+        assertEquals(1, list.get(0));
         assertEquals(2, list.get(1));
+        assertEquals(3, list.get(2));
     }
 
     @Test
     public void testToArray() {
-        BoundedList<Integer> list = new BoundedList<>(3, Arrays.asList(1, 2, 3));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(2);
+        list.add(3);
         assertArrayEquals(new Integer[] {1, 2, 3}, list.toArray());
         assertArrayEquals(new Integer[] {1, 2, 3}, list.toArray(new Integer[3]));
     }
 
     @Test
     public void testAddAll() {
-        ArrayList<String> underlying = new ArrayList<>(Arrays.asList("a", "b", "c"));
-        BoundedList<String> list = new BoundedList<>(5, underlying);
+        BoundedList<String> list = BoundedList.newArrayBacked(5);
+        list.add("a");
+        list.add("b");
+        list.add("c");
         assertEquals("Cannot add another 3 element(s) to the list because it would exceed the " +
             "maximum length of 5",
                 assertThrows(BoundedListTooLongException.class,
-                    () -> list.addAll(Arrays.asList("d", "e", "f"))).
-                        getMessage());
+                        () -> list.addAll(Arrays.asList("d", "e", "f"))).getMessage());
         assertEquals("Cannot add another 3 element(s) to the list because it would exceed the " +
             "maximum length of 5",
                 assertThrows(BoundedListTooLongException.class,
-                        () -> list.addAll(0, Arrays.asList("d", "e", "f"))).
-                        getMessage());
+                        () -> list.addAll(0, Arrays.asList("d", "e", "f"))).getMessage());
         list.addAll(Arrays.asList("d", "e"));
         assertEquals(Arrays.asList("a", "b", "c", "d", "e"), list);
     }
 
     @Test
     public void testIterator() {
-        BoundedList<Integer> list = new BoundedList<>(3, Arrays.asList(1, 2, 3));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(2);
+        list.add(3);
         assertEquals(1, list.iterator().next());
         assertEquals(1, list.listIterator().next());
         assertEquals(3, list.listIterator(2).next());
@@ -165,7 +189,10 @@ public void testIterator() {
 
     @Test
     public void testIteratorIsImmutable() {
-        BoundedList<Integer> list = new BoundedList<>(3, new ArrayList<>(Arrays.asList(1, 2, 3)));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(2);
+        list.add(3);
         assertThrows(UnsupportedOperationException.class,
             () -> list.iterator().remove());
         assertThrows(UnsupportedOperationException.class,
@@ -174,7 +201,10 @@ public void testIteratorIsImmutable() {
 
     @Test
     public void testSubList() {
-        BoundedList<Integer> list = new BoundedList<>(3, new ArrayList<>(Arrays.asList(1, 2, 3)));
+        BoundedList<Integer> list = BoundedList.newArrayBacked(3);
+        list.add(1);
+        list.add(2);
+        list.add(3);
         assertEquals(Arrays.asList(2), list.subList(1, 2));
         assertThrows(UnsupportedOperationException.class,
             () -> list.subList(1, 2).remove(2));

From 7945d322f6b1a61e477e992768668d28f594a072 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 15 Mar 2024 06:19:01 -0700
Subject: [PATCH 161/258] =?UTF-8?q?KAFKA-16267:=20Update=20consumer=5Fgrou?=
 =?UTF-8?q?p=5Fcommand=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20grou?=
 =?UTF-8?q?p=20protocol=20config=20(#15537)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* KAFKA-16267: Update consumer_group_command_test.py to support KIP-848’s group protocol config

Added a new optional group_protocol parameter to the test methods, then passed that down to the setup_consumer method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Note: this requires #15330.

* Update consumer_group_command_test.py

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../tests/core/consumer_group_command_test.py | 29 ++++++++++---------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/tests/kafkatest/tests/core/consumer_group_command_test.py b/tests/kafkatest/tests/core/consumer_group_command_test.py
index 2340677054823..7f1d79574d355 100644
--- a/tests/kafkatest/tests/core/consumer_group_command_test.py
+++ b/tests/kafkatest/tests/core/consumer_group_command_test.py
@@ -20,7 +20,7 @@
 from ducktape.mark.resource import cluster
 
 from kafkatest.services.zookeeper import ZookeeperService
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.console_consumer import ConsoleConsumer
 from kafkatest.services.security.security_config import SecurityConfig
 
@@ -59,14 +59,15 @@ def start_kafka(self, security_protocol, interbroker_security_protocol):
             controller_num_nodes_override=self.num_zk)
         self.kafka.start()
 
-    def start_consumer(self):
+    def start_consumer(self, group_protocol=None):
+        consumer_properties = consumer_group.maybe_set_group_protocol(group_protocol)
         self.consumer = ConsoleConsumer(self.test_context, num_nodes=self.num_brokers, kafka=self.kafka, topic=TOPIC,
-                                        consumer_timeout_ms=None)
+                                        consumer_timeout_ms=None, consumer_properties=consumer_properties)
         self.consumer.start()
 
-    def setup_and_verify(self, security_protocol, group=None):
+    def setup_and_verify(self, security_protocol, group=None, group_protocol=None):
         self.start_kafka(security_protocol, security_protocol)
-        self.start_consumer()
+        self.start_consumer(group_protocol=group_protocol)
         consumer_node = self.consumer.nodes[0]
         wait_until(lambda: self.consumer.alive(consumer_node),
                    timeout_sec=20, backoff_sec=.2, err_msg="Consumer was too slow to start")
@@ -92,35 +93,37 @@ def setup_and_verify(self, security_protocol, group=None):
     @cluster(num_nodes=3)
     @matrix(
         security_protocol=['PLAINTEXT', 'SSL'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         security_protocol=['PLAINTEXT', 'SSL'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_list_consumer_groups(self, security_protocol='PLAINTEXT', metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_list_consumer_groups(self, security_protocol='PLAINTEXT', metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Tests if ConsumerGroupCommand is listing correct consumer groups
         :return: None
         """
-        self.setup_and_verify(security_protocol)
+        self.setup_and_verify(security_protocol, group_protocol=group_protocol)
 
     @cluster(num_nodes=3)
     @matrix(
         security_protocol=['PLAINTEXT', 'SSL'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         security_protocol=['PLAINTEXT', 'SSL'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_describe_consumer_group(self, security_protocol='PLAINTEXT', metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_describe_consumer_group(self, security_protocol='PLAINTEXT', metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Tests if ConsumerGroupCommand is describing a consumer group correctly
         :return: None
         """
-        self.setup_and_verify(security_protocol, group="test-consumer-group")
+        self.setup_and_verify(security_protocol, group="test-consumer-group", group_protocol=group_protocol)

From 8359f947a7a438595a75b9bcfcdd9239587f7205 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 15 Mar 2024 06:19:56 -0700
Subject: [PATCH 162/258] =?UTF-8?q?KAFKA-16268:=20Update=20fetch=5Ffrom=5F?=
 =?UTF-8?q?follower=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20group?=
 =?UTF-8?q?=20protocol=20config=20(#15539)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional `group_protocol` parameter to the test methods, then passed that down to the `setup_consumer` method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new `@matrix` block instead of adding the `group_protocol=["classic", "consumer"]` to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../tests/core/fetch_from_follower_test.py       | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tests/kafkatest/tests/core/fetch_from_follower_test.py b/tests/kafkatest/tests/core/fetch_from_follower_test.py
index 96f3d16cdd1c7..6a096d7b92be8 100644
--- a/tests/kafkatest/tests/core/fetch_from_follower_test.py
+++ b/tests/kafkatest/tests/core/fetch_from_follower_test.py
@@ -20,7 +20,7 @@
 from ducktape.mark.resource import cluster
 
 from kafkatest.services.console_consumer import ConsoleConsumer
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.monitor.jmx import JmxTool
 from kafkatest.services.verifiable_producer import VerifiableProducer
 from kafkatest.services.zookeeper import ZookeeperService
@@ -71,14 +71,15 @@ def setUp(self):
 
     @cluster(num_nodes=9)
     @matrix(
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_consumer_preferred_read_replica(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_consumer_preferred_read_replica(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         This test starts up brokers with "broker.rack" and "replica.selector.class" configurations set. The replica
         selector is set to the rack-aware implementation. One of the brokers has a different rack than the other two.
@@ -98,10 +99,15 @@ def test_consumer_preferred_read_replica(self, metadata_quorum=quorum.zk, use_ne
 
         self.producer = VerifiableProducer(self.test_context, self.num_producers, self.kafka, self.topic,
                                            throughput=self.producer_throughput)
+        consumer_properties = consumer_group.maybe_set_group_protocol(group_protocol,
+                                                                      config={
+                                                                          "client.rack": non_leader_rack,
+                                                                          "metadata.max.age.ms": self.METADATA_MAX_AGE_MS
+                                                                      })
         self.consumer = ConsoleConsumer(self.test_context, self.num_consumers, self.kafka, self.topic,
                                         client_id="console-consumer", group_id="test-consumer-group-1",
                                         consumer_timeout_ms=60000, message_validator=is_int,
-                                        consumer_properties={"client.rack": non_leader_rack, "metadata.max.age.ms": self.METADATA_MAX_AGE_MS})
+                                        consumer_properties=consumer_properties)
 
         # Start up and let some data get produced
         self.start_producer_and_consumer()

From 57557df3eda02e0157c2b10f14cdbf1c7e944814 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 15 Mar 2024 06:20:46 -0700
Subject: [PATCH 163/258] =?UTF-8?q?KAFKA-16269:=20Update=20reassign=5Fpart?=
 =?UTF-8?q?itions=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20group=20p?=
 =?UTF-8?q?rotocol=20config=20(#15540)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional `group_protocol` parameter to the test methods, then passed that down to the `setup_consumer` method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new `@matrix` block instead of adding the `group_protocol=["classic", "consumer"]` to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../kafkatest/tests/core/reassign_partitions_test.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/kafkatest/tests/core/reassign_partitions_test.py b/tests/kafkatest/tests/core/reassign_partitions_test.py
index 0e02c3535e787..b5c83b6a740e0 100644
--- a/tests/kafkatest/tests/core/reassign_partitions_test.py
+++ b/tests/kafkatest/tests/core/reassign_partitions_test.py
@@ -19,7 +19,7 @@
 
 from kafkatest.services.kafka import config_property
 from kafkatest.services.zookeeper import ZookeeperService
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.verifiable_producer import VerifiableProducer
 from kafkatest.services.console_consumer import ConsoleConsumer
 from kafkatest.tests.produce_consume_validate import ProduceConsumeValidateTest
@@ -135,16 +135,17 @@ def move_start_offset(self):
     @matrix(
         bounce_brokers=[True, False],
         reassign_from_offset_zero=[True, False],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         bounce_brokers=[True, False],
         reassign_from_offset_zero=[True, False],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_reassign_partitions(self, bounce_brokers, reassign_from_offset_zero, metadata_quorum, use_new_coordinator=False):
+    def test_reassign_partitions(self, bounce_brokers, reassign_from_offset_zero, metadata_quorum, use_new_coordinator=False, group_protocol=None):
         """Reassign partitions tests.
         Setup: 1 zk, 4 kafka nodes, 1 topic with partitions=20, replication-factor=3,
         and min.insync.replicas=3
@@ -167,7 +168,8 @@ def test_reassign_partitions(self, bounce_brokers, reassign_from_offset_zero, me
         self.consumer = ConsoleConsumer(self.test_context, self.num_consumers,
                                         self.kafka, self.topic,
                                         consumer_timeout_ms=60000,
-                                        message_validator=is_int)
+                                        message_validator=is_int,
+                                        consumer_properties=consumer_group.maybe_set_group_protocol(group_protocol))
 
         self.enable_idempotence=True
         self.run_produce_consume_validate(core_test_action=lambda: self.reassign_partitions(bounce_brokers))

From 386347bbcaa089056b3d854e55d5cac3b0ae24cb Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 15 Mar 2024 06:21:41 -0700
Subject: [PATCH 164/258] =?UTF-8?q?KAFKA-16270:=20Update=20snapshot=5Ftest?=
 =?UTF-8?q?.py=20to=20support=20KIP-848=E2=80=99s=20group=20protocol=20con?=
 =?UTF-8?q?fig=20(#15538)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional `group_protocol` parameter to the test methods, then passed that down to the `setup_consumer` method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new `@matrix` block instead of adding the `group_protocol=["classic", "consumer"]` to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 tests/kafkatest/tests/core/snapshot_test.py | 28 +++++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/tests/kafkatest/tests/core/snapshot_test.py b/tests/kafkatest/tests/core/snapshot_test.py
index 5368b48e5c11f..25ffd4227e312 100644
--- a/tests/kafkatest/tests/core/snapshot_test.py
+++ b/tests/kafkatest/tests/core/snapshot_test.py
@@ -21,6 +21,7 @@
 from kafkatest.services.console_consumer import ConsoleConsumer
 from kafkatest.services.kafka import KafkaService
 from kafkatest.services.kafka import config_property
+from kafkatest.services.kafka import consumer_group
 from kafkatest.services.verifiable_producer import VerifiableProducer
 from kafkatest.tests.produce_consume_validate import ProduceConsumeValidateTest
 from kafkatest.utils import is_int
@@ -125,7 +126,7 @@ def file_exists(self, node, file_path):
             self.logger.debug("File %s was found" % file_path)
             return True
 
-    def validate_success(self, topic = None):
+    def validate_success(self, topic = None, group_protocol=None):
         if topic is None:
             # Create a new topic
             topic = "%s%d" % (TestSnapshots.TOPIC_NAME_PREFIX, self.topics_created)
@@ -138,7 +139,8 @@ def validate_success(self, topic = None):
 
         self.consumer = ConsoleConsumer(self.test_context, self.num_consumers, self.kafka,
                                         topic, consumer_timeout_ms=30000,
-                                        message_validator=is_int)
+                                        message_validator=is_int,
+                                        consumer_properties=consumer_group.maybe_set_group_protocol(group_protocol))
         self.start_producer_and_consumer()
         self.stop_producer_and_consumer()
         self.validate()
@@ -146,9 +148,14 @@ def validate_success(self, topic = None):
     @cluster(num_nodes=9)
     @matrix(
         metadata_quorum=quorum.all_kraft,
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
     )
-    def test_broker(self, metadata_quorum=quorum.combined_kraft, use_new_coordinator=False):
+    @matrix(
+        metadata_quorum=quorum.all_kraft,
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
+    )
+    def test_broker(self, metadata_quorum=quorum.combined_kraft, use_new_coordinator=False, group_protocol=None):
         """ Test the ability of a broker to consume metadata snapshots
         and to recover the cluster metadata state using them
 
@@ -204,14 +211,19 @@ def test_broker(self, metadata_quorum=quorum.combined_kraft, use_new_coordinator
         self.kafka.create_topic(topic_cfg)
 
         # Produce to the newly created topic and make sure it works.
-        self.validate_success(broker_topic)
+        self.validate_success(broker_topic, group_protocol=group_protocol)
 
     @cluster(num_nodes=9)
     @matrix(
         metadata_quorum=quorum.all_kraft,
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        metadata_quorum=quorum.all_kraft,
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_controller(self, metadata_quorum=quorum.combined_kraft, use_new_coordinator=False):
+    def test_controller(self, metadata_quorum=quorum.combined_kraft, use_new_coordinator=False, group_protocol=None):
         """ Test the ability of controllers to consume metadata snapshots
         and to recover the cluster metadata state using them
 
@@ -254,4 +266,4 @@ def test_controller(self, metadata_quorum=quorum.combined_kraft, use_new_coordin
             self.kafka.controller_quorum.start_node(node)
 
         # Produce to a newly created topic and make sure it works.
-        self.validate_success()
+        self.validate_success(group_protocol=group_protocol)

From 3f666026260472179c17ded967a8d9577107e1f0 Mon Sep 17 00:00:00 2001
From: TapDang <89607407+phong260702@users.noreply.github.com>
Date: Fri, 15 Mar 2024 13:41:25 -0400
Subject: [PATCH 165/258] KAFKA-16190: Member should send full heartbeat when
 rejoining (#15401)

When the consumer rejoins, heartbeat request builder make sure that all fields are sent in the heartbeat request.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/HeartbeatRequestManager.java     |  8 +++++---
 .../internals/HeartbeatRequestManagerTest.java | 18 ++++++++++++++++--
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 19e4a8b7132ba..12920d8b2dbcc 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -532,8 +532,10 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
             // InstanceId - set if present
             membershipManager.groupInstanceId().ifPresent(data::setInstanceId);
 
+            boolean sendAllFields = membershipManager.state() == MemberState.JOINING;
+
             // RebalanceTimeoutMs - only sent if has changed since the last heartbeat
-            if (sentFields.rebalanceTimeoutMs != rebalanceTimeoutMs) {
+            if (sendAllFields || sentFields.rebalanceTimeoutMs != rebalanceTimeoutMs) {
                 data.setRebalanceTimeoutMs(rebalanceTimeoutMs);
                 sentFields.rebalanceTimeoutMs = rebalanceTimeoutMs;
             }
@@ -541,7 +543,7 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
             if (!this.subscriptions.hasPatternSubscription()) {
                 // SubscribedTopicNames - only sent if has changed since the last heartbeat
                 TreeSet<String> subscribedTopicNames = new TreeSet<>(this.subscriptions.subscription());
-                if (!subscribedTopicNames.equals(sentFields.subscribedTopicNames)) {
+                if (sendAllFields || !subscribedTopicNames.equals(sentFields.subscribedTopicNames)) {
                     data.setSubscribedTopicNames(new ArrayList<>(this.subscriptions.subscription()));
                     sentFields.subscribedTopicNames = subscribedTopicNames;
                 }
@@ -552,7 +554,7 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
 
             // ServerAssignor - only sent if has changed since the last heartbeat
             this.membershipManager.serverAssignor().ifPresent(serverAssignor -> {
-                if (!serverAssignor.equals(sentFields.serverAssignor)) {
+                if (sendAllFields || !serverAssignor.equals(sentFields.serverAssignor)) {
                     data.setServerAssignor(serverAssignor);
                     sentFields.serverAssignor = serverAssignor;
                 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 8e05e505be471..3409871813391 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -80,6 +80,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+
 public class HeartbeatRequestManagerTest {
     private long retryBackoffMs = DEFAULT_RETRY_BACKOFF_MS;
     private int heartbeatIntervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS;
@@ -428,9 +429,22 @@ public void testHeartbeatState() {
         assertEquals(memberId, data.memberId());
         assertEquals(0, data.memberEpoch());
         assertNull(data.instanceId());
-        assertEquals(-1, data.rebalanceTimeoutMs());
+        assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, data.rebalanceTimeoutMs());
         assertEquals(Collections.singletonList(topic), data.subscribedTopicNames());
-        assertNull(data.serverAssignor());
+        assertEquals(ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR, data.serverAssignor());
+        assertNull(data.topicPartitions());
+        membershipManager.onHeartbeatRequestSent();
+        assertEquals(MemberState.JOINING, membershipManager.state());
+
+        membershipManager.transitionToFenced();
+        data = heartbeatState.buildRequestData();
+        assertEquals(ConsumerTestBuilder.DEFAULT_GROUP_ID, data.groupId());
+        assertEquals(memberId, data.memberId());
+        assertEquals(0, data.memberEpoch());
+        assertNull(data.instanceId());
+        assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, data.rebalanceTimeoutMs());
+        assertEquals(Collections.singletonList(topic), data.subscribedTopicNames());
+        assertEquals(ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR, data.serverAssignor());
         assertNull(data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.JOINING, membershipManager.state());

From 313574e329ef55ce2da3641430add3c4473d621a Mon Sep 17 00:00:00 2001
From: "Matthias J. Sax" <matthias@confluent.io>
Date: Sat, 16 Mar 2024 18:06:45 -0700
Subject: [PATCH 166/258] MINOR: fix flaky EosIntegrationTest (#15494)

Bumping some timeout due to slow Jenkins build.

Reviewers: Bruno Cadonna <bruno@confluent.io>
---
 .../org/apache/kafka/streams/integration/EosIntegrationTest.java | 1 +
 1 file changed, 1 insertion(+)

diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
index 7989de0e76bb7..20199267ecac8 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java
@@ -189,6 +189,7 @@ public static Collection<Object[]> data() {
     public void createTopics() throws Exception {
         applicationId = "appId-" + TEST_NUMBER.getAndIncrement();
         CLUSTER.deleteTopicsAndWait(
+            60_000L,
             SINGLE_PARTITION_INPUT_TOPIC, MULTI_PARTITION_INPUT_TOPIC,
             SINGLE_PARTITION_THROUGH_TOPIC, MULTI_PARTITION_THROUGH_TOPIC,
             SINGLE_PARTITION_OUTPUT_TOPIC, MULTI_PARTITION_OUTPUT_TOPIC);

From 359982328840226fecd9e655197fee07a6bdc538 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Mon, 18 Mar 2024 00:52:08 -0700
Subject: [PATCH 167/258] MINOR: Remove unused client side assignor
 fields/classes (#15545)

In https://github.com/apache/kafka/pull/15364, we introduced, thoughtfully, a non-backward compatible record change for the new consumer group protocol. So it is a good opportunity for cleaning unused fields, mainly related to the client side assignor logic which is not implemented yet. It is better to introduce them when we need them and more importantly when we implement it.

Note that starting from 3.8, we won't make such changes anymore. Non-backward compatible changes are still acceptable now because we clearly said that upgrade won't be supported from the KIP-848 EA.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../coordinator/group/RecordHelpers.java      |  11 +-
 .../group/consumer/Assignment.java            |  68 +-------
 .../group/consumer/ClientAssignor.java        | 164 ------------------
 .../group/consumer/ConsumerGroupMember.java   |  33 ----
 .../group/consumer/VersionedMetadata.java     |  87 ----------
 .../ConsumerGroupMemberMetadataValue.json     |  17 +-
 ...sumerGroupTargetAssignmentMemberValue.json |   8 +-
 .../coordinator/group/RecordHelpersTest.java  |  20 +--
 .../group/consumer/AssignmentTest.java        |  71 +-------
 .../group/consumer/ClientAssignorTest.java    | 133 --------------
 .../consumer/ConsumerGroupMemberTest.java     |  69 +-------
 .../consumer/TargetAssignmentBuilderTest.java |   6 +-
 .../group/consumer/VersionedMetadataTest.java |  59 -------
 13 files changed, 16 insertions(+), 730 deletions(-)
 delete mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ClientAssignor.java
 delete mode 100644 group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadata.java
 delete mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ClientAssignorTest.java
 delete mode 100644 group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadataTest.java

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
index b162100be3745..c2bd093d6a8cb 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
@@ -82,16 +82,7 @@ public static Record newMemberSubscriptionRecord(
                     .setSubscribedTopicNames(member.subscribedTopicNames())
                     .setSubscribedTopicRegex(member.subscribedTopicRegex())
                     .setServerAssignor(member.serverAssignorName().orElse(null))
-                    .setRebalanceTimeoutMs(member.rebalanceTimeoutMs())
-                    .setAssignors(member.clientAssignors().stream().map(assignorState ->
-                        new ConsumerGroupMemberMetadataValue.Assignor()
-                            .setName(assignorState.name())
-                            .setReason(assignorState.reason())
-                            .setMinimumVersion(assignorState.minimumVersion())
-                            .setMaximumVersion(assignorState.maximumVersion())
-                            .setVersion(assignorState.metadata().version())
-                            .setMetadata(assignorState.metadata().metadata().array())
-                    ).collect(Collectors.toList())),
+                    .setRebalanceTimeoutMs(member.rebalanceTimeoutMs()),
                 (short) 0
             )
         );
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/Assignment.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/Assignment.java
index 9cf6b521dd26f..f70601ec0cc1c 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/Assignment.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/Assignment.java
@@ -19,7 +19,6 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberValue;
 
-import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
@@ -31,52 +30,17 @@
  * An immutable assignment for a member.
  */
 public class Assignment {
-    public static final Assignment EMPTY = new Assignment(
-        (byte) 0,
-        Collections.emptyMap(),
-        VersionedMetadata.EMPTY
-    );
-
-    /**
-     * The error assigned to the member.
-     */
-    private final byte error;
+    public static final Assignment EMPTY = new Assignment(Collections.emptyMap());
 
     /**
      * The partitions assigned to the member.
      */
     private final Map<Uuid, Set<Integer>> partitions;
 
-    /**
-     * The metadata assigned to the member.
-     */
-    private final VersionedMetadata metadata;
-
     public Assignment(
         Map<Uuid, Set<Integer>> partitions
     ) {
-        this(
-            (byte) 0,
-            partitions,
-            VersionedMetadata.EMPTY
-        );
-    }
-
-    public Assignment(
-        byte error,
-        Map<Uuid, Set<Integer>> partitions,
-        VersionedMetadata metadata
-    ) {
-        this.error = error;
         this.partitions = Collections.unmodifiableMap(Objects.requireNonNull(partitions));
-        this.metadata = Objects.requireNonNull(metadata);
-    }
-
-    /**
-     * @return The error.
-     */
-    public byte error() {
-        return error;
     }
 
     /**
@@ -86,40 +50,22 @@ public Map<Uuid, Set<Integer>> partitions() {
         return partitions;
     }
 
-    /**
-     * @return The metadata.
-     */
-    public VersionedMetadata metadata() {
-        return metadata;
-    }
-
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         Assignment that = (Assignment) o;
-
-        if (error != that.error) return false;
-        if (!partitions.equals(that.partitions)) return false;
-        return metadata.equals(that.metadata);
+        return partitions.equals(that.partitions);
     }
 
     @Override
     public int hashCode() {
-        int result = error;
-        result = 31 * result + partitions.hashCode();
-        result = 31 * result + metadata.hashCode();
-        return result;
+        return partitions.hashCode();
     }
 
     @Override
     public String toString() {
-        return "Assignment(" +
-            "error=" + error +
-            ", partitions=" + partitions +
-            ", metadata=" + metadata +
-            ')';
+        return "Assignment(partitions=" + partitions + ')';
     }
 
     /**
@@ -132,13 +78,9 @@ public static Assignment fromRecord(
         ConsumerGroupTargetAssignmentMemberValue record
     ) {
         return new Assignment(
-            record.error(),
             record.topicPartitions().stream().collect(Collectors.toMap(
                 ConsumerGroupTargetAssignmentMemberValue.TopicPartition::topicId,
-                topicPartitions -> new HashSet<>(topicPartitions.partitions()))),
-            new VersionedMetadata(
-                record.metadataVersion(),
-                ByteBuffer.wrap(record.metadataBytes()))
+                topicPartitions -> new HashSet<>(topicPartitions.partitions())))
         );
     }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ClientAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ClientAssignor.java
deleted file mode 100644
index fed73cd58d39f..0000000000000
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ClientAssignor.java
+++ /dev/null
@@ -1,164 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.coordinator.group.consumer;
-
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
-
-import java.nio.ByteBuffer;
-import java.util.Objects;
-
-/**
- * An immutable representation of a client side assignor within a consumer group member.
- */
-public class ClientAssignor {
-    /**
-     * The name of the assignor.
-     */
-    private final String name;
-
-    /**
-     * The reason reported by the assignor.
-     */
-    private final byte reason;
-
-    /**
-     * The minimum metadata version supported by the assignor.
-     */
-    private final short minimumVersion;
-
-    /**
-     * The maximum metadata version supported by the assignor.
-     */
-    private final short maximumVersion;
-
-    /**
-     * The versioned metadata.
-     */
-    private final VersionedMetadata metadata;
-
-    public ClientAssignor(
-        String name,
-        byte reason,
-        short minimumVersion,
-        short maximumVersion,
-        VersionedMetadata metadata
-    ) {
-        this.name = Objects.requireNonNull(name);
-        if (name.isEmpty()) {
-            throw new IllegalArgumentException("Assignor name cannot be empty.");
-        }
-        this.reason = reason;
-        this.minimumVersion = minimumVersion;
-        if (minimumVersion < -1) {
-            // -1 is supported as part of the upgrade from the old protocol to the new protocol. It
-            // basically means that the assignor supports metadata from the old client assignor.
-            throw new IllegalArgumentException("Assignor minimum version must be greater than -1.");
-        }
-        this.maximumVersion = maximumVersion;
-        if (maximumVersion < 0) {
-            throw new IllegalArgumentException("Assignor maximum version must be greater than or equals to 0.");
-        } else if (maximumVersion < minimumVersion) {
-            throw new IllegalArgumentException("Assignor maximum version must be greater than or equals to "
-                + "the minimum version.");
-        }
-        this.metadata = Objects.requireNonNull(metadata);
-    }
-
-    /**
-     * @return The client side assignor name.
-     */
-    public String name() {
-        return this.name;
-    }
-
-    /**
-     * @return The current reason reported by the assignor.
-     */
-    public byte reason() {
-        return this.reason;
-    }
-
-    /**
-     * @return The minimum version supported by the assignor.
-     */
-    public short minimumVersion() {
-        return this.minimumVersion;
-    }
-
-    /**
-     * @return The maximum version supported by the assignor.
-     */
-    public short maximumVersion() {
-        return this.maximumVersion;
-    }
-
-    /**
-     * @return The versioned metadata.
-     */
-    public VersionedMetadata metadata() {
-        return metadata;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-
-        ClientAssignor that = (ClientAssignor) o;
-
-        if (reason != that.reason) return false;
-        if (minimumVersion != that.minimumVersion) return false;
-        if (maximumVersion != that.maximumVersion) return false;
-        if (!name.equals(that.name)) return false;
-        return metadata.equals(that.metadata);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = name.hashCode();
-        result = 31 * result + (int) reason;
-        result = 31 * result + (int) minimumVersion;
-        result = 31 * result + (int) maximumVersion;
-        result = 31 * result + metadata.hashCode();
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "ClientAssignor(name=" + name +
-            ", reason=" + reason +
-            ", minimumVersion=" + minimumVersion +
-            ", maximumVersion=" + maximumVersion +
-            ", metadata=" + metadata +
-            ')';
-    }
-
-    public static ClientAssignor fromRecord(
-        ConsumerGroupMemberMetadataValue.Assignor record
-    ) {
-        return new ClientAssignor(
-            record.name(),
-            record.reason(),
-            record.minimumVersion(),
-            record.maximumVersion(),
-            new VersionedMetadata(
-                record.version(),
-                ByteBuffer.wrap(record.metadata())
-            )
-        );
-    }
-}
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
index fdf919cc81cf7..3bf87caef9bcc 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
@@ -62,7 +62,6 @@ public static class Builder {
         private List<String> subscribedTopicNames = Collections.emptyList();
         private String subscribedTopicRegex = "";
         private String serverAssignorName = null;
-        private List<ClientAssignor> clientAssignors = Collections.emptyList();
         private Map<Uuid, Set<Integer>> assignedPartitions = Collections.emptyMap();
         private Map<Uuid, Set<Integer>> partitionsPendingRevocation = Collections.emptyMap();
 
@@ -84,7 +83,6 @@ public Builder(ConsumerGroupMember member) {
             this.subscribedTopicNames = member.subscribedTopicNames;
             this.subscribedTopicRegex = member.subscribedTopicRegex;
             this.serverAssignorName = member.serverAssignorName;
-            this.clientAssignors = member.clientAssignors;
             this.state = member.state;
             this.assignedPartitions = member.assignedPartitions;
             this.partitionsPendingRevocation = member.partitionsPendingRevocation;
@@ -179,16 +177,6 @@ public Builder maybeUpdateServerAssignorName(Optional<String> serverAssignorName
             return this;
         }
 
-        public Builder setClientAssignors(List<ClientAssignor> clientAssignors) {
-            this.clientAssignors = clientAssignors;
-            return this;
-        }
-
-        public Builder maybeUpdateClientAssignors(Optional<List<ClientAssignor>> clientAssignors) {
-            this.clientAssignors = clientAssignors.orElse(this.clientAssignors);
-            return this;
-        }
-
         public Builder setState(MemberState state) {
             this.state = state;
             return this;
@@ -213,9 +201,6 @@ public Builder updateWith(ConsumerGroupMemberMetadataValue record) {
             setSubscribedTopicRegex(record.subscribedTopicRegex());
             setRebalanceTimeoutMs(record.rebalanceTimeoutMs());
             setServerAssignorName(record.serverAssignor());
-            setClientAssignors(record.assignors().stream()
-                .map(ClientAssignor::fromRecord)
-                .collect(Collectors.toList()));
             return this;
         }
 
@@ -249,7 +234,6 @@ public ConsumerGroupMember build() {
                 subscribedTopicNames,
                 subscribedTopicRegex,
                 serverAssignorName,
-                clientAssignors,
                 state,
                 assignedPartitions,
                 partitionsPendingRevocation
@@ -317,11 +301,6 @@ public ConsumerGroupMember build() {
      */
     private final String serverAssignorName;
 
-    /**
-     * The states of the client side assignors of the member.
-     */
-    private final List<ClientAssignor> clientAssignors;
-
     /**
      * The partitions assigned to this member.
      */
@@ -344,7 +323,6 @@ private ConsumerGroupMember(
         List<String> subscribedTopicNames,
         String subscribedTopicRegex,
         String serverAssignorName,
-        List<ClientAssignor> clientAssignors,
         MemberState state,
         Map<Uuid, Set<Integer>> assignedPartitions,
         Map<Uuid, Set<Integer>> partitionsPendingRevocation
@@ -361,7 +339,6 @@ private ConsumerGroupMember(
         this.subscribedTopicNames = subscribedTopicNames;
         this.subscribedTopicRegex = subscribedTopicRegex;
         this.serverAssignorName = serverAssignorName;
-        this.clientAssignors = clientAssignors;
         this.assignedPartitions = assignedPartitions;
         this.partitionsPendingRevocation = partitionsPendingRevocation;
     }
@@ -443,13 +420,6 @@ public Optional<String> serverAssignorName() {
         return Optional.ofNullable(serverAssignorName);
     }
 
-    /**
-     * @return The list of client side assignors.
-     */
-    public List<ClientAssignor> clientAssignors() {
-        return clientAssignors;
-    }
-
     /**
      * @return The current state.
      */
@@ -551,7 +521,6 @@ public boolean equals(Object o) {
             && Objects.equals(subscribedTopicNames, that.subscribedTopicNames)
             && Objects.equals(subscribedTopicRegex, that.subscribedTopicRegex)
             && Objects.equals(serverAssignorName, that.serverAssignorName)
-            && Objects.equals(clientAssignors, that.clientAssignors)
             && Objects.equals(assignedPartitions, that.assignedPartitions)
             && Objects.equals(partitionsPendingRevocation, that.partitionsPendingRevocation);
     }
@@ -570,7 +539,6 @@ public int hashCode() {
         result = 31 * result + Objects.hashCode(subscribedTopicNames);
         result = 31 * result + Objects.hashCode(subscribedTopicRegex);
         result = 31 * result + Objects.hashCode(serverAssignorName);
-        result = 31 * result + Objects.hashCode(clientAssignors);
         result = 31 * result + Objects.hashCode(assignedPartitions);
         result = 31 * result + Objects.hashCode(partitionsPendingRevocation);
         return result;
@@ -591,7 +559,6 @@ public String toString() {
             ", subscribedTopicNames=" + subscribedTopicNames +
             ", subscribedTopicRegex='" + subscribedTopicRegex + '\'' +
             ", serverAssignorName='" + serverAssignorName + '\'' +
-            ", clientAssignors=" + clientAssignors +
             ", assignedPartitions=" + assignedPartitions +
             ", partitionsPendingRevocation=" + partitionsPendingRevocation +
             ')';
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadata.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadata.java
deleted file mode 100644
index 89baeb82bfa6d..0000000000000
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadata.java
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.coordinator.group.consumer;
-
-import java.nio.ByteBuffer;
-import java.util.Objects;
-
-/**
- * Immutable versioned metadata. It contains a bunch of bytes tagged with a version. The
- * format of the bytes is unspecified. This is mainly used by client side assignors to
- * exchange arbitrary metadata between the members and the assignor and vice versa.
- */
-public class VersionedMetadata {
-    public static final VersionedMetadata EMPTY = new VersionedMetadata((short) 0, ByteBuffer.allocate(0));
-
-    /**
-     * The version of the metadata encoded in {{@link VersionedMetadata#metadata}}.
-     */
-    private final short version;
-
-    /**
-     * The metadata bytes.
-     */
-    private final ByteBuffer metadata;
-
-    public VersionedMetadata(
-        short version,
-        ByteBuffer metadata
-    ) {
-        this.version = version;
-        this.metadata = Objects.requireNonNull(metadata);
-    }
-
-    /**
-     * @return The version of the metadata.
-     */
-    public short version() {
-        return this.version;
-    }
-
-    /**
-     * @return The ByteBuffer holding the metadata.
-     */
-    public ByteBuffer metadata() {
-        return this.metadata;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-
-        VersionedMetadata that = (VersionedMetadata) o;
-
-        if (version != that.version) return false;
-        return metadata.equals(that.metadata);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = version;
-        result = 31 * result + metadata.hashCode();
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "VersionedMetadata(" +
-            "version=" + version +
-            ", metadata=" + metadata +
-            ')';
-    }
-}
diff --git a/group-coordinator/src/main/resources/common/message/ConsumerGroupMemberMetadataValue.json b/group-coordinator/src/main/resources/common/message/ConsumerGroupMemberMetadataValue.json
index 3b83f95985835..4886b7ad9d320 100644
--- a/group-coordinator/src/main/resources/common/message/ConsumerGroupMemberMetadataValue.json
+++ b/group-coordinator/src/main/resources/common/message/ConsumerGroupMemberMetadataValue.json
@@ -35,21 +35,6 @@
     { "name": "RebalanceTimeoutMs", "type": "int32", "versions": "0+", "default": -1,
       "about": "The rebalance timeout" },
     { "name": "ServerAssignor", "versions": "0+", "nullableVersions": "0+", "type": "string",
-      "about": "The server assignor to use; or null if not used." },
-    { "name": "Assignors", "versions": "0+", "type": "[]Assignor",
-      "about": "The list of assignors.", "fields": [
-      { "name": "Name", "versions": "0+", "type": "string",
-        "about": "The assignor name." },
-      { "name": "MinimumVersion", "versions": "0+", "type": "int16",
-        "about": "The minimum version supported by the assignor." },
-      { "name": "MaximumVersion", "versions": "0+", "type": "int16",
-        "about": "The maximum version supported by the assignor." },
-      { "name": "Reason", "versions": "0+", "type": "int8",
-        "about": "The reason reported by the assignor." },
-      { "name": "Version", "versions": "0+", "type": "int16",
-        "about": "The version used to serialize the metadata." },
-      { "name": "Metadata", "versions": "0+", "type": "bytes",
-        "about": "The metadata." }
-    ]}
+      "about": "The server assignor to use; or null if not used." }
   ]
 }
diff --git a/group-coordinator/src/main/resources/common/message/ConsumerGroupTargetAssignmentMemberValue.json b/group-coordinator/src/main/resources/common/message/ConsumerGroupTargetAssignmentMemberValue.json
index d056eab6d8511..e05c28928e878 100644
--- a/group-coordinator/src/main/resources/common/message/ConsumerGroupTargetAssignmentMemberValue.json
+++ b/group-coordinator/src/main/resources/common/message/ConsumerGroupTargetAssignmentMemberValue.json
@@ -20,16 +20,10 @@
   "validVersions": "0",
   "flexibleVersions": "0+",
   "fields": [
-    { "name": "Error", "versions": "0+", "type": "int8",
-      "about": "The assigned error."},
     { "name": "TopicPartitions", "versions": "0+", "type": "[]TopicPartition",
       "about": "The assigned partitions.", "fields": [
       { "name": "TopicId", "versions": "0+", "type": "uuid" },
       { "name": "Partitions", "versions": "0+", "type": "[]int32" }
-    ]},
-    { "name": "MetadataVersion", "versions": "0+", "type": "int16",
-      "about": "The version of the assigned metadata." },
-    { "name": "MetadataBytes", "versions": "0+", "type": "bytes",
-      "about": "The assigned metadata." }
+    ]}
   ]
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
index 50bfcbd52356c..17d24959a1d66 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
@@ -22,11 +22,9 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.coordinator.group.consumer.ClientAssignor;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
-import org.apache.kafka.coordinator.group.consumer.VersionedMetadata;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
@@ -55,8 +53,6 @@
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -105,14 +101,6 @@ public void testNewMemberSubscriptionRecord() {
             .setSubscribedTopicNames(Arrays.asList("foo", "zar", "bar"))
             .setSubscribedTopicRegex("regex")
             .setServerAssignorName("range")
-            .setClientAssignors(Collections.singletonList(new ClientAssignor(
-                "assignor",
-                (byte) 0,
-                (byte) 1,
-                (byte) 10,
-                new VersionedMetadata(
-                    (byte) 5,
-                    ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))))))
             .build();
 
         Record expectedRecord = new Record(
@@ -130,13 +118,7 @@ public void testNewMemberSubscriptionRecord() {
                     .setClientHost("client-host")
                     .setSubscribedTopicNames(Arrays.asList("bar", "foo", "zar"))
                     .setSubscribedTopicRegex("regex")
-                    .setServerAssignor("range")
-                    .setAssignors(Collections.singletonList(new ConsumerGroupMemberMetadataValue.Assignor()
-                        .setName("assignor")
-                        .setMinimumVersion((short) 1)
-                        .setMaximumVersion((short) 10)
-                        .setVersion((short) 5)
-                        .setMetadata("hello".getBytes(StandardCharsets.UTF_8)))),
+                    .setServerAssignor("range"),
                 (short) 0));
 
         assertEquals(expectedRecord, newMemberSubscriptionRecord(
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java
index 2cff004ded22f..536ec71cacb7f 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/AssignmentTest.java
@@ -20,11 +20,8 @@
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberValue;
 import org.junit.jupiter.api.Test;
 
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -32,29 +29,13 @@
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class AssignmentTest {
 
     @Test
-    public void testPartitionsAndMetadataCannotBeNull() {
-        assertThrows(NullPointerException.class, () -> new Assignment(
-            (byte) 1,
-            null,
-            new VersionedMetadata(
-                (short) 1,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        ));
-
-        assertThrows(NullPointerException.class, () -> new Assignment(
-            (byte) 1,
-            mkAssignment(
-                mkTopicAssignment(Uuid.randomUuid(), 1, 2, 3)
-            ),
-            null
-        ));
+    public void testPartitionsCannotBeNull() {
+        assertThrows(NullPointerException.class, () -> new Assignment(null));
     }
 
     @Test
@@ -62,21 +43,8 @@ public void testAttributes() {
         Map<Uuid, Set<Integer>> partitions = mkAssignment(
             mkTopicAssignment(Uuid.randomUuid(), 1, 2, 3)
         );
-
-        VersionedMetadata metadata = new VersionedMetadata(
-            (short) 1,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        );
-
-        Assignment assignment = new Assignment(
-            (byte) 1,
-            partitions,
-            metadata
-        );
-
-        assertEquals((byte) 1, assignment.error());
+        Assignment assignment = new Assignment(partitions);
         assertEquals(partitions, assignment.partitions());
-        assertEquals(metadata, assignment.metadata());
     }
 
     @Test
@@ -93,22 +61,14 @@ public void testFromTargetAssignmentRecord() {
             .setPartitions(Arrays.asList(4, 5, 6)));
 
         ConsumerGroupTargetAssignmentMemberValue record = new ConsumerGroupTargetAssignmentMemberValue()
-            .setError((byte) 1)
-            .setTopicPartitions(partitions)
-            .setMetadataVersion((short) 2)
-            .setMetadataBytes("foo".getBytes(StandardCharsets.UTF_8));
+            .setTopicPartitions(partitions);
 
         Assignment assignment = Assignment.fromRecord(record);
 
-        assertEquals((short) 1, assignment.error());
         assertEquals(mkAssignment(
             mkTopicAssignment(topicId1, 1, 2, 3),
             mkTopicAssignment(topicId2, 4, 5, 6)
         ), assignment.partitions());
-        assertEquals(new VersionedMetadata(
-            (short) 2,
-            ByteBuffer.wrap("foo".getBytes(StandardCharsets.UTF_8))
-        ), assignment.metadata());
     }
 
     @Test
@@ -117,27 +77,6 @@ public void testEquals() {
             mkTopicAssignment(Uuid.randomUuid(), 1, 2, 3)
         );
 
-        VersionedMetadata metadata = new VersionedMetadata(
-            (short) 1,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        );
-
-        Assignment assignment = new Assignment(
-            (byte) 1,
-            partitions,
-            metadata
-        );
-
-        assertEquals(new Assignment(
-            (byte) 1,
-            partitions,
-            metadata
-        ), assignment);
-
-        assertNotEquals(new Assignment(
-            (byte) 1,
-            Collections.emptyMap(),
-            metadata
-        ), assignment);
+        assertEquals(new Assignment(partitions), new Assignment(partitions));
     }
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ClientAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ClientAssignorTest.java
deleted file mode 100644
index d136a411cd8ec..0000000000000
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ClientAssignorTest.java
+++ /dev/null
@@ -1,133 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.coordinator.group.consumer;
-
-import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
-import org.junit.jupiter.api.Test;
-
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-public class ClientAssignorTest {
-
-    @Test
-    public void testNameAndMetadataCannotBeNull() {
-        assertThrows(NullPointerException.class, () -> new ClientAssignor(
-            "range",
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            null
-        ));
-
-        assertThrows(NullPointerException.class, () -> new ClientAssignor(
-            null,
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            new VersionedMetadata(
-                (short) 8,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        ));
-    }
-
-    @Test
-    public void testAttributes() {
-        ClientAssignor clientAssignor = new ClientAssignor(
-            "range",
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            new VersionedMetadata(
-                (short) 8,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        );
-
-        assertEquals("range", clientAssignor.name());
-        assertEquals((byte) 2, clientAssignor.reason());
-        assertEquals((short) 5, clientAssignor.minimumVersion());
-        assertEquals((short) 10, clientAssignor.maximumVersion());
-        assertEquals(new VersionedMetadata(
-            (short) 8,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        ), clientAssignor.metadata());
-    }
-
-    @Test
-    public void testFromRecord() {
-        ConsumerGroupMemberMetadataValue.Assignor record = new ConsumerGroupMemberMetadataValue.Assignor()
-            .setName("range")
-            .setReason((byte) 2)
-            .setMinimumVersion((byte) 5)
-            .setMaximumVersion((byte) 10)
-            .setVersion((byte) 8)
-            .setMetadata("hello".getBytes(StandardCharsets.UTF_8));
-
-        ClientAssignor clientAssignor = ClientAssignor.fromRecord(record);
-
-        assertEquals("range", clientAssignor.name());
-        assertEquals((byte) 2, clientAssignor.reason());
-        assertEquals((short) 5, clientAssignor.minimumVersion());
-        assertEquals((short) 10, clientAssignor.maximumVersion());
-        assertEquals(new VersionedMetadata(
-            (short) 8,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        ), clientAssignor.metadata());
-    }
-
-    @Test
-    public void testEquals() {
-        ClientAssignor clientAssignor = new ClientAssignor(
-            "range",
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            new VersionedMetadata(
-                (short) 8,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        );
-
-        assertEquals(new ClientAssignor(
-            "range",
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            new VersionedMetadata(
-                (short) 8,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        ), clientAssignor);
-
-        assertNotEquals(new ClientAssignor(
-            "uniform",
-            (byte) 2,
-            (short) 5,
-            (short) 10,
-            new VersionedMetadata(
-                (short) 8,
-                ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-            )
-        ), clientAssignor);
-    }
-}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
index 3baa6f77af4ee..cfb9a79a2d6d3 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMemberTest.java
@@ -24,7 +24,6 @@
 import org.apache.kafka.image.MetadataImage;
 import org.junit.jupiter.api.Test;
 
-import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -40,6 +39,7 @@
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+
 public class ConsumerGroupMemberTest {
 
     @Test
@@ -59,15 +59,6 @@ public void testNewMember() {
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setSubscribedTopicRegex("regex")
             .setServerAssignorName("range")
-            .setClientAssignors(Collections.singletonList(
-                new ClientAssignor(
-                    "assignor",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 1,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))))
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
@@ -85,17 +76,6 @@ public void testNewMember() {
         assertEquals(Arrays.asList("bar", "foo"), member.subscribedTopicNames());
         assertEquals("regex", member.subscribedTopicRegex());
         assertEquals("range", member.serverAssignorName().get());
-        assertEquals(
-            Collections.singletonList(
-                new ClientAssignor(
-                    "assignor",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 1,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))),
-            member.clientAssignors());
         assertEquals(mkAssignment(mkTopicAssignment(topicId1, 1, 2, 3)), member.assignedPartitions());
         assertEquals(mkAssignment(mkTopicAssignment(topicId2, 4, 5, 6)), member.partitionsPendingRevocation());
     }
@@ -117,15 +97,6 @@ public void testEquals() {
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setSubscribedTopicRegex("regex")
             .setServerAssignorName("range")
-            .setClientAssignors(Collections.singletonList(
-                new ClientAssignor(
-                    "assignor",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 1,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))))
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
@@ -143,15 +114,6 @@ public void testEquals() {
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setSubscribedTopicRegex("regex")
             .setServerAssignorName("range")
-            .setClientAssignors(Collections.singletonList(
-                new ClientAssignor(
-                    "assignor",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 1,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))))
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
@@ -178,15 +140,6 @@ public void testUpdateMember() {
             .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
             .setSubscribedTopicRegex("regex")
             .setServerAssignorName("range")
-            .setClientAssignors(Collections.singletonList(
-                new ClientAssignor(
-                    "assignor",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 1,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))))
             .setAssignedPartitions(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3)))
             .setPartitionsPendingRevocation(mkAssignment(
@@ -201,7 +154,6 @@ public void testUpdateMember() {
             .maybeUpdateSubscribedTopicNames(Optional.empty())
             .maybeUpdateSubscribedTopicRegex(Optional.empty())
             .maybeUpdateRebalanceTimeoutMs(OptionalInt.empty())
-            .maybeUpdateClientAssignors(Optional.empty())
             .build();
 
         assertEquals(member, updatedMember);
@@ -213,7 +165,6 @@ public void testUpdateMember() {
             .maybeUpdateSubscribedTopicNames(Optional.of(Arrays.asList("zar")))
             .maybeUpdateSubscribedTopicRegex(Optional.of("new-regex"))
             .maybeUpdateRebalanceTimeoutMs(OptionalInt.of(6000))
-            .maybeUpdateClientAssignors(Optional.of(Collections.emptyList()))
             .build();
 
         assertEquals("new-instance-id", updatedMember.instanceId());
@@ -222,18 +173,11 @@ public void testUpdateMember() {
         assertEquals(Arrays.asList("zar"), updatedMember.subscribedTopicNames());
         assertEquals("new-regex", updatedMember.subscribedTopicRegex());
         assertEquals("new-assignor", updatedMember.serverAssignorName().get());
-        assertEquals(Collections.emptyList(), updatedMember.clientAssignors());
     }
 
     @Test
     public void testUpdateWithConsumerGroupMemberMetadataValue() {
         ConsumerGroupMemberMetadataValue record = new ConsumerGroupMemberMetadataValue()
-            .setAssignors(Collections.singletonList(new ConsumerGroupMemberMetadataValue.Assignor()
-                .setName("client")
-                .setMinimumVersion((short) 0)
-                .setMaximumVersion((short) 2)
-                .setVersion((short) 1)
-                .setMetadata(new byte[0])))
             .setServerAssignor("range")
             .setClientId("client-id")
             .setClientHost("host-id")
@@ -255,17 +199,6 @@ public void testUpdateWithConsumerGroupMemberMetadataValue() {
         assertEquals(Arrays.asList("bar", "foo"), member.subscribedTopicNames());
         assertEquals("regex", member.subscribedTopicRegex());
         assertEquals("range", member.serverAssignorName().get());
-        assertEquals(
-            Collections.singletonList(
-                new ClientAssignor(
-                    "client",
-                    (byte) 0,
-                    (byte) 0,
-                    (byte) 2,
-                    new VersionedMetadata(
-                        (byte) 1,
-                        ByteBuffer.allocate(0)))),
-            member.clientAssignors());
     }
 
     @Test
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
index af94fede50257..2a73e6fb54a9d 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/TargetAssignmentBuilderTest.java
@@ -90,11 +90,7 @@ private void addGroupMember(
                 staticMembers.put(instanceId, memberId);
             }
             members.put(memberId, memberBuilder.build());
-            targetAssignment.put(memberId, new Assignment(
-                (byte) 0,
-                targetPartitions,
-                VersionedMetadata.EMPTY
-            ));
+            targetAssignment.put(memberId, new Assignment(targetPartitions));
         }
 
         public Uuid addTopicMetadata(
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadataTest.java
deleted file mode 100644
index 10997d015cb63..0000000000000
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/VersionedMetadataTest.java
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.coordinator.group.consumer;
-
-import org.junit.jupiter.api.Test;
-
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-public class VersionedMetadataTest {
-    @Test
-    public void testAttributes() {
-        VersionedMetadata metadata = new VersionedMetadata(
-            (short) 1,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        );
-
-        assertEquals((short) 1, metadata.version());
-        assertEquals(ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8)), metadata.metadata());
-    }
-
-    @Test
-    public void testMetadataCannotBeNull() {
-        assertThrows(NullPointerException.class, () -> new VersionedMetadata((short) 1, null));
-    }
-
-    @Test
-    public void testEquals() {
-        VersionedMetadata metadata = new VersionedMetadata(
-            (short) 1,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        );
-
-        assertEquals(new VersionedMetadata(
-            (short) 1,
-            ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))
-        ), metadata);
-
-        assertNotEquals(VersionedMetadata.EMPTY, metadata);
-    }
-}

From e9c50b1f4b0e15b5b016256fb8e31ce704a8c166 Mon Sep 17 00:00:00 2001
From: Edoardo Comar <ecomar@uk.ibm.com>
Date: Mon, 18 Mar 2024 10:14:43 +0000
Subject: [PATCH 168/258] KAFKA-16369: Broker may not shut down when
 SocketServer fails to bind as Address already in use (#15530)

* KAFKA-16369: wait on enableRequestProcessingFuture

Add a Wait in in KafkaServer (ZK mode) for all the SocketServer ports
to be open, and the Acceptors to be started

The BrokerServer (KRaft mode) had such a wait,
which was missing from the KafkaServer (ZK mode).

Add unit test.
---
 .../main/scala/kafka/server/KafkaServer.scala |  9 +++++-
 .../unit/kafka/server/KafkaServerTest.scala   | 30 +++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index bea0a252fb77d..55c773b0a46db 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -617,7 +617,7 @@ class KafkaServer(
             }
           }
         }
-        socketServer.enableRequestProcessing(authorizerFutures)
+        val enableRequestProcessingFuture = socketServer.enableRequestProcessing(authorizerFutures)
         // Block here until all the authorizer futures are complete
         try {
           CompletableFuture.allOf(authorizerFutures.values.toSeq: _*).join()
@@ -625,6 +625,13 @@ class KafkaServer(
           case t: Throwable => throw new RuntimeException("Received a fatal error while " +
             "waiting for all of the authorizer futures to be completed.", t)
         }
+        // Wait for all the SocketServer ports to be open, and the Acceptors to be started.
+        try {
+          enableRequestProcessingFuture.join()
+        } catch {
+          case t: Throwable => throw new RuntimeException("Received a fatal error while " +
+            "waiting for the SocketServer Acceptors to be started.", t)
+        }
 
         _brokerState = BrokerState.RUNNING
         shutdownLatch = new CountDownLatch(1)
diff --git a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
index 1372690afeb25..bf6312f86319a 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
@@ -17,12 +17,13 @@
 
 package kafka.server
 
-import kafka.utils.TestUtils
+import kafka.utils.{CoreUtils, TestUtils}
 import org.apache.kafka.common.security.JaasUtils
 import org.junit.jupiter.api.Assertions.{assertEquals, assertNull, assertThrows, fail}
 import org.junit.jupiter.api.Test
-import java.util.Properties
 
+import java.util.Properties
+import java.net.{InetAddress, ServerSocket}
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 
@@ -42,6 +43,24 @@ class KafkaServerTest extends QuorumTestHarness {
     TestUtils.shutdownServers(Seq(server1, server2))
   }
 
+  @Test
+  def testListenerPortAlreadyInUse(): Unit = {
+    val serverSocket = new ServerSocket(0, 0, InetAddress.getLoopbackAddress)
+
+    var kafkaServer : Option[KafkaServer] = None
+    try {
+      TestUtils.waitUntilTrue(() => serverSocket.isBound, "Server socket failed to bind.")
+      // start a server with listener on the port already bound
+      assertThrows(classOf[RuntimeException],
+        () => kafkaServer = Option(createServerWithListenerOnPort(serverSocket.getLocalPort)),
+        "Expected RuntimeException due to address already in use during KafkaServer startup"
+      )
+    } finally {
+      CoreUtils.swallow(serverSocket.close(), this);
+      TestUtils.shutdownServers(kafkaServer.toList)
+    }
+  }
+
   @Test
   def testCreatesProperZkConfigWhenSaslDisabled(): Unit = {
     val props = new Properties
@@ -161,4 +180,11 @@ class KafkaServerTest extends QuorumTestHarness {
     TestUtils.createServer(kafkaConfig)
   }
 
+  def createServerWithListenerOnPort(port: Int): KafkaServer = {
+    val props = TestUtils.createBrokerConfig(0, zkConnect)
+    props.put(KafkaConfig.ListenersProp, s"PLAINTEXT://localhost:$port")
+    val kafkaConfig = KafkaConfig.fromProps(props)
+    TestUtils.createServer(kafkaConfig)
+  }
+
 }

From 5c929874b88b3b96f650de0f733d93d42ac535a4 Mon Sep 17 00:00:00 2001
From: Lucas Brutschy <lbrutschy@confluent.io>
Date: Mon, 18 Mar 2024 11:52:23 +0100
Subject: [PATCH 169/258] KAFKA-16312, KAFKA-16185: Local epochs in
 reconciliation (#15511)

The goal of this commit is to change the following internals of the reconciliation:

- Introduce a "local epoch" to the local target assignment. When a new target is received by the server, we compare it with the current value. If it is the same, no change. Otherwise, we bump the local epoch and store the new target assignment. Then, on the reconciliation, we also store the epoch in the reconciled assignment and keep using target != current to trigger the reconciliation.
- When we are not in a group (we have not received an assignment), we use null to represent the local target assignment instead of an empty list, to avoid confusions with an empty assignment received by the server. Similarly, we use null to represent the current assignment, when we haven't reconciled the assignment yet.
We also carry the new epoch into the request builder to ensure that we report the owned partitions for the last local epoch.
- To address KAFKA-16312 (call onPartitionsAssigned on empty assignments after joining), we apply the initial assignment returned by the group coordinator (whether empty or not) as a normal reconciliation. This avoids introducing another code path to trigger rebalance listeners - reconciliation is the only way to transition to STABLE. The unneeded parts of reconciliation (autocommit, revocation) will be skipped in the existing. Since a lot of unit tests assumed that not reconciliation behavior is invoked when joining the group with an empty assignment, this required a lot of the changes in the unit tests.

Reviewers: Lianet Magrans <lianetmr@gmail.com>, David Jacot <djacot@confluent.io>
---
 .../internals/HeartbeatRequestManager.java    |  33 +--
 .../consumer/internals/MembershipManager.java |  91 +++++-
 .../internals/MembershipManagerImpl.java      | 122 ++++----
 .../HeartbeatRequestManagerTest.java          | 149 ++++++----
 .../internals/MembershipManagerImplTest.java  | 280 ++++++++++++++----
 5 files changed, 485 insertions(+), 190 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
index 12920d8b2dbcc..5b54e8a4369ba 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.internals.MembershipManager.LocalAssignment;
 import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
@@ -534,25 +535,25 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
 
             boolean sendAllFields = membershipManager.state() == MemberState.JOINING;
 
-            // RebalanceTimeoutMs - only sent if has changed since the last heartbeat
+            // RebalanceTimeoutMs - only sent when joining or if it has changed since the last heartbeat
             if (sendAllFields || sentFields.rebalanceTimeoutMs != rebalanceTimeoutMs) {
                 data.setRebalanceTimeoutMs(rebalanceTimeoutMs);
                 sentFields.rebalanceTimeoutMs = rebalanceTimeoutMs;
             }
 
             if (!this.subscriptions.hasPatternSubscription()) {
-                // SubscribedTopicNames - only sent if has changed since the last heartbeat
+                // SubscribedTopicNames - only sent when joining or if it has changed since the last heartbeat
                 TreeSet<String> subscribedTopicNames = new TreeSet<>(this.subscriptions.subscription());
                 if (sendAllFields || !subscribedTopicNames.equals(sentFields.subscribedTopicNames)) {
                     data.setSubscribedTopicNames(new ArrayList<>(this.subscriptions.subscription()));
                     sentFields.subscribedTopicNames = subscribedTopicNames;
                 }
             } else {
-                // SubscribedTopicRegex - only sent if has changed since the last heartbeat
+                // SubscribedTopicRegex - only sent if it has changed since the last heartbeat
                 //                      - not supported yet
             }
 
-            // ServerAssignor - only sent if has changed since the last heartbeat
+            // ServerAssignor - sent when joining or if it has changed since the last heartbeat
             this.membershipManager.serverAssignor().ifPresent(serverAssignor -> {
                 if (sendAllFields || !serverAssignor.equals(sentFields.serverAssignor)) {
                     data.setServerAssignor(serverAssignor);
@@ -562,18 +563,16 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
 
             // ClientAssignors - not supported yet
 
-            // TopicPartitions - only sent if it has changed since the last heartbeat. Note that
-            // the string consists of just the topic ID and the partitions. When an assignment is
-            // received, we might not yet know the topic name, and then it is learnt subsequently
-            // by a metadata update.
-            TreeSet<String> assignedPartitions = membershipManager.currentAssignment().entrySet().stream()
-                .map(entry -> entry.getKey() + "-" + entry.getValue())
-                .collect(Collectors.toCollection(TreeSet::new));
-            if (!assignedPartitions.equals(sentFields.topicPartitions)) {
+            // TopicPartitions - sent when joining or with the first heartbeat after a new assignment from
+            // the server was reconciled. This is ensured by resending the topic partitions whenever the
+            // local assignment, including its local epoch is changed (although the local epoch is not sent
+            // in the heartbeat).
+            LocalAssignment local = membershipManager.currentAssignment();
+            if (sendAllFields || !local.equals(sentFields.localAssignment)) {
                 List<ConsumerGroupHeartbeatRequestData.TopicPartitions> topicPartitions =
-                    buildTopicPartitionsList(membershipManager.currentAssignment());
+                    buildTopicPartitionsList(local.partitions);
                 data.setTopicPartitions(topicPartitions);
-                sentFields.topicPartitions = assignedPartitions;
+                sentFields.localAssignment = local;
             }
 
             return data;
@@ -592,15 +591,15 @@ static class SentFields {
             private int rebalanceTimeoutMs = -1;
             private TreeSet<String> subscribedTopicNames = null;
             private String serverAssignor = null;
-            private TreeSet<String> topicPartitions = null;
+            private LocalAssignment localAssignment = null;
 
             SentFields() {}
 
             void reset() {
-                rebalanceTimeoutMs = -1;
                 subscribedTopicNames = null;
+                rebalanceTimeoutMs = -1;
                 serverAssignor = null;
-                topicPartitions = null;
+                localAssignment = null;
             }
         }
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
index f0a641d140fb2..c33ac396d237a 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManager.java
@@ -18,12 +18,17 @@
 
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
+import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
 
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.SortedSet;
+import java.util.TreeSet;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -99,7 +104,7 @@ public interface MembershipManager extends RequestManager {
      * @return Current assignment for the member as received from the broker (topic IDs and
      * partitions). This is the last assignment that the member has successfully reconciled.
      */
-    Map<Uuid, SortedSet<Integer>> currentAssignment();
+    LocalAssignment currentAssignment();
 
     /**
      * Transition the member to the FENCED state, where the member will release the assignment by
@@ -185,4 +190,88 @@ public interface MembershipManager extends RequestManager {
      * releasing its assignment. This is expected to be used when the poll timer is reset.
      */
     void maybeRejoinStaleMember();
+
+    /**
+     * A data structure to represent the current assignment, and current target assignment of a member in a consumer group.
+     *
+     * Besides the assigned partitions, it contains a local epoch that is bumped whenever the assignment changes, to ensure
+     * that two assignments with the same partitions but different local epochs are not considered equal.
+     */
+    final class LocalAssignment {
+
+        public static final long NONE_EPOCH = -1;
+
+        public static final LocalAssignment NONE = new LocalAssignment(NONE_EPOCH, Collections.emptyMap());
+
+        public final long localEpoch;
+
+        public final Map<Uuid, SortedSet<Integer>> partitions;
+
+        public LocalAssignment(long localEpoch, Map<Uuid, SortedSet<Integer>> partitions) {
+            this.localEpoch = localEpoch;
+            this.partitions = partitions;
+            if (localEpoch == NONE_EPOCH && !partitions.isEmpty()) {
+                throw new IllegalArgumentException("Local epoch must be set if there are partitions");
+            }
+        }
+
+        public LocalAssignment(long localEpoch, SortedSet<TopicIdPartition> topicIdPartitions) {
+            this.localEpoch = localEpoch;
+            this.partitions = new HashMap<>();
+            if (localEpoch == NONE_EPOCH && !topicIdPartitions.isEmpty()) {
+                throw new IllegalArgumentException("Local epoch must be set if there are partitions");
+            }
+            topicIdPartitions.forEach(topicIdPartition -> {
+                Uuid topicId = topicIdPartition.topicId();
+                partitions.computeIfAbsent(topicId, k -> new TreeSet<>()).add(topicIdPartition.partition());
+            });
+        }
+
+        public String toString() {
+            return "LocalAssignment{" +
+                "localEpoch=" + localEpoch +
+                ", partitions=" + partitions +
+                '}';
+        }
+
+        public boolean equals(final Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            final LocalAssignment that = (LocalAssignment) o;
+            return localEpoch == that.localEpoch && Objects.equals(partitions, that.partitions);
+        }
+
+        public int hashCode() {
+            return Objects.hash(localEpoch, partitions);
+        }
+
+        public boolean isNone() {
+            return localEpoch == NONE_EPOCH;
+        }
+
+        Optional<LocalAssignment> updateWith(ConsumerGroupHeartbeatResponseData.Assignment assignment) {
+            // Return if we have an assignment, and it is the same as current assignment; comparison without creating a new collection
+            if (localEpoch != NONE_EPOCH) {
+                if (partitions.size() == assignment.topicPartitions().size() &&
+                    assignment.topicPartitions().stream().allMatch(
+                        tp -> partitions.containsKey(tp.topicId()) &&
+                            partitions.get(tp.topicId()).size() == tp.partitions().size() &&
+                            partitions.get(tp.topicId()).containsAll(tp.partitions()))) {
+                    return Optional.empty();
+                }
+            }
+
+            // Bump local epoch and replace assignment
+            long nextLocalEpoch = localEpoch + 1;
+            HashMap<Uuid, SortedSet<Integer>> partitions = new HashMap<>();
+            assignment.topicPartitions().forEach(topicPartitions ->
+                partitions.put(topicPartitions.topicId(), new TreeSet<>(topicPartitions.partitions())));
+            return Optional.of(new LocalAssignment(nextLocalEpoch, partitions));
+
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index c2bdd3f860991..8dc033ea5fcf9 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -172,9 +172,12 @@ public class MembershipManagerImpl implements MembershipManager {
     private final Optional<String> serverAssignor;
 
     /**
-     * Assignment that the member received from the server and successfully processed.
+     * Assignment that the member received from the server and successfully processed, together with
+     * its local epoch.
+     *
+     * This is equal to LocalAssignment.NONE when we are not in a group, or we haven't reconciled any assignment yet.
      */
-    private Map<Uuid, SortedSet<Integer>> currentAssignment;
+    private LocalAssignment currentAssignment;
 
     /**
      * Subscription state object holding the current assignment the member has for the topics it
@@ -210,12 +213,12 @@ public class MembershipManagerImpl implements MembershipManager {
     private final Map<Uuid, String> assignedTopicNamesCache;
 
     /**
-     * Topic IDs and partitions received in the last target assignment. Items are added to this set
-     * every time a target assignment is received. This is where the member collects the assignment
-     * received from the broker, even though it may not be ready to fully reconcile due to missing
-     * metadata.
+     * Topic IDs and partitions received in the last target assignment, together with its local epoch.
+     *
+     * This member variable is reassigned every time a new assignment is received.
+     * It is equal to LocalAssignment.NONE whenever we are not in a group.
      */
-    private final Map<Uuid, SortedSet<Integer>> currentTargetAssignment;
+    private LocalAssignment currentTargetAssignment;
 
     /**     
      * If there is a reconciliation running (triggering commit, callbacks) for the
@@ -329,8 +332,8 @@ public MembershipManagerImpl(String groupId,
         this.commitRequestManager = commitRequestManager;
         this.metadata = metadata;
         this.assignedTopicNamesCache = new HashMap<>();
-        this.currentTargetAssignment = new HashMap<>();
-        this.currentAssignment = new HashMap<>();
+        this.currentTargetAssignment = LocalAssignment.NONE;
+        this.currentAssignment = LocalAssignment.NONE;
         this.log = logContext.logger(MembershipManagerImpl.class);
         this.stateUpdatesListeners = new ArrayList<>();
         this.clientTelemetryReporter = clientTelemetryReporter;
@@ -451,9 +454,6 @@ public void onHeartbeatSuccess(ConsumerGroupHeartbeatResponseData response) {
                 return;
             }
             processAssignmentReceived(assignment);
-
-        } else if (targetAssignmentReconciled()) {
-            transitionTo(MemberState.STABLE);
         }
     }
 
@@ -508,9 +508,11 @@ private void processAssignmentReceived(ConsumerGroupHeartbeatResponseData.Assign
      */
     private void replaceTargetAssignmentWithNewAssignment(
             ConsumerGroupHeartbeatResponseData.Assignment assignment) {
-        currentTargetAssignment.clear();
-        assignment.topicPartitions().forEach(topicPartitions ->
-            currentTargetAssignment.put(topicPartitions.topicId(), new TreeSet<>(topicPartitions.partitions())));
+        currentTargetAssignment.updateWith(assignment).ifPresent(updatedAssignment -> {
+            log.debug("Target assignment updated from {} to {}. Member will reconcile it on the next poll.",
+                currentTargetAssignment, updatedAssignment);
+            currentTargetAssignment = updatedAssignment;
+        });
     }
 
     /**
@@ -568,7 +570,11 @@ public void transitionToFenced() {
     public void transitionToFatal() {
         MemberState previousState = state;
         transitionTo(MemberState.FATAL);
-        log.error("Member {} with epoch {} transitioned to {} state", memberId, memberEpoch, MemberState.FATAL);
+        if (memberId.isEmpty()) {
+            log.error("Member {} with epoch {} transitioned to {} state", memberId, memberEpoch, MemberState.FATAL);
+        } else {
+            log.error("Non-member transitioned to {} state", MemberState.FATAL);
+        }
         notifyEpochChange(Optional.empty(), Optional.empty());
 
         if (previousState == MemberState.UNSUBSCRIBED) {
@@ -604,7 +610,7 @@ private void clearSubscription() {
         if (subscriptions.hasAutoAssignedPartitions()) {
             subscriptions.assignFromSubscribed(Collections.emptySet());
         }
-        updateCurrentAssignment(Collections.emptySet());
+        currentAssignment = LocalAssignment.NONE;
         clearPendingAssignmentsAndLocalNamesCache();
     }
 
@@ -620,7 +626,6 @@ private void updateSubscriptionAwaitingCallback(SortedSet<TopicIdPartition> assi
                                                     SortedSet<TopicPartition> addedPartitions) {
         Collection<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedPartitions);
         subscriptions.assignFromSubscribedAwaitingCallback(assignedTopicPartitions, addedPartitions);
-        updateCurrentAssignment(assignedPartitions);
     }
 
     /**
@@ -748,7 +753,7 @@ public void transitionToSendingLeaveGroup(boolean dueToExpiredPollTimer) {
                 ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH :
                 ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH;
         updateMemberEpoch(leaveEpoch);
-        currentAssignment = new HashMap<>();
+        currentAssignment = LocalAssignment.NONE;
         transitionTo(MemberState.LEAVING);
     }
 
@@ -889,12 +894,12 @@ private void transitionToStale() {
      */
     void maybeReconcile() {
         if (targetAssignmentReconciled()) {
-            log.debug("Ignoring reconciliation attempt. Target assignment is equal to the " +
+            log.trace("Ignoring reconciliation attempt. Target assignment is equal to the " +
                     "current assignment.");
             return;
         }
         if (reconciliationInProgress) {
-            log.debug("Ignoring reconciliation attempt. Another reconciliation is already in progress. Assignment " +
+            log.trace("Ignoring reconciliation attempt. Another reconciliation is already in progress. Assignment " +
                 currentTargetAssignment + " will be handled in the next reconciliation loop.");
             return;
         }
@@ -902,29 +907,25 @@ void maybeReconcile() {
         // Find the subset of the target assignment that can be resolved to topic names, and trigger a metadata update
         // if some topic IDs are not resolvable.
         SortedSet<TopicIdPartition> assignedTopicIdPartitions = findResolvableAssignmentAndTriggerMetadataUpdate();
+        final LocalAssignment resolvedAssignment = new LocalAssignment(currentTargetAssignment.localEpoch, assignedTopicIdPartitions);
+
+        if (!currentAssignment.isNone() && resolvedAssignment.partitions.equals(currentAssignment.partitions)) {
+            log.debug("There are unresolved partitions, and the resolvable fragment of the target assignment {} is equal to the current "
+                + "assignment. Bumping the local epoch of the assignment and acknowledging the partially resolved assignment",
+                resolvedAssignment.partitions);
+            currentAssignment = resolvedAssignment;
+            transitionTo(MemberState.ACKNOWLEDGING);
+            return;
+        }
 
-        SortedSet<TopicPartition> ownedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
-        ownedPartitions.addAll(subscriptions.assignedPartitions());
+        markReconciliationInProgress();
 
         // Keep copy of assigned TopicPartitions created from the TopicIdPartitions that are
         // being reconciled. Needed for interactions with the centralized subscription state that
         // does not support topic IDs yet, and for the callbacks.
         SortedSet<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedTopicIdPartitions);
-
-        // Check same assignment. Based on topic names for now, until topic IDs are properly
-        // supported in the centralized subscription state object. Note that this check is
-        // required to make sure that reconciliation is not triggered if the assignment ready to
-        // be reconciled is the same as the current one (even though the member may remain
-        // in RECONCILING state if it has some unresolved assignments).
-        boolean sameAssignmentReceived = assignedTopicPartitions.equals(ownedPartitions);
-
-        if (sameAssignmentReceived) {
-            log.debug("Ignoring reconciliation attempt. Target assignment ready to reconcile {} " +
-                    "is equal to the member current assignment {}.", assignedTopicPartitions, ownedPartitions);
-            return;
-        }
-
-        markReconciliationInProgress();
+        SortedSet<TopicPartition> ownedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+        ownedPartitions.addAll(subscriptions.assignedPartitions());
 
         // Partitions to assign (not previously owned)
         SortedSet<TopicPartition> addedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
@@ -936,12 +937,13 @@ void maybeReconcile() {
         revokedPartitions.addAll(ownedPartitions);
         revokedPartitions.removeAll(assignedTopicPartitions);
 
-        log.info("Updating assignment with\n" +
+        log.info("Updating assignment with local epoch {}\n" +
                         "\tAssigned partitions:                       {}\n" +
                         "\tCurrent owned partitions:                  {}\n" +
                         "\tAdded partitions (assigned - owned):       {}\n" +
                         "\tRevoked partitions (owned - assigned):     {}\n",
-                assignedTopicIdPartitions,
+                resolvedAssignment.localEpoch,
+                assignedTopicPartitions,
                 ownedPartitions,
                 addedPartitions,
                 revokedPartitions
@@ -970,7 +972,12 @@ void maybeReconcile() {
                 log.debug("Auto-commit before reconciling new assignment completed successfully.");
             }
 
-            revokeAndAssign(assignedTopicIdPartitions, revokedPartitions, addedPartitions);
+            revokeAndAssign(resolvedAssignment, assignedTopicIdPartitions, revokedPartitions, addedPartitions);
+        }).exceptionally(error -> {
+            if (error != null) {
+                log.error("Reconciliation failed.", error);
+            }
+            return null;
         });
     }
 
@@ -988,7 +995,8 @@ long getExpirationTimeForTimeout(final long timeoutMs) {
      * then complete the reconciliation by updating the assignment and making the appropriate state
      * transition. Note that if any of the 2 callbacks fails, the reconciliation should fail.
      */
-    private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitions,
+    private void revokeAndAssign(LocalAssignment resolvedAssignment,
+                                 SortedSet<TopicIdPartition> assignedTopicIdPartitions,
                                  SortedSet<TopicPartition> revokedPartitions,
                                  SortedSet<TopicPartition> addedPartitions) {
         CompletableFuture<Void> revocationResult;
@@ -1028,8 +1036,9 @@ private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitio
                 // RECONCILING -> FENCED transition.
                 log.error("Reconciliation failed.", error);
             } else {
-                boolean memberHasRejoined = memberEpochOnReconciliationStart != memberEpoch;
-                if (state == MemberState.RECONCILING && !memberHasRejoined) {
+                if (state == MemberState.RECONCILING) {
+                    currentAssignment = resolvedAssignment;
+
                     // Reschedule the auto commit starting from now that the member has a new assignment.
                     commitRequestManager.resetAutoCommitTimer();
 
@@ -1045,12 +1054,8 @@ private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitio
     }
 
     // Visible for testing.
-    void updateCurrentAssignment(Set<TopicIdPartition> assignedTopicIdPartitions) {
-        currentAssignment.clear();
-        assignedTopicIdPartitions.forEach(topicIdPartition -> {
-            Uuid topicId = topicIdPartition.topicId();
-            currentAssignment.computeIfAbsent(topicId, k -> new TreeSet<>()).add(topicIdPartition.partition());
-        });
+    void updateAssignment(Map<Uuid, SortedSet<Integer>> partitions) {
+        currentAssignment = new LocalAssignment(0, partitions);
     }
 
     /**
@@ -1109,7 +1114,7 @@ void markReconciliationCompleted() {
      */
     private SortedSet<TopicIdPartition> findResolvableAssignmentAndTriggerMetadataUpdate() {
         final SortedSet<TopicIdPartition> assignmentReadyToReconcile = new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR);
-        final HashMap<Uuid, SortedSet<Integer>> unresolved = new HashMap<>(currentTargetAssignment);
+        final HashMap<Uuid, SortedSet<Integer>> unresolved = new HashMap<>(currentTargetAssignment.partitions);
 
         // Try to resolve topic names from metadata cache or subscription cache, and move
         // assignments from the "unresolved" collection, to the "assignmentReadyToReconcile" one.
@@ -1382,7 +1387,7 @@ private void logPausedPartitionsBeingRevoked(Set<TopicPartition> partitionsToRev
      * or the next reconciliation loop). Remove all elements from the topic names cache.
      */
     private void clearPendingAssignmentsAndLocalNamesCache() {
-        currentTargetAssignment.clear();
+        currentTargetAssignment = LocalAssignment.NONE;
         assignedTopicNamesCache.clear();
     }
 
@@ -1424,7 +1429,7 @@ public Optional<String> serverAssignor() {
      * {@inheritDoc}
      */
     @Override
-    public Map<Uuid, SortedSet<Integer>> currentAssignment() {
+    public LocalAssignment currentAssignment() {
         return this.currentAssignment;
     }
 
@@ -1449,9 +1454,15 @@ Set<Uuid> topicsAwaitingReconciliation() {
      * Visible for testing.
      */
     Map<Uuid, SortedSet<Integer>> topicPartitionsAwaitingReconciliation() {
+        if (currentTargetAssignment == LocalAssignment.NONE) {
+            return Collections.emptyMap();
+        }
+        if (currentAssignment == LocalAssignment.NONE) {
+            return currentTargetAssignment.partitions;
+        }
         final Map<Uuid, SortedSet<Integer>> topicPartitionMap = new HashMap<>();
-        currentTargetAssignment.forEach((topicId, targetPartitions) -> {
-            final SortedSet<Integer> reconciledPartitions = currentAssignment.get(topicId);
+        currentTargetAssignment.partitions.forEach((topicId, targetPartitions) -> {
+            final SortedSet<Integer> reconciledPartitions = currentAssignment.partitions.get(topicId);
             if (!targetPartitions.equals(reconciledPartitions)) {
                 final TreeSet<Integer> missingPartitions = new TreeSet<>(targetPartitions);
                 if (reconciledPartitions != null) {
@@ -1497,4 +1508,5 @@ public PollResult poll(final long currentTimeMs) {
     List<MemberStateListener> stateListeners() {
         return unmodifiableList(stateUpdatesListeners);
     }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
index 3409871813391..653394cd4f1e5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java
@@ -19,7 +19,9 @@
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.HeartbeatRequestManager.HeartbeatRequestState;
+import org.apache.kafka.clients.consumer.internals.HeartbeatRequestManager.HeartbeatState;
+import org.apache.kafka.clients.consumer.internals.MembershipManager.LocalAssignment;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
@@ -27,11 +29,13 @@
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData.Assignment;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest;
+import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.Builder;
 import org.apache.kafka.common.requests.ConsumerGroupHeartbeatResponse;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -53,10 +57,11 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Random;
-import java.util.concurrent.BlockingQueue;
+import java.util.SortedSet;
 import java.util.concurrent.TimeUnit;
 
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_GROUP_INSTANCE_ID;
@@ -65,7 +70,9 @@
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR;
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_RETRY_BACKOFF_MAX_MS;
 import static org.apache.kafka.clients.consumer.internals.ConsumerTestBuilder.DEFAULT_RETRY_BACKOFF_MS;
+import static org.apache.kafka.common.utils.Utils.mkSortedSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -76,16 +83,11 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 
 public class HeartbeatRequestManagerTest {
-    private long retryBackoffMs = DEFAULT_RETRY_BACKOFF_MS;
-    private int heartbeatIntervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS;
-    private int maxPollIntervalMs = DEFAULT_MAX_POLL_INTERVAL_MS;
-    private long retryBackoffMaxMs = DEFAULT_RETRY_BACKOFF_MAX_MS;
     private static final String DEFAULT_GROUP_ID = "groupId";
     private static final String CONSUMER_COORDINATOR_METRICS = "consumer-coordinator-metrics";
 
@@ -102,7 +104,6 @@ public class HeartbeatRequestManagerTest {
     private final String memberId = "member-id";
     private final int memberEpoch = 1;
     private BackgroundEventHandler backgroundEventHandler;
-    private BlockingQueue<BackgroundEvent> backgroundEventQueue;
     private Metrics metrics;
 
     @BeforeEach
@@ -118,7 +119,6 @@ private void setUp(Optional<ConsumerTestBuilder.GroupInformation> groupInfo) {
         heartbeatRequestState = testBuilder.heartbeatRequestState.orElseThrow(IllegalStateException::new);
         heartbeatState = testBuilder.heartbeatState.orElseThrow(IllegalStateException::new);
         backgroundEventHandler = testBuilder.backgroundEventHandler;
-        backgroundEventQueue = testBuilder.backgroundEventQueue;
         subscriptions = testBuilder.subscriptions;
         membershipManager = testBuilder.membershipManager.orElseThrow(IllegalStateException::new);
         metadata = testBuilder.metadata;
@@ -177,7 +177,7 @@ public void testFirstHeartbeatIncludesRequiredInfoToJoinGroupAndGetAssignments(s
         NetworkClientDelegate.PollResult pollResult = heartbeatRequestManager.poll(time.milliseconds());
         assertEquals(1, pollResult.unsentRequests.size());
         NetworkClientDelegate.UnsentRequest request = pollResult.unsentRequests.get(0);
-        assertTrue(request.requestBuilder() instanceof ConsumerGroupHeartbeatRequest.Builder);
+        assertInstanceOf(Builder.class, request.requestBuilder());
 
         ConsumerGroupHeartbeatRequest heartbeatRequest =
                 (ConsumerGroupHeartbeatRequest) request.requestBuilder().build(version);
@@ -186,7 +186,7 @@ public void testFirstHeartbeatIncludesRequiredInfoToJoinGroupAndGetAssignments(s
         assertTrue(heartbeatRequest.data().memberId().isEmpty());
         assertEquals(0, heartbeatRequest.data().memberEpoch());
 
-        // Should include subscription and group basic info to start getting assignments.
+        // Should include subscription and group basic info to start getting assignments, as well as rebalanceTimeoutMs
         assertEquals(Collections.singletonList(topic), heartbeatRequest.data().subscribedTopicNames());
         assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, heartbeatRequest.data().rebalanceTimeoutMs());
         assertEquals(DEFAULT_GROUP_ID, heartbeatRequest.data().groupId());
@@ -312,19 +312,75 @@ public void testValidateConsumerGroupHeartbeatRequest(final short version) {
         NetworkClientDelegate.PollResult pollResult = heartbeatRequestManager.poll(time.milliseconds());
         assertEquals(1, pollResult.unsentRequests.size());
         NetworkClientDelegate.UnsentRequest request = pollResult.unsentRequests.get(0);
-        assertTrue(request.requestBuilder() instanceof ConsumerGroupHeartbeatRequest.Builder);
+        assertInstanceOf(Builder.class, request.requestBuilder());
 
         ConsumerGroupHeartbeatRequest heartbeatRequest =
             (ConsumerGroupHeartbeatRequest) request.requestBuilder().build(version);
         assertEquals(DEFAULT_GROUP_ID, heartbeatRequest.data().groupId());
         assertEquals(memberId, heartbeatRequest.data().memberId());
         assertEquals(memberEpoch, heartbeatRequest.data().memberEpoch());
-        assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, heartbeatRequest.data().rebalanceTimeoutMs());
+        assertEquals(10000, heartbeatRequest.data().rebalanceTimeoutMs());
         assertEquals(subscribedTopics, heartbeatRequest.data().subscribedTopicNames());
         assertEquals(DEFAULT_GROUP_INSTANCE_ID, heartbeatRequest.data().instanceId());
         assertEquals(DEFAULT_REMOTE_ASSIGNOR, heartbeatRequest.data().serverAssignor());
     }
 
+    @ParameterizedTest
+    @ApiKeyVersionsSource(apiKey = ApiKeys.CONSUMER_GROUP_HEARTBEAT)
+    public void testValidateConsumerGroupHeartbeatRequestAssignmentSentWhenLocalEpochChanges(final short version) {
+        CoordinatorRequestManager coordinatorRequestManager = mock(CoordinatorRequestManager.class);
+        MembershipManager membershipManager = mock(MembershipManager.class);
+        BackgroundEventHandler backgroundEventHandler = mock(BackgroundEventHandler.class);
+        SubscriptionState subscriptionState = mock(SubscriptionState.class);
+        HeartbeatRequestState requestState = mock(HeartbeatRequestState.class);
+        HeartbeatState heartbeatState = new HeartbeatState(subscriptionState, membershipManager, DEFAULT_MAX_POLL_INTERVAL_MS);
+
+        HeartbeatRequestManager heartbeatRequestManager = createHeartbeatRequestManager(
+            coordinatorRequestManager,
+            membershipManager,
+            heartbeatState,
+            requestState,
+            backgroundEventHandler
+        );
+
+        when(membershipManager.shouldHeartbeatNow()).thenReturn(true);
+        when(coordinatorRequestManager.coordinator()).thenReturn(Optional.of(new Node(1, "localhost", 9999)));
+
+        Uuid topicId = Uuid.randomUuid();
+        ConsumerGroupHeartbeatRequestData.TopicPartitions expectedTopicPartitions =
+            new ConsumerGroupHeartbeatRequestData.TopicPartitions();
+        Map<Uuid, SortedSet<Integer>> testAssignment = Collections.singletonMap(
+            topicId, mkSortedSet(0)
+        );
+        expectedTopicPartitions.setTopicId(topicId);
+        expectedTopicPartitions.setPartitions(Collections.singletonList(0));
+
+        // First heartbeat, include assignment
+        when(membershipManager.currentAssignment()).thenReturn(new LocalAssignment(0, testAssignment));
+
+        ConsumerGroupHeartbeatRequest heartbeatRequest1 = getHeartbeatRequest(heartbeatRequestManager, version);
+        assertEquals(Collections.singletonList(expectedTopicPartitions), heartbeatRequest1.data().topicPartitions());
+
+        // Assignment did not change, so no assignment should be sent
+        ConsumerGroupHeartbeatRequest heartbeatRequest2 = getHeartbeatRequest(heartbeatRequestManager, version);
+        assertNull(heartbeatRequest2.data().topicPartitions());
+
+        // Local epoch bumped, so assignment should be sent
+        when(membershipManager.currentAssignment()).thenReturn(new LocalAssignment(1, testAssignment));
+
+        ConsumerGroupHeartbeatRequest heartbeatRequest3 = getHeartbeatRequest(heartbeatRequestManager, version);
+        assertEquals(Collections.singletonList(expectedTopicPartitions), heartbeatRequest3.data().topicPartitions());
+    }
+
+    private ConsumerGroupHeartbeatRequest getHeartbeatRequest(HeartbeatRequestManager heartbeatRequestManager, final short version) {
+        // Create a ConsumerHeartbeatRequest and verify the payload -- no assignment should be sent
+        NetworkClientDelegate.PollResult pollResult = heartbeatRequestManager.poll(time.milliseconds());
+        assertEquals(1, pollResult.unsentRequests.size());
+        NetworkClientDelegate.UnsentRequest request = pollResult.unsentRequests.get(0);
+        assertInstanceOf(Builder.class, request.requestBuilder());
+        return (ConsumerGroupHeartbeatRequest) request.requestBuilder().build(version);
+    }
+
     @ParameterizedTest
     @MethodSource("errorProvider")
     public void testHeartbeatResponseOnErrorHandling(final Errors error, final boolean isFatal) {
@@ -345,7 +401,7 @@ public void testHeartbeatResponseOnErrorHandling(final Errors error, final boole
 
         switch (error) {
             case NONE:
-                verify(membershipManager, times(2)).onHeartbeatSuccess(mockResponse.data());
+                verify(membershipManager).onHeartbeatSuccess(mockResponse.data());
                 assertEquals(DEFAULT_HEARTBEAT_INTERVAL_MS, heartbeatRequestState.nextHeartbeatMs(time.milliseconds()));
                 break;
 
@@ -415,7 +471,7 @@ public void testHeartbeatState() {
         assertEquals(-1, data.rebalanceTimeoutMs());
         assertNull(data.subscribedTopicNames());
         assertNull(data.serverAssignor());
-        assertNull(data.topicPartitions());
+        assertEquals(data.topicPartitions(), Collections.emptyList());
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.STABLE, membershipManager.state());
 
@@ -432,7 +488,7 @@ public void testHeartbeatState() {
         assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, data.rebalanceTimeoutMs());
         assertEquals(Collections.singletonList(topic), data.subscribedTopicNames());
         assertEquals(ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR, data.serverAssignor());
-        assertNull(data.topicPartitions());
+        assertEquals(Collections.emptyList(), data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.JOINING, membershipManager.state());
 
@@ -445,7 +501,7 @@ public void testHeartbeatState() {
         assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, data.rebalanceTimeoutMs());
         assertEquals(Collections.singletonList(topic), data.subscribedTopicNames());
         assertEquals(ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR, data.serverAssignor());
-        assertNull(data.topicPartitions());
+        assertEquals(Collections.emptyList(), data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.JOINING, membershipManager.state());
 
@@ -478,9 +534,9 @@ public void testPollTimerExpiration() {
         heartbeatRequestState = spy(new HeartbeatRequestManager.HeartbeatRequestState(
                 new LogContext(),
                 time,
-                heartbeatIntervalMs,
-                retryBackoffMs,
-                retryBackoffMaxMs,
+                DEFAULT_HEARTBEAT_INTERVAL_MS,
+                DEFAULT_RETRY_BACKOFF_MS,
+                DEFAULT_RETRY_BACKOFF_MAX_MS,
                 0));
         backgroundEventHandler = mock(BackgroundEventHandler.class);
 
@@ -495,8 +551,8 @@ public void testPollTimerExpiration() {
 
         // On poll timer expiration, the member should send a last heartbeat to leave the group
         // and notify the membership manager
-        time.sleep(maxPollIntervalMs);
-        assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
+        time.sleep(DEFAULT_MAX_POLL_INTERVAL_MS);
+        assertHeartbeat(heartbeatRequestManager, DEFAULT_HEARTBEAT_INTERVAL_MS);
         verify(membershipManager).transitionToSendingLeaveGroup(true);
         verify(heartbeatState).reset();
         verify(heartbeatRequestState).reset();
@@ -509,7 +565,7 @@ public void testPollTimerExpiration() {
         assertTrue(pollTimer.notExpired());
         verify(membershipManager).maybeRejoinStaleMember();
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
-        assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
+        assertHeartbeat(heartbeatRequestManager, DEFAULT_HEARTBEAT_INTERVAL_MS);
     }
 
     /**
@@ -524,7 +580,7 @@ public void testPollTimerExpirationShouldNotMarkMemberStaleIfMemberAlreadyLeavin
         when(membershipManager.shouldSkipHeartbeat()).thenReturn(false);
         when(membershipManager.isLeavingGroup()).thenReturn(true);
 
-        time.sleep(maxPollIntervalMs);
+        time.sleep(DEFAULT_MAX_POLL_INTERVAL_MS);
         NetworkClientDelegate.PollResult result = heartbeatRequestManager.poll(time.milliseconds());
 
         // No transition to leave due to stale member should be triggered, because the member is
@@ -547,8 +603,8 @@ public void testHeartbeatMetrics() {
             new LogContext(),
             time,
             0, // This initial interval should be 0 to ensure heartbeat on the clock
-            retryBackoffMs,
-            retryBackoffMaxMs,
+            DEFAULT_RETRY_BACKOFF_MS,
+            DEFAULT_RETRY_BACKOFF_MAX_MS,
             0);
         backgroundEventHandler = mock(BackgroundEventHandler.class);
         heartbeatRequestManager = createHeartbeatRequestManager(
@@ -567,11 +623,11 @@ public void testHeartbeatMetrics() {
 
         // test poll
         assertHeartbeat(heartbeatRequestManager, 0);
-        time.sleep(heartbeatIntervalMs);
+        time.sleep(DEFAULT_HEARTBEAT_INTERVAL_MS);
         assertEquals(1.0, getMetric("heartbeat-total").metricValue());
-        assertEquals((double) TimeUnit.MILLISECONDS.toSeconds(heartbeatIntervalMs), getMetric("last-heartbeat-seconds-ago").metricValue());
+        assertEquals((double) TimeUnit.MILLISECONDS.toSeconds(DEFAULT_HEARTBEAT_INTERVAL_MS), getMetric("last-heartbeat-seconds-ago").metricValue());
 
-        assertHeartbeat(heartbeatRequestManager, heartbeatIntervalMs);
+        assertHeartbeat(heartbeatRequestManager, DEFAULT_HEARTBEAT_INTERVAL_MS);
         assertEquals(0.06d, (double) getMetric("heartbeat-rate").metricValue(), 0.005d);
         assertEquals(2.0, getMetric("heartbeat-total").metricValue());
 
@@ -625,11 +681,17 @@ private void assertNoHeartbeat(HeartbeatRequestManager hrm) {
     private void mockStableMember() {
         membershipManager.onSubscriptionUpdated();
         // Heartbeat response without assignment to set the state to STABLE.
+        when(subscriptions.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptions.rebalanceListener()).thenReturn(Optional.empty());
         ConsumerGroupHeartbeatResponse rs1 = new ConsumerGroupHeartbeatResponse(new ConsumerGroupHeartbeatResponseData()
                 .setHeartbeatIntervalMs(DEFAULT_HEARTBEAT_INTERVAL_MS)
                 .setMemberId(memberId)
-                .setMemberEpoch(memberEpoch));
+                .setMemberEpoch(memberEpoch)
+                .setAssignment(new Assignment())
+        );
         membershipManager.onHeartbeatSuccess(rs1.data());
+        membershipManager.poll(time.milliseconds());
+        membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.STABLE, membershipManager.state());
     }
 
@@ -666,23 +728,6 @@ private static Collection<Arguments> errorProvider() {
     private ClientResponse createHeartbeatResponse(
         final NetworkClientDelegate.UnsentRequest request,
         final Errors error
-    ) {
-        return createHeartbeatResponse(request, error, memberId, memberEpoch);
-    }
-
-    private ClientResponse createHeartbeatResponseWithMemberIdNull(
-        final NetworkClientDelegate.UnsentRequest request,
-        final Errors error,
-        final int memberEpoch
-    ) {
-        return createHeartbeatResponse(request, error, null, memberEpoch);
-    }
-
-    private ClientResponse createHeartbeatResponse(
-        final NetworkClientDelegate.UnsentRequest request,
-        final Errors error,
-        final String memberId,
-        final int memberEpoch
     ) {
         ConsumerGroupHeartbeatResponseData data = new ConsumerGroupHeartbeatResponseData()
             .setErrorCode(error.code())
@@ -711,10 +756,10 @@ private ConsumerConfig config() {
         prop.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
         prop.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
 
-        prop.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, String.valueOf(maxPollIntervalMs));
-        prop.setProperty(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(retryBackoffMs));
-        prop.setProperty(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG, String.valueOf(retryBackoffMaxMs));
-        prop.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, String.valueOf(heartbeatIntervalMs));
+        prop.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, String.valueOf(DEFAULT_MAX_POLL_INTERVAL_MS));
+        prop.setProperty(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(DEFAULT_RETRY_BACKOFF_MS));
+        prop.setProperty(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG, String.valueOf(DEFAULT_RETRY_BACKOFF_MAX_MS));
+        prop.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, String.valueOf(DEFAULT_HEARTBEAT_INTERVAL_MS));
         return new ConsumerConfig(prop);
     }
 
@@ -729,7 +774,7 @@ private HeartbeatRequestManager createHeartbeatRequestManager(
             final HeartbeatRequestManager.HeartbeatRequestState heartbeatRequestState,
             final BackgroundEventHandler backgroundEventHandler) {
         LogContext logContext = new LogContext();
-        pollTimer = time.timer(maxPollIntervalMs);
+        pollTimer = time.timer(DEFAULT_MAX_POLL_INTERVAL_MS);
         return new HeartbeatRequestManager(
                 logContext,
                 pollTimer,
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index bc7be88c3bd8d..4f893a4efbc11 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -29,6 +29,8 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData.Assignment;
+import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData.TopicPartitions;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest;
@@ -201,14 +203,14 @@ public void testReconcilingWhenReceivingAssignmentFoundInMetadata() {
     }
 
     @Test
-    public void testTransitionToReconcilingOnlyIfAssignmentReceived() {
+    public void testTransitionToReconcilingIfEmptyAssignmentReceived() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         assertEquals(MemberState.JOINING, membershipManager.state());
 
         ConsumerGroupHeartbeatResponse responseWithoutAssignment =
-                createConsumerGroupHeartbeatResponse(null);
+                createConsumerGroupHeartbeatResponse(new Assignment());
         membershipManager.onHeartbeatSuccess(responseWithoutAssignment.data());
-        assertNotEquals(MemberState.RECONCILING, membershipManager.state());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
 
         ConsumerGroupHeartbeatResponse responseWithAssignment =
                 createConsumerGroupHeartbeatResponse(createAssignment(true));
@@ -218,10 +220,7 @@ public void testTransitionToReconcilingOnlyIfAssignmentReceived() {
 
     @Test
     public void testMemberIdAndEpochResetOnFencedMembers() {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
-        assertEquals(MemberState.STABLE, membershipManager.state());
+        MembershipManagerImpl membershipManager = createMemberInStableState();
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
 
@@ -234,11 +233,7 @@ public void testMemberIdAndEpochResetOnFencedMembers() {
 
     @Test
     public void testTransitionToFatal() {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse =
-                createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
-        assertEquals(MemberState.STABLE, membershipManager.state());
+        MembershipManagerImpl membershipManager = createMemberInStableState(null);
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
 
@@ -271,6 +266,7 @@ public void testFencingWhenStateIsStable() {
 
     @Test
     public void testListenersGetNotifiedOnTransitionsToFatal() {
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty());
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         MemberStateListener listener = mock(MemberStateListener.class);
         membershipManager.registerStateListener(listener);
@@ -286,6 +282,7 @@ public void testListenersGetNotifiedOnTransitionsToFatal() {
 
     @Test
     public void testListenersGetNotifiedOnTransitionsToLeavingGroup() {
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty());
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         MemberStateListener listener = mock(MemberStateListener.class);
         membershipManager.registerStateListener(listener);
@@ -322,8 +319,11 @@ public void testListenersGetNotifiedOfMemberEpochUpdatesOnlyIfItChanges() {
     }
 
     private void mockStableMember(MembershipManagerImpl membershipManager) {
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(new Assignment());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+        membershipManager.poll(time.milliseconds());
+        membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.STABLE, membershipManager.state());
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(MEMBER_EPOCH, membershipManager.memberEpoch());
@@ -517,6 +517,114 @@ public void testDelayedReconciliationResultDiscardedIfMemberRejoins() {
         assertEquals(assignmentAfterRejoin, membershipManager.topicPartitionsAwaitingReconciliation());
     }
 
+    /**
+     * This is the case where a member rejoins (due to fence). If
+     * the member gets the same assignment again, it should still reconcile and ack the assignment.
+     */
+    @Test
+    public void testSameAssignmentReconciledAgainWhenFenced() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        Uuid topic1 = Uuid.randomUuid();
+        final Assignment assignment1 = new ConsumerGroupHeartbeatResponseData.Assignment();
+        final Assignment assignment2 = new ConsumerGroupHeartbeatResponseData.Assignment()
+            .setTopicPartitions(Collections.singletonList(
+                new TopicPartitions()
+                    .setTopicId(topic1)
+                    .setPartitions(Arrays.asList(0, 1, 2))
+            ));
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topic1, "topic1"));
+        assertEquals(toTopicIdPartitionMap(assignment1), membershipManager.currentAssignment().partitions);
+
+        // Receive assignment, wait on commit
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment2).data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        CompletableFuture<Void> commitResult = new CompletableFuture<>();
+        when(commitRequestManager.maybeAutoCommitSyncNow(anyLong())).thenReturn(commitResult);
+        membershipManager.poll(time.milliseconds());
+
+        // Get fenced, commit completes
+        membershipManager.transitionToFenced();
+
+        assertEquals(MemberState.JOINING, membershipManager.state());
+        assertTrue(membershipManager.currentAssignment().isNone());
+        assertTrue(subscriptionState.assignedPartitions().isEmpty());
+
+        commitResult.complete(null);
+
+        assertEquals(MemberState.JOINING, membershipManager.state());
+        assertTrue(membershipManager.currentAssignment().isNone());
+        assertTrue(subscriptionState.assignedPartitions().isEmpty());
+
+        // We have to reconcile & ack the assignment again
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment1).data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        membershipManager.poll(time.milliseconds());
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        membershipManager.onHeartbeatRequestSent();
+        assertEquals(MemberState.STABLE, membershipManager.state());
+        assertEquals(toTopicIdPartitionMap(assignment1), membershipManager.currentAssignment().partitions);
+    }
+
+    /**
+     * This is the case where we receive a new assignment while reconciling an existing one. The intermediate assignment
+     * is not applied, and a new assignment containing the same partitions is received and reconciled. In all assignments,
+     * one topic is not resolvable.
+     *
+     * We need to make sure that the last assignment is acked and applied, even though the set of partitions does not change.
+     * In this case, no rebalance listeners are run.
+     */
+    @Test
+    public void testSameAssignmentReconciledAgainWithMissingTopic() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        Uuid topic1 = Uuid.randomUuid();
+        Uuid topic2 = Uuid.randomUuid();
+        final Assignment assignment1 = new ConsumerGroupHeartbeatResponseData.Assignment()
+            .setTopicPartitions(Arrays.asList(
+                new TopicPartitions().setTopicId(topic1).setPartitions(Collections.singletonList(0)),
+                new TopicPartitions().setTopicId(topic2).setPartitions(Collections.singletonList(0))
+            ));
+        final Assignment assignment2 = new ConsumerGroupHeartbeatResponseData.Assignment()
+            .setTopicPartitions(Arrays.asList(
+                new TopicPartitions().setTopicId(topic1).setPartitions(Arrays.asList(0, 1)),
+                new TopicPartitions().setTopicId(topic2).setPartitions(Collections.singletonList(0))
+            ));
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topic1, "topic1"));
+
+        // Receive assignment - full reconciliation triggered
+        // stay in RECONCILING state, since an unresolved topic is assigned
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment1).data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        membershipManager.poll(time.milliseconds());
+        verifyReconciliationTriggeredAndCompleted(membershipManager,
+            Collections.singletonList(new TopicIdPartition(topic1, new TopicPartition("topic1", 0)))
+        );
+        membershipManager.onHeartbeatRequestSent();
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        clearInvocations(membershipManager);
+
+        // Receive extended assignment - assignment received but no reconciliation triggered
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment2).data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        verifyReconciliationNotTriggered(membershipManager);
+
+        // Receive original assignment again - full reconciliation not triggered but assignment is acked again
+        membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment1).data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        membershipManager.poll(time.milliseconds());
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        verifyReconciliationNotTriggered(membershipManager);
+        assertEquals(Collections.singletonMap(topic1, mkSortedSet(0)), membershipManager.currentAssignment().partitions);
+        assertEquals(mkSet(topic2), membershipManager.topicsAwaitingReconciliation());
+    }
+
+    private Map<Uuid, SortedSet<Integer>> toTopicIdPartitionMap(final Assignment assignment) {
+        Map<Uuid, SortedSet<Integer>> result = new HashMap<>();
+        for (TopicPartitions topicPartitions : assignment.topicPartitions()) {
+            result.put(topicPartitions.topicId(), new TreeSet<>(topicPartitions.partitions()));
+        }
+        return result;
+    }
+
     /**
      * This is the case where a member is stuck reconciling an assignment A (waiting on metadata, commit or callbacks), and the target
      * assignment changes due to new topics. If the reconciliation of A completes it should be applied (should update the assignment
@@ -683,6 +791,12 @@ public void testDelayedMetadataUsedToCompleteAssignment() {
         receiveAssignment(newAssignment, membershipManager);
         membershipManager.poll(time.milliseconds());
 
+        // No full reconciliation triggered, but assignment needs to be acknowledged.
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        assertTrue(membershipManager.shouldHeartbeatNow());
+
+        membershipManager.onHeartbeatRequestSent();
+
         verifyReconciliationNotTriggered(membershipManager);
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertEquals(Collections.singleton(topicId2), membershipManager.topicsAwaitingReconciliation());
@@ -721,7 +835,7 @@ public void testIgnoreHeartbeatWhenLeavingGroup() {
         assertEquals(MemberState.LEAVING, membershipManager.state());
         assertEquals(-1, membershipManager.memberEpoch());
         assertEquals(MEMBER_ID, membershipManager.memberId());
-        assertTrue(membershipManager.currentAssignment().isEmpty());
+        assertTrue(membershipManager.currentAssignment().isNone());
         assertFalse(leaveResult.isDone(), "Leave group result should not complete until the " +
             "heartbeat request to leave is sent out.");
     }
@@ -766,7 +880,7 @@ public void testLeaveGroupWhenMemberOwnsAssignment() {
             new TopicIdPartition(topicId, new TopicPartition(topicName, 1)));
         verifyReconciliationTriggeredAndCompleted(membershipManager, assignedPartitions);
 
-        assertEquals(1, membershipManager.currentAssignment().size());
+        assertEquals(1, membershipManager.currentAssignment().partitions.size());
 
         testLeaveGroupReleasesAssignmentAndResetsEpochToSendLeaveGroup(membershipManager);
     }
@@ -881,10 +995,7 @@ public void testFatalFailureWhenStateIsUnjoined() {
 
     @Test
     public void testFatalFailureWhenStateIsStable() {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
-        assertEquals(MemberState.STABLE, membershipManager.state());
+        MembershipManagerImpl membershipManager = createMemberInStableState(null);
 
         testStateUpdateOnFatalFailure(membershipManager);
     }
@@ -1027,10 +1138,14 @@ public void testNewEmptyAssignmentReplacesPreviousOneWaitingOnMetadata() {
         receiveEmptyAssignment(membershipManager);
 
         verifyReconciliationNotTriggered(membershipManager);
+
         membershipManager.poll(time.milliseconds());
 
+        verifyReconciliationTriggeredAndCompleted(membershipManager, Collections.emptyList());
+
+        membershipManager.onHeartbeatRequestSent();
+
         assertEquals(MemberState.STABLE, membershipManager.state());
-        verify(subscriptionState, never()).assignFromSubscribed(any());
     }
 
     @Test
@@ -1049,6 +1164,7 @@ public void testNewAssignmentNotInMetadataReplacesPreviousOneWaitingOnMetadata()
 
         verifyReconciliationNotTriggered(membershipManager);
         membershipManager.poll(time.milliseconds());
+        membershipManager.onHeartbeatRequestSent();
 
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         assertFalse(membershipManager.topicsAwaitingReconciliation().isEmpty());
@@ -1282,7 +1398,7 @@ public void testReconcileNewPartitionsAssignedAndRevoked() {
         membershipManager.poll(time.milliseconds());
 
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
-        assertEquals(topicIdPartitionsMap(topicId, 1, 2), membershipManager.currentAssignment());
+        assertEquals(topicIdPartitionsMap(topicId, 1, 2), membershipManager.currentAssignment().partitions);
         assertFalse(membershipManager.reconciliationInProgress());
 
         verify(subscriptionState).assignFromSubscribed(anyCollection());
@@ -1298,7 +1414,8 @@ public void testMetadataUpdatesReconcilesUnresolvedAssignments() {
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                                 .setTopicId(topicId)
                                 .setPartitions(Arrays.asList(0, 1))));
-        MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(false, targetAssignment);
+        MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true, targetAssignment);
+        membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.RECONCILING, membershipManager.state());
 
         // Should not trigger reconciliation, and request a metadata update.
@@ -1327,7 +1444,8 @@ public void testMetadataUpdatesRequestsAnotherUpdateIfNeeded() {
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
                                 .setTopicId(topicId)
                                 .setPartitions(Arrays.asList(0, 1))));
-        MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(false, targetAssignment);
+        MembershipManagerImpl membershipManager = mockJoinAndReceiveAssignment(true, targetAssignment);
+        membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.RECONCILING, membershipManager.state());
 
         // Should not trigger reconciliation, and request a metadata update.
@@ -1366,7 +1484,7 @@ public void testRevokePartitionsUsesTopicNamesLocalCacheWhenMetadataNotAvailable
             partitions.stream().map(p -> new TopicPartition(topicName, p)).collect(Collectors.toSet());
         Map<Uuid, SortedSet<Integer>> assignedTopicIdPartitions = Collections.singletonMap(topicId,
             new TreeSet<>(partitions));
-        assertEquals(assignedTopicIdPartitions, membershipManager.currentAssignment());
+        assertEquals(assignedTopicIdPartitions, membershipManager.currentAssignment().partitions);
         assertFalse(membershipManager.reconciliationInProgress());
 
         mockAckSent(membershipManager);
@@ -1393,8 +1511,6 @@ public void testRevokePartitionsUsesTopicNamesLocalCacheWhenMetadataNotAvailable
     @Test
     public void testOnSubscriptionUpdatedTransitionsToJoiningOnlyIfNotInGroup() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
-        verify(membershipManager).transitionToJoining();
-        clearInvocations(membershipManager);
         membershipManager.onSubscriptionUpdated();
         verify(membershipManager, never()).transitionToJoining();
     }
@@ -1443,7 +1559,7 @@ public void testListenerCallbacksBasic() {
         // Step 4: Send ack and make sure we're done and our listener was called appropriately
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.STABLE, membershipManager.state());
-        assertEquals(topicIdPartitionsMap(topicId, 0, 1), membershipManager.currentAssignment());
+        assertEquals(topicIdPartitionsMap(topicId, 0, 1), membershipManager.currentAssignment().partitions);
 
         assertEquals(0, listener.revokedCount());
         assertEquals(1, listener.assignedCount());
@@ -1512,7 +1628,7 @@ public void testListenerCallbacksThrowsErrorOnPartitionsRevoked() {
         membershipManager.poll(time.milliseconds());
 
         assertEquals(MemberState.RECONCILING, membershipManager.state());
-        assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment());
+        assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment().partitions);
         assertTrue(membershipManager.reconciliationInProgress());
         assertEquals(0, listener.revokedCount());
         assertEquals(0, listener.assignedCount());
@@ -1565,7 +1681,7 @@ public void testListenerCallbacksThrowsErrorOnPartitionsAssigned() {
         membershipManager.poll(time.milliseconds());
 
         assertEquals(MemberState.RECONCILING, membershipManager.state());
-        assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment());
+        assertEquals(topicIdPartitionsMap(topicId, 0), membershipManager.currentAssignment().partitions);
         assertTrue(membershipManager.reconciliationInProgress());
         assertEquals(0, listener.revokedCount());
         assertEquals(0, listener.assignedCount());
@@ -1666,19 +1782,11 @@ public void testAddedPartitionsNotEnabledAfterFailedOnPartitionsAssignedCallback
 
     @Test
     public void testOnPartitionsLostNoError() {
-        MembershipManagerImpl membershipManager = createMemberInStableState();
-        String topicName = "topic1";
-        Uuid topicId = Uuid.randomUuid();
-        mockOwnedPartition(membershipManager, topicId, topicName);
         testOnPartitionsLost(Optional.empty());
     }
 
     @Test
     public void testOnPartitionsLostError() {
-        MembershipManagerImpl membershipManager = createMemberInStableState();
-        String topicName = "topic1";
-        Uuid topicId = Uuid.randomUuid();
-        mockOwnedPartition(membershipManager, topicId, topicName);
         testOnPartitionsLost(Optional.of(new KafkaException("Intentional error for test")));
     }
 
@@ -1709,12 +1817,7 @@ public void testTransitionToLeavingWhileJoiningDueToStaleMember() {
 
     @Test
     public void testTransitionToLeavingWhileStableDueToStaleMember() {
-        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
-        membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
-        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
-        doNothing().when(subscriptionState).assignFromSubscribed(any());
-        assertEquals(MemberState.STABLE, membershipManager.state());
+        MembershipManagerImpl membershipManager = createMemberInStableState(null);
         assertLeaveGroupDueToExpiredPollAndTransitionToStale(membershipManager);
     }
 
@@ -1807,7 +1910,7 @@ private void mockPartitionOwnedAndNewPartitionAdded(String topicName,
         Uuid topicId = Uuid.randomUuid();
         TopicPartition owned = new TopicPartition(topicName, partitionOwned);
         when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(owned));
-        membershipManager.updateCurrentAssignment(Collections.singleton(new TopicIdPartition(topicId, owned)));
+        membershipManager.updateAssignment(Collections.singletonMap(topicId, mkSortedSet(partitionOwned)));
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.ofNullable(listener));
@@ -1819,6 +1922,9 @@ private void mockPartitionOwnedAndNewPartitionAdded(String topicName,
     private void testOnPartitionsLost(Optional<RuntimeException> lostError) {
         // Step 1: set up mocks
         MembershipManagerImpl membershipManager = createMemberInStableState();
+        String topicName = "topic1";
+        Uuid topicId = Uuid.randomUuid();
+        mockOwnedPartition(membershipManager, topicId, topicName);
         CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener(
                 Optional.empty(),
                 Optional.empty(),
@@ -1832,7 +1938,6 @@ private void testOnPartitionsLost(Optional<RuntimeException> lostError) {
         // Step 2: put the state machine into the appropriate... state
         membershipManager.transitionToFenced();
         assertEquals(MemberState.FENCED, membershipManager.state());
-        assertEquals(Collections.emptyMap(), membershipManager.currentAssignment());
         assertEquals(0, listener.revokedCount());
         assertEquals(0, listener.assignedCount());
         assertEquals(0, listener.lostCount());
@@ -1848,6 +1953,8 @@ private void testOnPartitionsLost(Optional<RuntimeException> lostError) {
                 true
         );
 
+        assertTrue(membershipManager.currentAssignment().isNone());
+
         // Step 4: Receive ack and make sure we're done and our listener was called appropriately
         membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.JOINING, membershipManager.state());
@@ -1952,7 +2059,7 @@ private void assertStaleMemberLeavesGroupAndClearsAssignment(MembershipManagerIm
         // Should reset epoch to leave the group and release the assignment (right away because
         // there is no onPartitionsLost callback defined)
         verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
-        assertTrue(membershipManager.currentAssignment().isEmpty());
+        assertTrue(membershipManager.currentAssignment().isNone());
         assertTrue(membershipManager.topicsAwaitingReconciliation().isEmpty());
         assertEquals(LEAVE_GROUP_MEMBER_EPOCH, membershipManager.memberEpoch());
     }
@@ -1960,29 +2067,62 @@ private void assertStaleMemberLeavesGroupAndClearsAssignment(MembershipManagerIm
     @Test
     public void testMemberJoiningTransitionsToStableWhenReceivingEmptyAssignment() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty());
         assertEquals(MemberState.JOINING, membershipManager.state());
         receiveEmptyAssignment(membershipManager);
 
         verifyReconciliationNotTriggered(membershipManager);
         membershipManager.poll(time.milliseconds());
+        membershipManager.onHeartbeatRequestSent();
 
         assertEquals(MemberState.STABLE, membershipManager.state());
     }
 
+    @Test
+    public void testMemberJoiningCallsRebalanceListenerWhenReceivingEmptyAssignment() {
+        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+        MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        receiveEmptyAssignment(membershipManager);
+
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        assertEquals(0, listener.assignedCount());
+        assertEquals(0, listener.revokedCount());
+        assertEquals(0, listener.lostCount());
+
+        verifyReconciliationNotTriggered(membershipManager);
+        membershipManager.poll(time.milliseconds());
+        performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
+            new TreeSet<>(Collections.emptyList()),
+            true
+        );
+
+        assertEquals(1, listener.assignedCount());
+        assertEquals(0, listener.revokedCount());
+        assertEquals(0, listener.lostCount());
+    }
+
     @Test
     public void testMetricsWhenHeartbeatFailed() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         membershipManager.onHeartbeatFailure();
 
-        // Not expecting rebalance failures without assignments being reconciled
-        assertEquals(0.0d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        // Not expecting rebalance failures with only the empty assignment being reconciled.
+        assertEquals(1.0d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
         assertEquals(0.0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
     }
 
     @Test
     public void testRebalanceMetricsOnSuccessfulRebalance() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(new Assignment());
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
         mockOwnedPartition(membershipManager, Uuid.randomUuid(), "topic1");
 
@@ -2067,12 +2207,12 @@ public void testRebalanceMetricsForMultipleReconcilations() {
 
         long secondRebalanceMs = listener.sleepMs;
         long total = firstRebalanaceTimesMs + secondRebalanceMs;
-        double avg = total / 2.0d;
+        double avg = total / 3.0d;
         long max = Math.max(firstRebalanaceTimesMs, secondRebalanceMs);
         assertEquals((double) total, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyTotal));
         assertEquals(avg, (double) getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyAvg), 1d);
         assertEquals((double) max, getMetricValue(metrics, rebalanceMetricsManager.rebalanceLatencyMax));
-        assertEquals(2d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
+        assertEquals(3d, getMetricValue(metrics, rebalanceMetricsManager.rebalanceTotal));
         // rate is not tested because it is subject to Rate implementation
         assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceRate));
         assertEquals(0d, getMetricValue(metrics, rebalanceMetricsManager.failedRebalanceTotal));
@@ -2083,7 +2223,7 @@ public void testRebalanceMetricsForMultipleReconcilations() {
     @Test
     public void testRebalanceMetricsOnFailedRebalance() {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(new Assignment());
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
 
         Uuid topicId = Uuid.randomUuid();
@@ -2180,7 +2320,7 @@ private void verifyReconciliationTriggeredAndCompleted(MembershipManagerImpl mem
         List<TopicPartition> expectedTopicPartitions = buildTopicPartitions(expectedAssignment);
         verify(subscriptionState).assignFromSubscribed(new HashSet<>(expectedTopicPartitions));
         Map<Uuid, SortedSet<Integer>> assignmentByTopicId = assignmentByTopicId(expectedAssignment);
-        assertEquals(assignmentByTopicId, membershipManager.currentAssignment());
+        assertEquals(assignmentByTopicId, membershipManager.currentAssignment().partitions);
 
         // The auto-commit interval should be reset (only once), when the reconciliation completes
         verify(commitRequestManager).resetAutoCommitTimer();
@@ -2231,7 +2371,7 @@ private void testRevocationCompleted(MembershipManagerImpl membershipManager,
                                          List<TopicIdPartition> expectedCurrentAssignment) {
         assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
         Map<Uuid, SortedSet<Integer>> assignmentByTopicId = assignmentByTopicId(expectedCurrentAssignment);
-        assertEquals(assignmentByTopicId, membershipManager.currentAssignment());
+        assertEquals(assignmentByTopicId, membershipManager.currentAssignment().partitions);
         assertFalse(membershipManager.reconciliationInProgress());
 
         verify(subscriptionState).markPendingRevocation(anySet());
@@ -2254,7 +2394,9 @@ private void mockOwnedPartitionAndAssignmentReceived(MembershipManagerImpl membe
                                                          String topicName,
                                                          Collection<TopicIdPartition> previouslyOwned) {
         when(subscriptionState.assignedPartitions()).thenReturn(getTopicPartitions(previouslyOwned));
-        membershipManager.updateCurrentAssignment(new HashSet<>(previouslyOwned));
+        HashMap<Uuid, SortedSet<Integer>> partitionsByTopicId = new HashMap<>();
+        partitionsByTopicId.put(topicId, new TreeSet<>(previouslyOwned.stream().map(TopicIdPartition::partition).collect(Collectors.toSet())));
+        membershipManager.updateAssignment(partitionsByTopicId);
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty()).thenReturn(Optional.empty());
@@ -2269,17 +2411,16 @@ private Set<TopicPartition> getTopicPartitions(Collection<TopicIdPartition> topi
     private void mockOwnedPartition(MembershipManagerImpl membershipManager, Uuid topicId, String topic) {
         int partition = 0;
         TopicPartition previouslyOwned = new TopicPartition(topic, partition);
-        membershipManager.updateCurrentAssignment(
-            Collections.singleton(new TopicIdPartition(topicId, new TopicPartition(topic, partition))));
+        membershipManager.updateAssignment(mkMap(mkEntry(topicId, new TreeSet<>(Collections.singletonList(partition)))));
         when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(previouslyOwned));
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
     }
 
-    private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean expectSubscriptionUpdated) {
-        return mockJoinAndReceiveAssignment(expectSubscriptionUpdated, createAssignment(expectSubscriptionUpdated));
+    private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean triggerReconciliation) {
+        return mockJoinAndReceiveAssignment(triggerReconciliation, createAssignment(triggerReconciliation));
     }
 
-    private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean expectSubscriptionUpdated,
+    private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean triggerReconciliation,
                                                                ConsumerGroupHeartbeatResponseData.Assignment assignment) {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
         ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(assignment);
@@ -2287,14 +2428,15 @@ private MembershipManagerImpl mockJoinAndReceiveAssignment(boolean expectSubscri
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty()).thenReturn(Optional.empty());
 
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
-        membershipManager.poll(time.milliseconds());
 
-        if (expectSubscriptionUpdated) {
+        if (triggerReconciliation) {
+            membershipManager.poll(time.milliseconds());
             verify(subscriptionState).assignFromSubscribed(anyCollection());
         } else {
             verify(subscriptionState, never()).assignFromSubscribed(anyCollection());
         }
 
+        clearInvocations(membershipManager, commitRequestManager);
         return membershipManager;
     }
 
@@ -2304,9 +2446,17 @@ private MembershipManagerImpl createMemberInStableState() {
 
     private MembershipManagerImpl createMemberInStableState(String groupInstanceId) {
         MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup(groupInstanceId, null);
-        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(null);
+        ConsumerGroupHeartbeatResponse heartbeatResponse = createConsumerGroupHeartbeatResponse(new Assignment());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.empty());
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+        membershipManager.poll(time.milliseconds());
+        assertEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        membershipManager.onHeartbeatRequestSent();
         assertEquals(MemberState.STABLE, membershipManager.state());
+
+        clearInvocations(subscriptionState, membershipManager, commitRequestManager);
         return membershipManager;
     }
 
@@ -2385,7 +2535,7 @@ private void testLeaveGroupReleasesAssignmentAndResetsEpochToSendLeaveGroup(Memb
         assertFalse(leaveResult.isCompletedExceptionally());
         assertEquals(MEMBER_ID, membershipManager.memberId());
         assertEquals(-1, membershipManager.memberEpoch());
-        assertTrue(membershipManager.currentAssignment().isEmpty());
+        assertTrue(membershipManager.currentAssignment().isNone());
         verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
     }
 
@@ -2508,7 +2658,7 @@ private MembershipManagerImpl memberJoinWithAssignment() {
         when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, "topic"));
         receiveAssignment(topicId, Collections.singletonList(0), membershipManager);
         membershipManager.onHeartbeatRequestSent();
-        assertFalse(membershipManager.currentAssignment().isEmpty());
+        assertFalse(membershipManager.currentAssignment().isNone());
         return membershipManager;
     }
 

From 67cb742b540796c80553f8ef03b8ccf10274ce7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Armando=20Garc=C3=ADa=20Sancio?=
 <jsancio@users.noreply.github.com>
Date: Mon, 18 Mar 2024 12:06:42 -0700
Subject: [PATCH 170/258] MINOR; Log reason for deleting a kraft snapshot
 (#15478)

There are three reasons why KRaft would delete a snapshot. One, it is older than the retention time. Two, the total number of bytes between the log and the snapshot excess the configuration. Three, the latest snapshot is newer than the log.

This change allows KRaft to log the exact reason why a snapshot is getting deleted.

Reviewers: David Arthur <mumrah@gmail.com>, Hailey Ni <hni@confluent.io>
---
 .../main/scala/kafka/MetadataLogConfig.scala  |  49 +++++++
 .../scala/kafka/raft/KafkaMetadataLog.scala   | 132 +++++++++++-------
 2 files changed, 131 insertions(+), 50 deletions(-)
 create mode 100755 core/src/main/scala/kafka/MetadataLogConfig.scala

diff --git a/core/src/main/scala/kafka/MetadataLogConfig.scala b/core/src/main/scala/kafka/MetadataLogConfig.scala
new file mode 100755
index 0000000000000..60afac20fa4d5
--- /dev/null
+++ b/core/src/main/scala/kafka/MetadataLogConfig.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.raft
+
+import kafka.server.KafkaConfig
+import org.apache.kafka.common.config.AbstractConfig
+import org.apache.kafka.storage.internals.log.LogConfig
+
+final case class MetadataLogConfig(
+  logSegmentBytes: Int,
+  logSegmentMinBytes: Int,
+  logSegmentMillis: Long,
+  retentionMaxBytes: Long,
+  retentionMillis: Long,
+  maxBatchSizeInBytes: Int,
+  maxFetchSizeInBytes: Int,
+  fileDeleteDelayMs: Long,
+  nodeId: Int
+)
+
+object MetadataLogConfig {
+  def apply(config: AbstractConfig, maxBatchSizeInBytes: Int, maxFetchSizeInBytes: Int): MetadataLogConfig = {
+    new MetadataLogConfig(
+      config.getInt(KafkaConfig.MetadataLogSegmentBytesProp),
+      config.getInt(KafkaConfig.MetadataLogSegmentMinBytesProp),
+      config.getLong(KafkaConfig.MetadataLogSegmentMillisProp),
+      config.getLong(KafkaConfig.MetadataMaxRetentionBytesProp),
+      config.getLong(KafkaConfig.MetadataMaxRetentionMillisProp),
+      maxBatchSizeInBytes,
+      maxFetchSizeInBytes,
+      LogConfig.DEFAULT_FILE_DELETE_DELAY_MS,
+      config.getInt(KafkaConfig.NodeIdProp)
+    )
+  }
+}
diff --git a/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala b/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala
index e41b650066537..514b08797a10a 100644
--- a/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala
+++ b/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala
@@ -17,10 +17,15 @@
 package kafka.raft
 
 import kafka.log.UnifiedLog
+import kafka.raft.KafkaMetadataLog.FullTruncation
+import kafka.raft.KafkaMetadataLog.RetentionMsBreach
+import kafka.raft.KafkaMetadataLog.RetentionSizeBreach
+import kafka.raft.KafkaMetadataLog.SnapshotDeletionReason
+import kafka.raft.KafkaMetadataLog.UnknownReason
 import kafka.server.KafkaConfig.{MetadataLogSegmentBytesProp, MetadataLogSegmentMinBytesProp}
-import kafka.server.{BrokerTopicStats, KafkaConfig, RequestLocal}
+import kafka.server.{BrokerTopicStats, RequestLocal}
 import kafka.utils.{CoreUtils, Logging}
-import org.apache.kafka.common.config.{AbstractConfig, TopicConfig}
+import org.apache.kafka.common.config.TopicConfig
 import org.apache.kafka.common.errors.InvalidConfigurationException
 import org.apache.kafka.common.record.{MemoryRecords, Records}
 import org.apache.kafka.common.utils.Time
@@ -179,7 +184,7 @@ final class KafkaMetadataLog private (
         (false, mutable.TreeMap.empty[OffsetAndEpoch, Option[FileRawSnapshotReader]])
     }
 
-    removeSnapshots(forgottenSnapshots)
+    removeSnapshots(forgottenSnapshots, FullTruncation)
     truncated
   }
 
@@ -340,6 +345,10 @@ final class KafkaMetadataLog private (
    * This method is thread-safe
    */
   override def deleteBeforeSnapshot(snapshotId: OffsetAndEpoch): Boolean = {
+    deleteBeforeSnapshot(snapshotId, UnknownReason)
+  }
+
+  private def deleteBeforeSnapshot(snapshotId: OffsetAndEpoch, reason: SnapshotDeletionReason): Boolean = {
     val (deleted, forgottenSnapshots) = snapshots synchronized {
       latestSnapshotId().asScala match {
         case Some(latestSnapshotId) if
@@ -348,14 +357,15 @@ final class KafkaMetadataLog private (
           snapshotId.offset <= latestSnapshotId.offset &&
           log.maybeIncrementLogStartOffset(snapshotId.offset, LogStartOffsetIncrementReason.SnapshotGenerated) =>
             // Delete all segments that have a "last offset" less than the log start offset
-            log.deleteOldSegments()
+            val deletedSegments = log.deleteOldSegments()
             // Remove older snapshots from the snapshots cache
-            (true, forgetSnapshotsBefore(snapshotId))
+            val forgottenSnapshots = forgetSnapshotsBefore(snapshotId)
+            (deletedSegments != 0 || forgottenSnapshots.nonEmpty, forgottenSnapshots)
         case _ =>
-            (false, mutable.TreeMap.empty[OffsetAndEpoch, Option[FileRawSnapshotReader]])
+          (false, mutable.TreeMap.empty[OffsetAndEpoch, Option[FileRawSnapshotReader]])
       }
     }
-    removeSnapshots(forgottenSnapshots)
+    removeSnapshots(forgottenSnapshots, reason)
     deleted
   }
 
@@ -404,21 +414,33 @@ final class KafkaMetadataLog private (
    * all cases.
    *
    * For the given predicate, we are testing if the snapshot identified by the first argument should be deleted.
+   * The predicate returns a Some with the reason if the snapshot should be deleted and a None if the snapshot
+   * should not be deleted
    */
-  private def cleanSnapshots(predicate: OffsetAndEpoch => Boolean): Boolean = {
-    if (snapshots.size < 2)
+  private def cleanSnapshots(predicate: OffsetAndEpoch => Option[SnapshotDeletionReason]): Boolean = {
+    if (snapshots.size < 2) {
       return false
+    }
 
     var didClean = false
     snapshots.keys.toSeq.sliding(2).foreach {
       case Seq(snapshot: OffsetAndEpoch, nextSnapshot: OffsetAndEpoch) =>
-        if (predicate(snapshot) && deleteBeforeSnapshot(nextSnapshot)) {
-          didClean = true
-        } else {
-          return didClean
+        predicate(snapshot) match {
+          case Some(reason) =>
+            if (deleteBeforeSnapshot(nextSnapshot, reason)) {
+              didClean = true
+            } else {
+              return didClean
+            }
+
+          case None =>
+            return didClean
+
         }
+
       case _ => false // Shouldn't get here with the sliding window
     }
+
     didClean
   }
 
@@ -427,13 +449,13 @@ final class KafkaMetadataLog private (
       return false
 
     // Keep deleting snapshots as long as the
-    def shouldClean(snapshotId: OffsetAndEpoch): Boolean = {
-      val now = time.milliseconds()
-      readSnapshotTimestamp(snapshotId).exists { timestamp =>
+    def shouldClean(snapshotId: OffsetAndEpoch): Option[SnapshotDeletionReason] = {
+      readSnapshotTimestamp(snapshotId).flatMap { timestamp =>
+        val now = time.milliseconds()
         if (now - timestamp > config.retentionMillis) {
-          true
+          Some(RetentionMsBreach(now, timestamp, config.retentionMillis))
         } else {
-          false
+          None
         }
       }
     }
@@ -450,13 +472,14 @@ final class KafkaMetadataLog private (
     var snapshotTotalSize: Long = snapshotSizes.values.sum
 
     // Keep deleting snapshots and segments as long as we exceed the retention size
-    def shouldClean(snapshotId: OffsetAndEpoch): Boolean = {
-      snapshotSizes.get(snapshotId).exists { snapshotSize =>
+    def shouldClean(snapshotId: OffsetAndEpoch): Option[SnapshotDeletionReason] = {
+      snapshotSizes.get(snapshotId).flatMap { snapshotSize =>
         if (log.size + snapshotTotalSize > config.retentionMaxBytes) {
+          val oldSnapshotTotalSize = snapshotTotalSize
           snapshotTotalSize -= snapshotSize
-          true
+          Some(RetentionSizeBreach(log.size, oldSnapshotTotalSize, config.retentionMaxBytes))
         } else {
-          false
+          None
         }
       }
     }
@@ -485,10 +508,11 @@ final class KafkaMetadataLog private (
    * given snapshots after some delay.
    */
   private def removeSnapshots(
-    expiredSnapshots: mutable.TreeMap[OffsetAndEpoch, Option[FileRawSnapshotReader]]
+    expiredSnapshots: mutable.TreeMap[OffsetAndEpoch, Option[FileRawSnapshotReader]],
+    reason: SnapshotDeletionReason,
   ): Unit = {
     expiredSnapshots.foreach { case (snapshotId, _) =>
-      info(s"Marking snapshot $snapshotId for deletion")
+      info(reason.reason(snapshotId))
       Snapshots.markForDelete(log.dir.toPath, snapshotId)
     }
 
@@ -516,32 +540,6 @@ final class KafkaMetadataLog private (
   }
 }
 
-object MetadataLogConfig {
-  def apply(config: AbstractConfig, maxBatchSizeInBytes: Int, maxFetchSizeInBytes: Int): MetadataLogConfig = {
-    new MetadataLogConfig(
-      config.getInt(KafkaConfig.MetadataLogSegmentBytesProp),
-      config.getInt(KafkaConfig.MetadataLogSegmentMinBytesProp),
-      config.getLong(KafkaConfig.MetadataLogSegmentMillisProp),
-      config.getLong(KafkaConfig.MetadataMaxRetentionBytesProp),
-      config.getLong(KafkaConfig.MetadataMaxRetentionMillisProp),
-      maxBatchSizeInBytes,
-      maxFetchSizeInBytes,
-      LogConfig.DEFAULT_FILE_DELETE_DELAY_MS,
-      config.getInt(KafkaConfig.NodeIdProp)
-    )
-  }
-}
-
-case class MetadataLogConfig(logSegmentBytes: Int,
-                             logSegmentMinBytes: Int,
-                             logSegmentMillis: Long,
-                             retentionMaxBytes: Long,
-                             retentionMillis: Long,
-                             maxBatchSizeInBytes: Int,
-                             maxFetchSizeInBytes: Int,
-                             fileDeleteDelayMs: Long,
-                             nodeId: Int)
-
 object KafkaMetadataLog extends Logging {
   def apply(
     topicPartition: TopicPartition,
@@ -677,4 +675,38 @@ object KafkaMetadataLog extends Logging {
       Snapshots.deleteIfExists(logDir, snapshotId)
     }
   }
+
+  private sealed trait SnapshotDeletionReason {
+    def reason(snapshotId: OffsetAndEpoch): String
+  }
+
+  private final case class RetentionMsBreach(now: Long, timestamp: Long, retentionMillis: Long) extends SnapshotDeletionReason {
+    override def reason(snapshotId: OffsetAndEpoch): String = {
+      s"""Marking snapshot $snapshotId for deletion because its timestamp ($timestamp) is now ($now) older than the
+          |retention ($retentionMillis)""".stripMargin
+    }
+  }
+
+  private final case class RetentionSizeBreach(
+    logSize: Long,
+    snapshotsSize: Long,
+    retentionMaxBytes: Long
+  ) extends SnapshotDeletionReason {
+    override def reason(snapshotId: OffsetAndEpoch): String = {
+      s"""Marking snapshot $snapshotId for deletion because the log size ($logSize) and snapshots size ($snapshotsSize)
+          |is greater than $retentionMaxBytes""".stripMargin
+    }
+  }
+
+  private final object FullTruncation extends SnapshotDeletionReason {
+    override def reason(snapshotId: OffsetAndEpoch): String = {
+      s"Marking snapshot $snapshotId for deletion because the partition was fully truncated"
+    }
+  }
+
+  private final object UnknownReason extends SnapshotDeletionReason {
+    override def reason(snapshotId: OffsetAndEpoch): String = {
+      s"Marking snapshot $snapshotId for deletion for unknown reason"
+    }
+  }
 }

From 1d6e0b872757317a23203e114303fbfb0a0a0c84 Mon Sep 17 00:00:00 2001
From: Artem Livshits <84364232+artemlivshits@users.noreply.github.com>
Date: Mon, 18 Mar 2024 19:08:55 -0700
Subject: [PATCH 171/258] KAFKA-16352: Txn may get get stuck in PrepareCommit
 or PrepareAbort state (#15524)

Now the removal of entries from the transactionsWithPendingMarkers map
checks the value and all pending marker operations keep the value along
with the operation state.  This way, the pending marker operation can
only delete the state it created and wouldn't accidentally delete the
state from a different epoch (which could lead to "stuck" transactions).

Reviewers: Justine Olshan <jolshan@confluent.io>
---
 .../TransactionMarkerChannelManager.scala     | 132 +++++++++++-------
 ...actionMarkerRequestCompletionHandler.scala |  38 ++---
 .../transaction/TransactionStateManager.scala |   1 +
 .../TransactionMarkerChannelManagerTest.scala |  83 ++++++++++-
 ...onMarkerRequestCompletionHandlerTest.scala |  23 +--
 5 files changed, 200 insertions(+), 77 deletions(-)

diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
index 5a72554c74d60..068dff4cca609 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
@@ -109,23 +109,43 @@ object TransactionMarkerChannelManager {
 
 }
 
-class TxnMarkerQueue(@volatile var destination: Node) {
+class TxnMarkerQueue(@volatile var destination: Node) extends Logging {
 
   // keep track of the requests per txn topic partition so we can easily clear the queue
   // during partition emigration
-  private val markersPerTxnTopicPartition = new ConcurrentHashMap[Int, BlockingQueue[TxnIdAndMarkerEntry]]().asScala
+  private val markersPerTxnTopicPartition = new ConcurrentHashMap[Int, BlockingQueue[PendingCompleteTxnAndMarkerEntry]]().asScala
 
-  def removeMarkersForTxnTopicPartition(partition: Int): Option[BlockingQueue[TxnIdAndMarkerEntry]] = {
+  def removeMarkersForTxnTopicPartition(partition: Int): Option[BlockingQueue[PendingCompleteTxnAndMarkerEntry]] = {
     markersPerTxnTopicPartition.remove(partition)
   }
 
-  def addMarkers(txnTopicPartition: Int, txnIdAndMarker: TxnIdAndMarkerEntry): Unit = {
-    val queue = CoreUtils.atomicGetOrUpdate(markersPerTxnTopicPartition, txnTopicPartition,
-        new LinkedBlockingQueue[TxnIdAndMarkerEntry]())
-    queue.add(txnIdAndMarker)
+  def addMarkers(txnTopicPartition: Int, pendingCompleteTxnAndMarker: PendingCompleteTxnAndMarkerEntry): Unit = {
+    val queue = CoreUtils.atomicGetOrUpdate(markersPerTxnTopicPartition, txnTopicPartition, {
+      // Note that this may get called more than once if threads have a close race while adding new queue.
+      info(s"Creating new marker queue for txn partition $txnTopicPartition to destination broker ${destination.id}")
+      new LinkedBlockingQueue[PendingCompleteTxnAndMarkerEntry]()
+    })
+    queue.add(pendingCompleteTxnAndMarker)
+
+    if (markersPerTxnTopicPartition.get(txnTopicPartition).orNull != queue) {
+      // This could happen if the queue got removed concurrently.
+      // Note that it could create an unexpected state when the queue is removed from
+      // removeMarkersForTxnTopicPartition, we could have:
+      //
+      // 1. [addMarkers] Retrieve queue.
+      // 2. [removeMarkersForTxnTopicPartition] Remove queue.
+      // 3. [removeMarkersForTxnTopicPartition] Iterate over queue, but not removeMarkersForTxn because queue is empty.
+      // 4. [addMarkers] Add markers to the queue.
+      //
+      // Now we've effectively removed the markers while transactionsWithPendingMarkers has an entry.
+      //
+      // While this could lead to an orphan entry in transactionsWithPendingMarkers, sending new markers
+      // will fix the state, so it shouldn't impact the state machine operation.
+      warn(s"Added $pendingCompleteTxnAndMarker to dead queue for txn partition $txnTopicPartition to destination broker ${destination.id}")
+    }
   }
 
-  def forEachTxnTopicPartition[B](f:(Int, BlockingQueue[TxnIdAndMarkerEntry]) => B): Unit =
+  def forEachTxnTopicPartition[B](f:(Int, BlockingQueue[PendingCompleteTxnAndMarkerEntry]) => B): Unit =
     markersPerTxnTopicPartition.forKeyValue { (partition, queue) =>
       if (!queue.isEmpty) f(partition, queue)
     }
@@ -187,17 +207,21 @@ class TransactionMarkerChannelManager(
   // visible for testing
   private[transaction] def queueForUnknownBroker = markersQueueForUnknownBroker
 
-  private[transaction] def addMarkersForBroker(broker: Node, txnTopicPartition: Int, txnIdAndMarker: TxnIdAndMarkerEntry): Unit = {
+  private[transaction] def addMarkersForBroker(broker: Node, txnTopicPartition: Int, pendingCompleteTxnAndMarker: PendingCompleteTxnAndMarkerEntry): Unit = {
     val brokerId = broker.id
 
     // we do not synchronize on the update of the broker node with the enqueuing,
     // since even if there is a race condition we will just retry
-    val brokerRequestQueue = CoreUtils.atomicGetOrUpdate(markersQueuePerBroker, brokerId,
-        new TxnMarkerQueue(broker))
+    val brokerRequestQueue = CoreUtils.atomicGetOrUpdate(markersQueuePerBroker, brokerId, {
+      // Note that this may get called more than once if threads have a close race while adding new queue.
+      info(s"Creating new marker queue map to destination broker $brokerId")
+      new TxnMarkerQueue(broker)
+    })
     brokerRequestQueue.destination = broker
-    brokerRequestQueue.addMarkers(txnTopicPartition, txnIdAndMarker)
+    brokerRequestQueue.addMarkers(txnTopicPartition, pendingCompleteTxnAndMarker)
 
-    trace(s"Added marker ${txnIdAndMarker.txnMarkerEntry} for transactional id ${txnIdAndMarker.txnId} to destination broker $brokerId")
+    trace(s"Added marker ${pendingCompleteTxnAndMarker.txnMarkerEntry} for transactional id" +
+      s" ${pendingCompleteTxnAndMarker.pendingCompleteTxn.transactionalId} to destination broker $brokerId")
   }
 
   private def retryLogAppends(): Unit = {
@@ -211,29 +235,28 @@ class TransactionMarkerChannelManager(
 
   override def generateRequests(): util.Collection[RequestAndCompletionHandler] = {
     retryLogAppends()
-    val txnIdAndMarkerEntries: util.List[TxnIdAndMarkerEntry] = new util.ArrayList[TxnIdAndMarkerEntry]()
+    val pendingCompleteTxnAndMarkerEntries = new util.ArrayList[PendingCompleteTxnAndMarkerEntry]()
     markersQueueForUnknownBroker.forEachTxnTopicPartition { case (_, queue) =>
-      queue.drainTo(txnIdAndMarkerEntries)
+      queue.drainTo(pendingCompleteTxnAndMarkerEntries)
     }
 
-    for (txnIdAndMarker: TxnIdAndMarkerEntry <- txnIdAndMarkerEntries.asScala) {
-      val transactionalId = txnIdAndMarker.txnId
-      val producerId = txnIdAndMarker.txnMarkerEntry.producerId
-      val producerEpoch = txnIdAndMarker.txnMarkerEntry.producerEpoch
-      val txnResult = txnIdAndMarker.txnMarkerEntry.transactionResult
-      val coordinatorEpoch = txnIdAndMarker.txnMarkerEntry.coordinatorEpoch
-      val topicPartitions = txnIdAndMarker.txnMarkerEntry.partitions.asScala.toSet
+    for (pendingCompleteTxnAndMarker: PendingCompleteTxnAndMarkerEntry <- pendingCompleteTxnAndMarkerEntries.asScala) {
+      val producerId = pendingCompleteTxnAndMarker.txnMarkerEntry.producerId
+      val producerEpoch = pendingCompleteTxnAndMarker.txnMarkerEntry.producerEpoch
+      val txnResult = pendingCompleteTxnAndMarker.txnMarkerEntry.transactionResult
+      val pendingCompleteTxn = pendingCompleteTxnAndMarker.pendingCompleteTxn
+      val topicPartitions = pendingCompleteTxnAndMarker.txnMarkerEntry.partitions.asScala.toSet
 
-      addTxnMarkersToBrokerQueue(transactionalId, producerId, producerEpoch, txnResult, coordinatorEpoch, topicPartitions)
+      addTxnMarkersToBrokerQueue(producerId, producerEpoch, txnResult, pendingCompleteTxn, topicPartitions)
     }
 
     val currentTimeMs = time.milliseconds()
     markersQueuePerBroker.values.map { brokerRequestQueue =>
-      val txnIdAndMarkerEntries = new util.ArrayList[TxnIdAndMarkerEntry]()
+      val pendingCompleteTxnAndMarkerEntries = new util.ArrayList[PendingCompleteTxnAndMarkerEntry]()
       brokerRequestQueue.forEachTxnTopicPartition { case (_, queue) =>
-        queue.drainTo(txnIdAndMarkerEntries)
+        queue.drainTo(pendingCompleteTxnAndMarkerEntries)
       }
-      (brokerRequestQueue.destination, txnIdAndMarkerEntries)
+      (brokerRequestQueue.destination, pendingCompleteTxnAndMarkerEntries)
     }.filter { case (_, entries) => !entries.isEmpty }.map { case (node, entries) =>
       val markersToSend = entries.asScala.map(_.txnMarkerEntry).asJava
       val requestCompletionHandler = new TransactionMarkerRequestCompletionHandler(node.id, txnStateManager, this, entries)
@@ -300,9 +323,12 @@ class TransactionMarkerChannelManager(
       txnMetadata,
       newMetadata)
 
-    transactionsWithPendingMarkers.put(transactionalId, pendingCompleteTxn)
-    addTxnMarkersToBrokerQueue(transactionalId, txnMetadata.producerId,
-      txnMetadata.producerEpoch, txnResult, coordinatorEpoch, txnMetadata.topicPartitions.toSet)
+    val prev = transactionsWithPendingMarkers.put(transactionalId, pendingCompleteTxn)
+    if (prev != null) {
+      info(s"Replaced an existing pending complete txn $prev with $pendingCompleteTxn while adding markers to send.")
+    }
+    addTxnMarkersToBrokerQueue(txnMetadata.producerId,
+      txnMetadata.producerEpoch, txnResult, pendingCompleteTxn, txnMetadata.topicPartitions.toSet)
     maybeWriteTxnCompletion(transactionalId)
   }
 
@@ -354,41 +380,42 @@ class TransactionMarkerChannelManager(
       txnLogAppend.newMetadata, appendCallback, _ == Errors.COORDINATOR_NOT_AVAILABLE, RequestLocal.NoCaching)
   }
 
-  def addTxnMarkersToBrokerQueue(transactionalId: String,
-                                 producerId: Long,
+  def addTxnMarkersToBrokerQueue(producerId: Long,
                                  producerEpoch: Short,
                                  result: TransactionResult,
-                                 coordinatorEpoch: Int,
+                                 pendingCompleteTxn: PendingCompleteTxn,
                                  topicPartitions: immutable.Set[TopicPartition]): Unit = {
-    val txnTopicPartition = txnStateManager.partitionFor(transactionalId)
+    val txnTopicPartition = txnStateManager.partitionFor(pendingCompleteTxn.transactionalId)
     val partitionsByDestination: immutable.Map[Option[Node], immutable.Set[TopicPartition]] = topicPartitions.groupBy { topicPartition: TopicPartition =>
       metadataCache.getPartitionLeaderEndpoint(topicPartition.topic, topicPartition.partition, interBrokerListenerName)
     }
 
+    val coordinatorEpoch = pendingCompleteTxn.coordinatorEpoch
     for ((broker: Option[Node], topicPartitions: immutable.Set[TopicPartition]) <- partitionsByDestination) {
       broker match {
         case Some(brokerNode) =>
           val marker = new TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, result, topicPartitions.toList.asJava)
-          val txnIdAndMarker = TxnIdAndMarkerEntry(transactionalId, marker)
+          val pendingCompleteTxnAndMarker = PendingCompleteTxnAndMarkerEntry(pendingCompleteTxn, marker)
 
           if (brokerNode == Node.noNode) {
             // if the leader of the partition is known but node not available, put it into an unknown broker queue
             // and let the sender thread to look for its broker and migrate them later
-            markersQueueForUnknownBroker.addMarkers(txnTopicPartition, txnIdAndMarker)
+            markersQueueForUnknownBroker.addMarkers(txnTopicPartition, pendingCompleteTxnAndMarker)
           } else {
-            addMarkersForBroker(brokerNode, txnTopicPartition, txnIdAndMarker)
+            addMarkersForBroker(brokerNode, txnTopicPartition, pendingCompleteTxnAndMarker)
           }
 
         case None =>
+          val transactionalId = pendingCompleteTxn.transactionalId
           txnStateManager.getTransactionState(transactionalId) match {
             case Left(error) =>
               info(s"Encountered $error trying to fetch transaction metadata for $transactionalId with coordinator epoch $coordinatorEpoch; cancel sending markers to its partition leaders")
-              transactionsWithPendingMarkers.remove(transactionalId)
+              transactionsWithPendingMarkers.remove(transactionalId, pendingCompleteTxn)
 
             case Right(Some(epochAndMetadata)) =>
               if (epochAndMetadata.coordinatorEpoch != coordinatorEpoch) {
                 info(s"The cached metadata has changed to $epochAndMetadata (old coordinator epoch is $coordinatorEpoch) since preparing to send markers; cancel sending markers to its partition leaders")
-                transactionsWithPendingMarkers.remove(transactionalId)
+                transactionsWithPendingMarkers.remove(transactionalId, pendingCompleteTxn)
               } else {
                 // if the leader of the partition is unknown, skip sending the txn marker since
                 // the partition is likely to be deleted already
@@ -419,25 +446,34 @@ class TransactionMarkerChannelManager(
 
   def removeMarkersForTxnTopicPartition(txnTopicPartitionId: Int): Unit = {
     markersQueueForUnknownBroker.removeMarkersForTxnTopicPartition(txnTopicPartitionId).foreach { queue =>
-      for (entry: TxnIdAndMarkerEntry <- queue.asScala)
-        removeMarkersForTxnId(entry.txnId)
+      for (entry <- queue.asScala) {
+        info(s"Removing $entry for txn partition $txnTopicPartitionId to destination broker -1")
+        removeMarkersForTxn(entry.pendingCompleteTxn)
+      }
     }
 
-    markersQueuePerBroker.foreach { case(_, brokerQueue) =>
+    markersQueuePerBroker.foreach { case(brokerId, brokerQueue) =>
       brokerQueue.removeMarkersForTxnTopicPartition(txnTopicPartitionId).foreach { queue =>
-        for (entry: TxnIdAndMarkerEntry <- queue.asScala)
-          removeMarkersForTxnId(entry.txnId)
+        for (entry <- queue.asScala) {
+          info(s"Removing $entry for txn partition $txnTopicPartitionId to destination broker $brokerId")
+          removeMarkersForTxn(entry.pendingCompleteTxn)
+        }
       }
     }
   }
 
-  def removeMarkersForTxnId(transactionalId: String): Unit = {
-    transactionsWithPendingMarkers.remove(transactionalId)
+  def removeMarkersForTxn(pendingCompleteTxn: PendingCompleteTxn): Unit = {
+    val transactionalId = pendingCompleteTxn.transactionalId
+    val removed = transactionsWithPendingMarkers.remove(transactionalId, pendingCompleteTxn)
+    if (!removed) {
+      val current = transactionsWithPendingMarkers.get(transactionalId)
+      if (current != null) {
+        info(s"Failed to remove pending marker entry $current trying to remove $pendingCompleteTxn")
+      }
+    }
   }
 }
 
-case class TxnIdAndMarkerEntry(txnId: String, txnMarkerEntry: TxnMarkerEntry)
-
 case class PendingCompleteTxn(transactionalId: String,
                               coordinatorEpoch: Int,
                               txnMetadata: TransactionMetadata,
@@ -451,3 +487,5 @@ case class PendingCompleteTxn(transactionalId: String,
       s"newMetadata=$newMetadata)"
   }
 }
+
+case class PendingCompleteTxnAndMarkerEntry(pendingCompleteTxn: PendingCompleteTxn, txnMarkerEntry: TxnMarkerEntry)
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
index 7a59139b17c76..d95dabab6c356 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
@@ -29,7 +29,7 @@ import scala.jdk.CollectionConverters._
 class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                                                 txnStateManager: TransactionStateManager,
                                                 txnMarkerChannelManager: TransactionMarkerChannelManager,
-                                                txnIdAndMarkerEntries: java.util.List[TxnIdAndMarkerEntry]) extends RequestCompletionHandler with Logging {
+                                                pendingCompleteTxnAndMarkerEntries: java.util.List[PendingCompleteTxnAndMarkerEntry]) extends RequestCompletionHandler with Logging {
 
   this.logIdent = "[Transaction Marker Request Completion Handler " + brokerId + "]: "
 
@@ -39,22 +39,23 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
     if (response.wasDisconnected) {
       trace(s"Cancelled request with header $requestHeader due to node ${response.destination} being disconnected")
 
-      for (txnIdAndMarker <- txnIdAndMarkerEntries.asScala) {
-        val transactionalId = txnIdAndMarker.txnId
-        val txnMarker = txnIdAndMarker.txnMarkerEntry
+      for (pendingCompleteTxnAndMarker <- pendingCompleteTxnAndMarkerEntries.asScala) {
+        val pendingCompleteTxn = pendingCompleteTxnAndMarker.pendingCompleteTxn
+        val transactionalId = pendingCompleteTxn.transactionalId
+        val txnMarker = pendingCompleteTxnAndMarker.txnMarkerEntry
 
         txnStateManager.getTransactionState(transactionalId) match {
 
           case Left(Errors.NOT_COORDINATOR) =>
             info(s"I am no longer the coordinator for $transactionalId; cancel sending transaction markers $txnMarker to the brokers")
 
-            txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+            txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
 
           case Left(Errors.COORDINATOR_LOAD_IN_PROGRESS) =>
             info(s"I am loading the transaction partition that contains $transactionalId which means the current markers have to be obsoleted; " +
               s"cancel sending transaction markers $txnMarker to the brokers")
 
-            txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+            txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
 
           case Left(unexpectedError) =>
             throw new IllegalStateException(s"Unhandled error $unexpectedError when fetching current transaction state")
@@ -69,17 +70,16 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
               info(s"Transaction coordinator epoch for $transactionalId has changed from ${txnMarker.coordinatorEpoch} to " +
                 s"${epochAndMetadata.coordinatorEpoch}; cancel sending transaction markers $txnMarker to the brokers")
 
-              txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+              txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
             } else {
               // re-enqueue the markers with possibly new destination brokers
               trace(s"Re-enqueuing ${txnMarker.transactionResult} transaction markers for transactional id $transactionalId " +
                 s"under coordinator epoch ${txnMarker.coordinatorEpoch}")
 
-              txnMarkerChannelManager.addTxnMarkersToBrokerQueue(transactionalId,
-                txnMarker.producerId,
+              txnMarkerChannelManager.addTxnMarkersToBrokerQueue(txnMarker.producerId,
                 txnMarker.producerEpoch,
                 txnMarker.transactionResult,
-                txnMarker.coordinatorEpoch,
+                pendingCompleteTxn,
                 txnMarker.partitions.asScala.toSet)
             }
         }
@@ -90,9 +90,10 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
       val writeTxnMarkerResponse = response.responseBody.asInstanceOf[WriteTxnMarkersResponse]
 
       val responseErrors = writeTxnMarkerResponse.errorsByProducerId
-      for (txnIdAndMarker <- txnIdAndMarkerEntries.asScala) {
-        val transactionalId = txnIdAndMarker.txnId
-        val txnMarker = txnIdAndMarker.txnMarkerEntry
+      for (pendingCompleteTxnAndMarker <- pendingCompleteTxnAndMarkerEntries.asScala) {
+        val pendingCompleteTxn = pendingCompleteTxnAndMarker.pendingCompleteTxn
+        val transactionalId = pendingCompleteTxn.transactionalId
+        val txnMarker = pendingCompleteTxnAndMarker.txnMarkerEntry
         val errors = responseErrors.get(txnMarker.producerId)
 
         if (errors == null)
@@ -102,13 +103,13 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
           case Left(Errors.NOT_COORDINATOR) =>
             info(s"I am no longer the coordinator for $transactionalId; cancel sending transaction markers $txnMarker to the brokers")
 
-            txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+            txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
 
           case Left(Errors.COORDINATOR_LOAD_IN_PROGRESS) =>
             info(s"I am loading the transaction partition that contains $transactionalId which means the current markers have to be obsoleted; " +
               s"cancel sending transaction markers $txnMarker to the brokers")
 
-            txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+            txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
 
           case Left(unexpectedError) =>
             throw new IllegalStateException(s"Unhandled error $unexpectedError when fetching current transaction state")
@@ -127,7 +128,7 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
               info(s"Transaction coordinator epoch for $transactionalId has changed from ${txnMarker.coordinatorEpoch} to " +
                 s"${epochAndMetadata.coordinatorEpoch}; cancel sending transaction markers $txnMarker to the brokers")
 
-              txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+              txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
               abortSending = true
             } else {
               txnMetadata.inLock {
@@ -161,7 +162,7 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                       info(s"Sending $transactionalId's transaction marker for partition $topicPartition has permanently failed with error ${error.exceptionName} " +
                         s"with the current coordinator epoch ${epochAndMetadata.coordinatorEpoch}; cancel sending any more transaction markers $txnMarker to the brokers")
 
-                      txnMarkerChannelManager.removeMarkersForTxnId(transactionalId)
+                      txnMarkerChannelManager.removeMarkersForTxn(pendingCompleteTxn)
                       abortSending = true
 
                     case Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT |
@@ -187,11 +188,10 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
 
                 // re-enqueue with possible new leaders of the partitions
                 txnMarkerChannelManager.addTxnMarkersToBrokerQueue(
-                  transactionalId,
                   txnMarker.producerId,
                   txnMarker.producerEpoch,
                   txnMarker.transactionResult,
-                  txnMarker.coordinatorEpoch,
+                  pendingCompleteTxn,
                   retryPartitions.toSet)
               } else {
                 txnMarkerChannelManager.maybeWriteTxnCompletion(transactionalId)
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
index 355e848980746..813fd75be36b6 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
@@ -557,6 +557,7 @@ class TransactionStateManager(brokerId: Int,
           loadingPartitions.remove(partitionAndLeaderEpoch)
 
           transactionsPendingForCompletion.foreach { txnTransitMetadata =>
+            info(s"Sending txn markers for $txnTransitMetadata after loading partition $partitionId")
             sendTxnMarkers(txnTransitMetadata.coordinatorEpoch, txnTransitMetadata.result,
               txnTransitMetadata.txnMetadata, txnTransitMetadata.transitMetadata)
           }
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
index de58f8ed7fa86..3356c4f9e372c 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
@@ -34,7 +34,7 @@ import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
 import org.mockito.ArgumentMatchers.any
 import org.mockito.{ArgumentCaptor, ArgumentMatchers}
-import org.mockito.Mockito.{mock, mockConstruction, times, verify, verifyNoMoreInteractions, when}
+import org.mockito.Mockito.{clearInvocations, mock, mockConstruction, times, verify, verifyNoMoreInteractions, when}
 
 import scala.jdk.CollectionConverters._
 import scala.collection.mutable
@@ -59,6 +59,7 @@ class TransactionMarkerChannelManagerTest {
   private val txnTopicPartition1 = 0
   private val txnTopicPartition2 = 1
   private val coordinatorEpoch = 0
+  private val coordinatorEpoch2 = 1
   private val txnTimeoutMs = 0
   private val txnResult = TransactionResult.COMMIT
   private val txnMetadata1 = new TransactionMetadata(transactionalId1, producerId1, producerId1, producerEpoch, lastProducerEpoch,
@@ -177,6 +178,86 @@ class TransactionMarkerChannelManagerTest {
       any())
   }
 
+  @Test
+  def shouldNotLoseTxnCompletionAfterLoad(): Unit = {
+    mockCache()
+
+    val expectedTransition = txnMetadata2.prepareComplete(time.milliseconds())
+
+    when(metadataCache.getPartitionLeaderEndpoint(
+      ArgumentMatchers.eq(partition1.topic),
+      ArgumentMatchers.eq(partition1.partition),
+      any())
+    ).thenReturn(Some(broker1))
+
+    // Build a successful client response.
+    val header = new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1)
+    val successfulResponse = new WriteTxnMarkersResponse(
+      Collections.singletonMap(producerId2: java.lang.Long, Collections.singletonMap(partition1, Errors.NONE)))
+    val successfulClientResponse = new ClientResponse(header, null, null,
+      time.milliseconds(), time.milliseconds(), false, null, null,
+      successfulResponse)
+
+    // Build a disconnected client response.
+    val disconnectedClientResponse = new ClientResponse(header, null, null,
+      time.milliseconds(), time.milliseconds(), true, null, null,
+      null)
+
+    // Test matrix to cover various scenarios:
+    val clientResponses = Seq(successfulClientResponse, disconnectedClientResponse)
+    val getTransactionStateResponses = Seq(
+      // NOT_COORDINATOR error case
+      Left(Errors.NOT_COORDINATOR),
+      // COORDINATOR_LOAD_IN_PROGRESS
+      Left(Errors.COORDINATOR_LOAD_IN_PROGRESS),
+      // "Newly loaded" transaction state with the new epoch.
+      Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch2, txnMetadata2)))
+    )
+
+    clientResponses.foreach { clientResponse =>
+      getTransactionStateResponses.foreach { getTransactionStateResponse =>
+        // Reset data from previous iteration.
+        txnMetadata2.topicPartitions.add(partition1)
+        clearInvocations(txnStateManager)
+        // Send out markers for a transaction before load.
+        channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult,
+          txnMetadata2, expectedTransition)
+
+        // Drain the marker to make it "in-flight".
+        val requests1 = channelManager.generateRequests().asScala
+        assertEquals(1, requests1.size)
+
+        // Simulate a partition load:
+        // 1. Remove the markers from the channel manager.
+        // 2. Simulate the corresponding test case scenario.
+        // 3. Add the markers back to the channel manager.
+        channelManager.removeMarkersForTxnTopicPartition(txnTopicPartition2)
+        when(txnStateManager.getTransactionState(ArgumentMatchers.eq(transactionalId2)))
+          .thenReturn(getTransactionStateResponse)
+        channelManager.addTxnMarkersToSend(coordinatorEpoch2, txnResult,
+          txnMetadata2, expectedTransition)
+
+        // Complete the marker from the previous epoch.
+        requests1.head.handler.onComplete(clientResponse)
+
+        // Now drain and complete the marker from the new epoch.
+        when(txnStateManager.getTransactionState(ArgumentMatchers.eq(transactionalId2)))
+          .thenReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch2, txnMetadata2))))
+        val requests2 = channelManager.generateRequests().asScala
+        assertEquals(1, requests2.size)
+        requests2.head.handler.onComplete(successfulClientResponse)
+
+        verify(txnStateManager).appendTransactionToLog(
+          ArgumentMatchers.eq(transactionalId2),
+          ArgumentMatchers.eq(coordinatorEpoch2),
+          ArgumentMatchers.eq(expectedTransition),
+          capturedErrorsCallback.capture(),
+          any(),
+          any())
+      }
+    }
+  }
+
   @Test
   def shouldGenerateEmptyMapWhenNoRequestsOutstanding(): Unit = {
     assertTrue(channelManager.generateRequests().isEmpty)
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
index aecf6542f7d7f..1004915f46cb9 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
@@ -18,7 +18,6 @@ package kafka.coordinator.transaction
 
 import java.{lang, util}
 import java.util.Arrays.asList
-
 import org.apache.kafka.clients.ClientResponse
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
@@ -43,18 +42,19 @@ class TransactionMarkerRequestCompletionHandlerTest {
   private val coordinatorEpoch = 0
   private val txnResult = TransactionResult.COMMIT
   private val topicPartition = new TopicPartition("topic1", 0)
-  private val txnIdAndMarkers = asList(
-      TxnIdAndMarkerEntry(transactionalId, new WriteTxnMarkersRequest.TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, txnResult, asList(topicPartition))))
-
   private val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, lastProducerEpoch,
     txnTimeoutMs, PrepareCommit, mutable.Set[TopicPartition](topicPartition), 0L, 0L)
+  private val pendingCompleteTxnAndMarkers = asList(
+    PendingCompleteTxnAndMarkerEntry(
+      PendingCompleteTxn(transactionalId, coordinatorEpoch, txnMetadata, txnMetadata.prepareComplete(42)),
+      new WriteTxnMarkersRequest.TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, txnResult, asList(topicPartition))))
 
   private val markerChannelManager: TransactionMarkerChannelManager =
     mock(classOf[TransactionMarkerChannelManager])
 
   private val txnStateManager: TransactionStateManager = mock(classOf[TransactionStateManager])
 
-  private val handler = new TransactionMarkerRequestCompletionHandler(brokerId, txnStateManager, markerChannelManager, txnIdAndMarkers)
+  private val handler = new TransactionMarkerRequestCompletionHandler(brokerId, txnStateManager, markerChannelManager, pendingCompleteTxnAndMarkers)
 
   private def mockCache(): Unit = {
     when(txnStateManager.partitionFor(transactionalId))
@@ -70,8 +70,9 @@ class TransactionMarkerRequestCompletionHandlerTest {
     handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
       null, null, 0, 0, true, null, null, null))
 
-    verify(markerChannelManager).addTxnMarkersToBrokerQueue(transactionalId,
-      producerId, producerEpoch, txnResult, coordinatorEpoch, Set[TopicPartition](topicPartition))
+    verify(markerChannelManager).addTxnMarkersToBrokerQueue(producerId,
+      producerEpoch, txnResult, pendingCompleteTxnAndMarkers.get(0).pendingCompleteTxn,
+      Set[TopicPartition](topicPartition))
   }
 
   @Test
@@ -193,8 +194,9 @@ class TransactionMarkerRequestCompletionHandlerTest {
       null, null, 0, 0, false, null, null, response))
 
     assertEquals(txnMetadata.topicPartitions, mutable.Set[TopicPartition](topicPartition))
-    verify(markerChannelManager).addTxnMarkersToBrokerQueue(transactionalId,
-      producerId, producerEpoch, txnResult, coordinatorEpoch, Set[TopicPartition](topicPartition))
+    verify(markerChannelManager).addTxnMarkersToBrokerQueue(producerId,
+      producerEpoch, txnResult, pendingCompleteTxnAndMarkers.get(0).pendingCompleteTxn,
+      Set[TopicPartition](topicPartition))
   }
 
   private def verifyThrowIllegalStateExceptionOnError(error: Errors) = {
@@ -222,7 +224,8 @@ class TransactionMarkerRequestCompletionHandlerTest {
   private def verifyRemoveDelayedOperationOnError(error: Errors): Unit = {
 
     var removed = false
-    when(markerChannelManager.removeMarkersForTxnId(transactionalId))
+    val pendingCompleteTxn = pendingCompleteTxnAndMarkers.get(0).pendingCompleteTxn
+    when(markerChannelManager.removeMarkersForTxn(pendingCompleteTxn))
       .thenAnswer(_ => removed = true)
 
     val response = new WriteTxnMarkersResponse(createProducerIdErrorMap(error))

From bf3f088c944763d8418764064f593d0bf06fbcc3 Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Tue, 19 Mar 2024 23:00:30 +0800
Subject: [PATCH 172/258] KAFKA-16341 fix the LogValidator for non-compressed
 type (#15476)

- Fix the verifying logic. If it's LOG_APPEND_TIME, we choose the offset of the first record. Else, we choose the record with the maxTimeStamp.
- rename the shallowOffsetOfMaxTimestamp to offsetOfMaxTimestamp

Reviewers: Jun Rao <junrao@gmail.com>, Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/common/record/MemoryRecords.java    | 22 ++++-----
 .../common/record/MemoryRecordsBuilder.java   |  6 +--
 .../record/MemoryRecordsBuilderTest.java      | 10 ++---
 .../common/record/MemoryRecordsTest.java      |  6 +--
 core/src/main/scala/kafka/log/LocalLog.scala  |  4 +-
 .../src/main/scala/kafka/log/LogCleaner.scala |  2 +-
 .../src/main/scala/kafka/log/UnifiedLog.scala |  2 +-
 .../admin/ListOffsetsIntegrationTest.scala    | 45 +++++++++++++------
 .../scala/unit/kafka/log/LocalLogTest.scala   |  2 +-
 .../scala/unit/kafka/log/LogSegmentTest.scala |  4 +-
 .../unit/kafka/log/LogValidatorTest.scala     | 30 +++++--------
 .../storage/internals/log/LogSegment.java     | 10 ++---
 .../storage/internals/log/LogValidator.java   | 16 +++----
 13 files changed, 84 insertions(+), 75 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java
index 3aa233c34e948..3ba60b09b30df 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecords.java
@@ -209,7 +209,7 @@ private static FilterResult filterTo(TopicPartition partition, Iterable<MutableR
                                 partition, batch.lastOffset(), maxRecordBatchSize, filteredBatchSize);
 
                         MemoryRecordsBuilder.RecordsInfo info = builder.info();
-                        filterResult.updateRetainedBatchMetadata(info.maxTimestamp, info.shallowOffsetOfMaxTimestamp,
+                        filterResult.updateRetainedBatchMetadata(info.maxTimestamp, info.offsetOfMaxTimestamp,
                             maxOffset, retainedRecords.size(), filteredBatchSize);
                     }
                 }
@@ -399,7 +399,7 @@ public static class FilterResult {
         private int bytesRetained = 0;
         private long maxOffset = -1L;
         private long maxTimestamp = RecordBatch.NO_TIMESTAMP;
-        private long shallowOffsetOfMaxTimestamp = -1L;
+        private long offsetOfMaxTimestamp = -1L;
 
         private FilterResult(ByteBuffer outputBuffer) {
             this.outputBuffer = outputBuffer;
@@ -411,21 +411,21 @@ private void updateRetainedBatchMetadata(MutableRecordBatch retainedBatch, int n
                     retainedBatch.lastOffset(), numMessagesInBatch, bytesRetained);
         }
 
-        private void updateRetainedBatchMetadata(long maxTimestamp, long shallowOffsetOfMaxTimestamp, long maxOffset,
-                                                int messagesRetained, int bytesRetained) {
-            validateBatchMetadata(maxTimestamp, shallowOffsetOfMaxTimestamp, maxOffset);
+        private void updateRetainedBatchMetadata(long maxTimestamp, long offsetOfMaxTimestamp, long maxOffset,
+                                                 int messagesRetained, int bytesRetained) {
+            validateBatchMetadata(maxTimestamp, offsetOfMaxTimestamp, maxOffset);
             if (maxTimestamp > this.maxTimestamp) {
                 this.maxTimestamp = maxTimestamp;
-                this.shallowOffsetOfMaxTimestamp = shallowOffsetOfMaxTimestamp;
+                this.offsetOfMaxTimestamp = offsetOfMaxTimestamp;
             }
             this.maxOffset = Math.max(maxOffset, this.maxOffset);
             this.messagesRetained += messagesRetained;
             this.bytesRetained += bytesRetained;
         }
 
-        private void validateBatchMetadata(long maxTimestamp, long shallowOffsetOfMaxTimestamp, long maxOffset) {
-            if (maxTimestamp != RecordBatch.NO_TIMESTAMP && shallowOffsetOfMaxTimestamp < 0)
-                throw new IllegalArgumentException("shallowOffset undefined for maximum timestamp " + maxTimestamp);
+        private void validateBatchMetadata(long maxTimestamp, long offsetOfMaxTimestamp, long maxOffset) {
+            if (maxTimestamp != RecordBatch.NO_TIMESTAMP && offsetOfMaxTimestamp < 0)
+                throw new IllegalArgumentException("offset undefined for maximum timestamp " + maxTimestamp);
             if (maxOffset < 0)
                 throw new IllegalArgumentException("maxOffset undefined");
         }
@@ -458,8 +458,8 @@ public long maxTimestamp() {
             return maxTimestamp;
         }
 
-        public long shallowOffsetOfMaxTimestamp() {
-            return shallowOffsetOfMaxTimestamp;
+        public long offsetOfMaxTimestamp() {
+            return offsetOfMaxTimestamp;
         }
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
index 3663c5ea7e339..6b53ee415959e 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/MemoryRecordsBuilder.java
@@ -851,12 +851,12 @@ private long nextSequentialOffset() {
 
     public static class RecordsInfo {
         public final long maxTimestamp;
-        public final long shallowOffsetOfMaxTimestamp;
+        public final long offsetOfMaxTimestamp;
 
         public RecordsInfo(long maxTimestamp,
-                           long shallowOffsetOfMaxTimestamp) {
+                           long offsetOfMaxTimestamp) {
             this.maxTimestamp = maxTimestamp;
-            this.shallowOffsetOfMaxTimestamp = shallowOffsetOfMaxTimestamp;
+            this.offsetOfMaxTimestamp = offsetOfMaxTimestamp;
         }
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
index 81aa162feee1f..eaaa95ff673cf 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsBuilderTest.java
@@ -378,7 +378,7 @@ public void buildUsingLogAppendTime(Args args) {
         MemoryRecordsBuilder.RecordsInfo info = builder.info();
         assertEquals(logAppendTime, info.maxTimestamp);
         // When logAppendTime is used, the first offset of the batch will be the offset of maxTimestamp
-        assertEquals(0L, info.shallowOffsetOfMaxTimestamp);
+        assertEquals(0L, info.offsetOfMaxTimestamp);
 
         for (RecordBatch batch : records.batches()) {
             if (magic == MAGIC_VALUE_V0) {
@@ -414,9 +414,9 @@ public void buildUsingCreateTime(Args args) {
 
         if (magic == MAGIC_VALUE_V0)
             // in MAGIC_VALUE_V0's case, we don't have timestamp info in records, so always return -1.
-            assertEquals(-1L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(-1L, info.offsetOfMaxTimestamp);
         else
-            assertEquals(1L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(1L, info.offsetOfMaxTimestamp);
 
         int i = 0;
         long[] expectedTimestamps = new long[] {0L, 2L, 1L};
@@ -495,10 +495,10 @@ public void writePastLimit(Args args) {
         MemoryRecordsBuilder.RecordsInfo info = builder.info();
         if (magic == MAGIC_VALUE_V0) {
             assertEquals(-1, info.maxTimestamp);
-            assertEquals(-1L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(-1L, info.offsetOfMaxTimestamp);
         } else {
             assertEquals(2L, info.maxTimestamp);
-            assertEquals(2L, info.shallowOffsetOfMaxTimestamp);
+            assertEquals(2L, info.offsetOfMaxTimestamp);
         }
 
         long i = 0L;
diff --git a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
index 50821af841c42..9e688fc3ab695 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/MemoryRecordsTest.java
@@ -352,7 +352,7 @@ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
                     assertEquals(0, filterResult.messagesRetained());
                     assertEquals(DefaultRecordBatch.RECORD_BATCH_OVERHEAD, filterResult.bytesRetained());
                     assertEquals(12, filterResult.maxTimestamp());
-                    assertEquals(baseOffset + 1, filterResult.shallowOffsetOfMaxTimestamp());
+                    assertEquals(baseOffset + 1, filterResult.offsetOfMaxTimestamp());
 
                     // Verify filtered records
                     filtered.flip();
@@ -413,7 +413,7 @@ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
         assertEquals(0, filterResult.messagesRetained());
         assertEquals(DefaultRecordBatch.RECORD_BATCH_OVERHEAD, filterResult.bytesRetained());
         assertEquals(timestamp, filterResult.maxTimestamp());
-        assertEquals(baseOffset, filterResult.shallowOffsetOfMaxTimestamp());
+        assertEquals(baseOffset, filterResult.offsetOfMaxTimestamp());
         assertTrue(filterResult.outputBuffer().position() > 0);
 
         // Verify filtered records
@@ -893,7 +893,7 @@ public void testFilterTo(Args args) {
         assertEquals(filtered.limit(), result.bytesRetained());
         if (magic > RecordBatch.MAGIC_VALUE_V0) {
             assertEquals(20L, result.maxTimestamp());
-            assertEquals(4L, result.shallowOffsetOfMaxTimestamp());
+            assertEquals(4L, result.offsetOfMaxTimestamp());
         }
 
         MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
diff --git a/core/src/main/scala/kafka/log/LocalLog.scala b/core/src/main/scala/kafka/log/LocalLog.scala
index b2121f5312d7b..a91bfae739e05 100644
--- a/core/src/main/scala/kafka/log/LocalLog.scala
+++ b/core/src/main/scala/kafka/log/LocalLog.scala
@@ -406,8 +406,8 @@ class LocalLog(@volatile private var _dir: File,
     }
   }
 
-  private[log] def append(lastOffset: Long, largestTimestamp: Long, shallowOffsetOfMaxTimestamp: Long, records: MemoryRecords): Unit = {
-    segments.activeSegment.append(lastOffset, largestTimestamp, shallowOffsetOfMaxTimestamp, records)
+  private[log] def append(lastOffset: Long, largestTimestamp: Long, offsetOfMaxTimestamp: Long, records: MemoryRecords): Unit = {
+    segments.activeSegment.append(lastOffset, largestTimestamp, offsetOfMaxTimestamp, records)
     updateLogEndOffset(lastOffset + 1)
   }
 
diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala
index b653f40b287a8..35dadc8672f08 100644
--- a/core/src/main/scala/kafka/log/LogCleaner.scala
+++ b/core/src/main/scala/kafka/log/LogCleaner.scala
@@ -812,7 +812,7 @@ private[log] class Cleaner(val id: Int,
         val retained = MemoryRecords.readableRecords(outputBuffer)
         // it's OK not to hold the Log's lock in this case, because this segment is only accessed by other threads
         // after `Log.replaceSegments` (which acquires the lock) is called
-        dest.append(result.maxOffset, result.maxTimestamp, result.shallowOffsetOfMaxTimestamp, retained)
+        dest.append(result.maxOffset, result.maxTimestamp, result.offsetOfMaxTimestamp, retained)
         throttler.maybeThrottle(outputBuffer.limit())
       }
 
diff --git a/core/src/main/scala/kafka/log/UnifiedLog.scala b/core/src/main/scala/kafka/log/UnifiedLog.scala
index f6198ce9d219e..0fdc236e720c2 100644
--- a/core/src/main/scala/kafka/log/UnifiedLog.scala
+++ b/core/src/main/scala/kafka/log/UnifiedLog.scala
@@ -819,7 +819,7 @@ class UnifiedLog(@volatile var logStartOffset: Long,
 
             validRecords = validateAndOffsetAssignResult.validatedRecords
             appendInfo.setMaxTimestamp(validateAndOffsetAssignResult.maxTimestampMs)
-            appendInfo.setOffsetOfMaxTimestamp(validateAndOffsetAssignResult.shallowOffsetOfMaxTimestampMs)
+            appendInfo.setOffsetOfMaxTimestamp(validateAndOffsetAssignResult.offsetOfMaxTimestampMs)
             appendInfo.setLastOffset(offset.value - 1)
             appendInfo.setRecordValidationStats(validateAndOffsetAssignResult.recordValidationStats)
             if (config.messageTimestampType == TimestampType.LOG_APPEND_TIME)
diff --git a/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala b/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
index e5e22e9dff97d..5362a1d5e35c9 100644
--- a/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/ListOffsetsIntegrationTest.scala
@@ -68,9 +68,7 @@ class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
     verifyListOffsets()
 
     // test LogAppendTime case
-    val props: Properties = new Properties()
-    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
-    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    setUpForLogAppendTimeCase()
     produceMessagesInOneBatch("gzip", topicNameWithCustomConfigs)
     // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
     // So in this one batch test, it'll be the first offset 0
@@ -79,9 +77,30 @@ class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
-  def testThreeRecordsInSeparateBatch(quorum: String): Unit = {
+  def testThreeNonCompressedRecordsInOneBatch(quorum: String): Unit = {
+    produceMessagesInOneBatch()
+    verifyListOffsets()
+
+    // test LogAppendTime case
+    setUpForLogAppendTimeCase()
+    produceMessagesInOneBatch(topic=topicNameWithCustomConfigs)
+    // In LogAppendTime's case, if the timestamps are the same, we choose the offset of the first record
+    // thus, the maxTimestampOffset should be the first record of the batch.
+    // So in this one batch test, it'll be the first offset which is 0
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 0)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
+  @ValueSource(strings = Array("zk", "kraft"))
+  def testThreeNonCompressedRecordsInSeparateBatch(quorum: String): Unit = {
     produceMessagesInSeparateBatch()
     verifyListOffsets()
+
+    // test LogAppendTime case
+    setUpForLogAppendTimeCase()
+    produceMessagesInSeparateBatch(topic = topicNameWithCustomConfigs)
+    // In LogAppendTime's case, if the timestamp is different, it should be the last one
+    verifyListOffsets(topic = topicNameWithCustomConfigs, 2)
   }
 
   // The message conversion test only run in ZK mode because KRaft mode doesn't support "inter.broker.protocol.version" < 3.0
@@ -93,9 +112,7 @@ class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
     verifyListOffsets()
 
     // test LogAppendTime case
-    val props: Properties = new Properties()
-    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
-    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    setUpForLogAppendTimeCase()
     produceMessagesInOneBatch(topic = topicNameWithCustomConfigs)
     // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
     // So in this one batch test, it'll be the first offset 0
@@ -111,9 +128,7 @@ class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
     verifyListOffsets()
 
     // test LogAppendTime case
-    val props: Properties = new Properties()
-    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
-    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    setUpForLogAppendTimeCase()
     produceMessagesInSeparateBatch(topic = topicNameWithCustomConfigs)
     // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
     // So in this separate batch test, it'll be the last offset 2
@@ -147,15 +162,19 @@ class ListOffsetsIntegrationTest extends KafkaServerTestHarness {
     verifyListOffsets()
 
     // test LogAppendTime case
-    val props: Properties = new Properties()
-    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
-    createTopicWithConfig(topicNameWithCustomConfigs, props)
+    setUpForLogAppendTimeCase()
     produceMessagesInSeparateBatch("gzip", topicNameWithCustomConfigs)
     // In LogAppendTime's case, the maxTimestampOffset should be the first message of the batch.
     // So in this separate batch test, it'll be the last offset 2
     verifyListOffsets(topic = topicNameWithCustomConfigs, 2)
   }
 
+  private def setUpForLogAppendTimeCase(): Unit = {
+    val props: Properties = new Properties()
+    props.setProperty(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG, "LogAppendTime")
+    createTopicWithConfig(topicNameWithCustomConfigs, props)
+  }
+
   private def createOldMessageFormatBrokers(): Unit = {
     setOldMessageFormat = true
     recreateBrokers(reconfigure = true, startup = true)
diff --git a/core/src/test/scala/unit/kafka/log/LocalLogTest.scala b/core/src/test/scala/unit/kafka/log/LocalLogTest.scala
index 29b5fd34f9091..bffd41156b357 100644
--- a/core/src/test/scala/unit/kafka/log/LocalLogTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LocalLogTest.scala
@@ -100,7 +100,7 @@ class LocalLogTest {
                             initialOffset: Long = 0L): Unit = {
     log.append(lastOffset = initialOffset + records.size - 1,
       largestTimestamp = records.head.timestamp,
-      shallowOffsetOfMaxTimestamp = initialOffset,
+      offsetOfMaxTimestamp = initialOffset,
       records = MemoryRecords.withRecords(initialOffset, CompressionType.NONE, 0, records.toList : _*))
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
index b4e278c380230..bb0c85a858360 100644
--- a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
@@ -86,10 +86,10 @@ class LogSegmentTest {
   def testAppendForLogSegmentOffsetOverflowException(baseOffset: Long, largestOffset: Long): Unit = {
     val seg = createSegment(baseOffset)
     val currentTime = Time.SYSTEM.milliseconds()
-    val shallowOffsetOfMaxTimestamp = largestOffset
+    val offsetOfMaxTimestamp = largestOffset
     val memoryRecords = records(0, "hello")
     assertThrows(classOf[LogSegmentOffsetOverflowException], () => {
-      seg.append(largestOffset, currentTime, shallowOffsetOfMaxTimestamp, memoryRecords)
+      seg.append(largestOffset, currentTime, offsetOfMaxTimestamp, memoryRecords)
     })
   }
 
diff --git a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
index ac6152b9b153c..53b385c62e86c 100644
--- a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
@@ -173,9 +173,9 @@ class LogValidatorTest {
     assertEquals(now, validatedResults.maxTimestampMs, s"Max timestamp should be $now")
     assertFalse(validatedResults.messageSizeMaybeChanged, "Message size should not have been changed")
 
-    // we index from last offset in version 2 instead of base offset
-    val expectedMaxTimestampOffset = if (magic >= RecordBatch.MAGIC_VALUE_V2) 2 else 0
-    assertEquals(expectedMaxTimestampOffset, validatedResults.shallowOffsetOfMaxTimestampMs,
+    // If it's LOG_APPEND_TIME, the offset will be the offset of the first record
+    val expectedMaxTimestampOffset = 0
+    assertEquals(expectedMaxTimestampOffset, validatedResults.offsetOfMaxTimestampMs,
       s"The offset of max timestamp should be $expectedMaxTimestampOffset")
     verifyRecordValidationStats(validatedResults.recordValidationStats, numConvertedRecords = 0, records,
       compressed = false)
@@ -219,7 +219,7 @@ class LogValidatorTest {
       "MessageSet should still valid")
     assertEquals(now, validatedResults.maxTimestampMs,
       s"Max timestamp should be $now")
-    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+    assertEquals(0, validatedResults.offsetOfMaxTimestampMs,
       s"The offset of max timestamp should be 0 if logAppendTime is used")
     assertTrue(validatedResults.messageSizeMaybeChanged,
       "Message size may have been changed")
@@ -271,7 +271,7 @@ class LogValidatorTest {
       "MessageSet should still valid")
     assertEquals(now, validatedResults.maxTimestampMs,
       s"Max timestamp should be $now")
-    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+    assertEquals(0, validatedResults.offsetOfMaxTimestampMs,
       s"The offset of max timestamp should be 0 if logAppendTime is used")
     assertFalse(validatedResults.messageSizeMaybeChanged,
       "Message size should not have been changed")
@@ -404,14 +404,8 @@ class LogValidatorTest {
     assertEquals(now + 1, validatingResults.maxTimestampMs,
       s"Max timestamp should be ${now + 1}")
 
-    val expectedShallowOffsetOfMaxTimestamp = if (magic >= RecordVersion.V2.value) {
-      // v2 records are always batched, even when not compressed.
-      // the shallow offset of max timestamp is the last offset of the batch
-      recordList.size - 1
-    } else {
-      1
-    }
-    assertEquals(expectedShallowOffsetOfMaxTimestamp, validatingResults.shallowOffsetOfMaxTimestampMs,
+    val expectedOffsetOfMaxTimestamp = 1
+    assertEquals(expectedOffsetOfMaxTimestamp, validatingResults.offsetOfMaxTimestampMs,
       s"Offset of max timestamp should be 1")
 
     assertFalse(validatingResults.messageSizeMaybeChanged,
@@ -486,7 +480,7 @@ class LogValidatorTest {
     }
     assertEquals(now + 1, validatingResults.maxTimestampMs,
       s"Max timestamp should be ${now + 1}")
-    assertEquals(1, validatingResults.shallowOffsetOfMaxTimestampMs,
+    assertEquals(1, validatingResults.offsetOfMaxTimestampMs,
       "Offset of max timestamp should be 1")
     assertTrue(validatingResults.messageSizeMaybeChanged,
       "Message size should have been changed")
@@ -538,7 +532,7 @@ class LogValidatorTest {
     }
     assertEquals(validatedResults.maxTimestampMs, RecordBatch.NO_TIMESTAMP,
       s"Max timestamp should be ${RecordBatch.NO_TIMESTAMP}")
-    assertEquals(-1, validatedResults.shallowOffsetOfMaxTimestampMs,
+    assertEquals(-1, validatedResults.offsetOfMaxTimestampMs,
       s"Offset of max timestamp should be -1")
     assertTrue(validatedResults.messageSizeMaybeChanged, "Message size should have been changed")
 
@@ -585,7 +579,7 @@ class LogValidatorTest {
       assertEquals(RecordBatch.NO_SEQUENCE, batch.baseSequence)
     }
     assertEquals(timestamp, validatedResults.maxTimestampMs)
-    assertEquals(0, validatedResults.shallowOffsetOfMaxTimestampMs,
+    assertEquals(0, validatedResults.offsetOfMaxTimestampMs,
       s"Offset of max timestamp should be 0 when multiple records having the same max timestamp.")
     assertTrue(validatedResults.messageSizeMaybeChanged, "Message size should have been changed")
 
@@ -657,8 +651,8 @@ class LogValidatorTest {
     }
     assertEquals(now + 1, validatedResults.maxTimestampMs, s"Max timestamp should be ${now + 1}")
 
-    val expectedShallowOffsetOfMaxTimestamp = 1
-    assertEquals(expectedShallowOffsetOfMaxTimestamp, validatedResults.shallowOffsetOfMaxTimestampMs,
+    val expectedOffsetOfMaxTimestamp = 1
+    assertEquals(expectedOffsetOfMaxTimestamp, validatedResults.offsetOfMaxTimestampMs,
       s"Offset of max timestamp should be 1")
     assertFalse(validatedResults.messageSizeMaybeChanged, "Message size should not have been changed")
 
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
index d42bf1124560d..464858b4cd783 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogSegment.java
@@ -232,17 +232,17 @@ private boolean canConvertToRelativeOffset(long offset) throws IOException {
      *
      * @param largestOffset The last offset in the message set
      * @param largestTimestampMs The largest timestamp in the message set.
-     * @param shallowOffsetOfMaxTimestamp The offset of the message that has the largest timestamp in the messages to append.
+     * @param offsetOfMaxTimestamp The offset of the message that has the largest timestamp in the messages to append.
      * @param records The log entries to append.
      * @throws LogSegmentOffsetOverflowException if the largest offset causes index offset overflow
      */
     public void append(long largestOffset,
                        long largestTimestampMs,
-                       long shallowOffsetOfMaxTimestamp,
+                       long offsetOfMaxTimestamp,
                        MemoryRecords records) throws IOException {
         if (records.sizeInBytes() > 0) {
-            LOGGER.trace("Inserting {} bytes at end offset {} at position {} with largest timestamp {} at shallow offset {}",
-                records.sizeInBytes(), largestOffset, log.sizeInBytes(), largestTimestampMs, shallowOffsetOfMaxTimestamp);
+            LOGGER.trace("Inserting {} bytes at end offset {} at position {} with largest timestamp {} at offset {}",
+                records.sizeInBytes(), largestOffset, log.sizeInBytes(), largestTimestampMs, offsetOfMaxTimestamp);
             int physicalPosition = log.sizeInBytes();
             if (physicalPosition == 0)
                 rollingBasedTimestamp = OptionalLong.of(largestTimestampMs);
@@ -254,7 +254,7 @@ public void append(long largestOffset,
             LOGGER.trace("Appended {} to {} at end offset {}", appendedBytes, log.file(), largestOffset);
             // Update the in memory max timestamp and corresponding offset.
             if (largestTimestampMs > maxTimestampSoFar()) {
-                maxTimestampAndOffsetSoFar = new TimestampOffset(largestTimestampMs, shallowOffsetOfMaxTimestamp);
+                maxTimestampAndOffsetSoFar = new TimestampOffset(largestTimestampMs, offsetOfMaxTimestamp);
             }
             // append an entry to the index (if needed)
             if (bytesSinceLastIndexEntry > indexIntervalBytes) {
diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
index 0cf9cd1c60f73..9aa1e06633b6d 100644
--- a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
+++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogValidator.java
@@ -68,17 +68,17 @@ public static class ValidationResult {
         public final long logAppendTimeMs;
         public final MemoryRecords validatedRecords;
         public final long maxTimestampMs;
-        public final long shallowOffsetOfMaxTimestampMs;
+        public final long offsetOfMaxTimestampMs;
         public final boolean messageSizeMaybeChanged;
         public final RecordValidationStats recordValidationStats;
 
         public ValidationResult(long logAppendTimeMs, MemoryRecords validatedRecords, long maxTimestampMs,
-                                long shallowOffsetOfMaxTimestampMs, boolean messageSizeMaybeChanged,
+                                long offsetOfMaxTimestampMs, boolean messageSizeMaybeChanged,
                                 RecordValidationStats recordValidationStats) {
             this.logAppendTimeMs = logAppendTimeMs;
             this.validatedRecords = validatedRecords;
             this.maxTimestampMs = maxTimestampMs;
-            this.shallowOffsetOfMaxTimestampMs = shallowOffsetOfMaxTimestampMs;
+            this.offsetOfMaxTimestampMs = offsetOfMaxTimestampMs;
             this.messageSizeMaybeChanged = messageSizeMaybeChanged;
             this.recordValidationStats = recordValidationStats;
         }
@@ -149,7 +149,7 @@ public LogValidator(MemoryRecords records,
      * avoid expensive re-compression.
      *
      * Returns a ValidationAndOffsetAssignResult containing the validated message set, maximum timestamp, the offset
-     * of the shallow message with the max timestamp and a boolean indicating whether the message sizes may have changed.
+     * of the message with the max timestamp and a boolean indicating whether the message sizes may have changed.
      */
     public ValidationResult validateMessagesAndAssignOffsets(PrimitiveRef.LongRef offsetCounter,
                                                              MetricsRecorder metricsRecorder,
@@ -232,7 +232,7 @@ private ValidationResult convertAndAssignOffsetsNonCompressed(LongRef offsetCoun
             now,
             convertedRecords,
             info.maxTimestamp,
-            info.shallowOffsetOfMaxTimestamp,
+            info.offsetOfMaxTimestamp,
             true,
             recordValidationStats);
     }
@@ -296,10 +296,6 @@ public ValidationResult assignOffsetsNonCompressed(LongRef offsetCounter,
             offsetOfMaxTimestamp = initialOffset;
         }
 
-        if (toMagic >= RecordBatch.MAGIC_VALUE_V2) {
-            offsetOfMaxTimestamp = offsetCounter.value - 1;
-        }
-
         return new ValidationResult(
             now,
             records,
@@ -480,7 +476,7 @@ private ValidationResult buildRecordsAndAssignOffsets(LongRef offsetCounter,
             logAppendTime,
             records,
             info.maxTimestamp,
-            info.shallowOffsetOfMaxTimestamp,
+            info.offsetOfMaxTimestamp,
             true,
             recordValidationStats);
     }

From c66d66dc67b3aacb60f438bb4c5c1c132e8be4f2 Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Tue, 19 Mar 2024 20:48:41 +0000
Subject: [PATCH 173/258] KAFKA-16367; Full ConsumerGroupHeartbeat response
 must be sent when full request is received (#15533)

This patch fixes a bug in the logic which decides when a full ConsumerGroupHeartbeat response must be returned to the client. Prior to it, the logic only relies on the `ownedTopicPartitions` field to check whether the response was a full response. This is not enough because `ownedTopicPartitions` is also set in different situations. This patch changes the logic to check `ownedTopicPartitions`, `subscribedTopicNames` and `rebalanceTimeoutMs` as they are the only three non optional fields.

Reviewers: Lianet Magrans <lianetmr@gmail.com>, Jeff Kim <jeff.kim@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../group/GroupMetadataManager.java           |  11 +-
 .../group/GroupMetadataManagerTest.java       | 114 +++++++++++++++---
 2 files changed, 105 insertions(+), 20 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 48f0618c55d34..0a789fa963089 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -1227,10 +1227,13 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
             .setHeartbeatIntervalMs(consumerGroupHeartbeatIntervalMs);
 
         // The assignment is only provided in the following cases:
-        // 1. The member reported its owned partitions;
-        // 2. The member just joined or rejoined to group (epoch equals to zero);
-        // 3. The member's assignment has been updated.
-        if (ownedTopicPartitions != null || memberEpoch == 0 || hasAssignedPartitionsChanged(member, updatedMember)) {
+        // 1. The member sent a full request. It does so when joining or rejoining the group with zero
+        //    as the member epoch; or on any errors (e.g. timeout). We use all the non-optional fields
+        //    (rebalanceTimeoutMs, subscribedTopicNames and ownedTopicPartitions) to detect a full request
+        //    as those must be set in a full request.
+        // 2. The member's assignment has been updated.
+        boolean isFullRequest = memberEpoch == 0 || (rebalanceTimeoutMs != -1 && subscribedTopicNames != null && ownedTopicPartitions != null);
+        if (isFullRequest || hasAssignedPartitionsChanged(member, updatedMember)) {
             response.setAssignment(createResponseAssignment(updatedMember));
         }
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index e9304407cdd20..43703059915a6 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -1650,6 +1650,102 @@ public void testShouldThrowFencedInstanceIdExceptionWhenStaticMemberWithDifferen
                 .setTopicPartitions(Collections.emptyList())));
     }
 
+    @Test
+    public void testConsumerGroupHeartbeatFullResponse() {
+        String groupId = "fooup";
+        String memberId = Uuid.randomUuid().toString();
+
+        Uuid fooTopicId = Uuid.randomUuid();
+        String fooTopicName = "foo";
+
+        // Create a context with an empty consumer group.
+        MockPartitionAssignor assignor = new MockPartitionAssignor("range");
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withAssignors(Collections.singletonList(assignor))
+            .withMetadataImage(new MetadataImageBuilder()
+                .addTopic(fooTopicId, fooTopicName, 2)
+                .addRacks()
+                .build())
+            .build();
+
+        // Prepare new assignment for the group.
+        assignor.prepareGroupAssignment(new GroupAssignment(
+            new HashMap<String, MemberAssignment>() {
+                {
+                    put(memberId, new MemberAssignment(mkAssignment(
+                        mkTopicAssignment(fooTopicId, 0, 1)
+                    )));
+                }
+            }
+        ));
+
+        CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result;
+
+        // A full response should be sent back on joining.
+        result = context.consumerGroupHeartbeat(
+            new ConsumerGroupHeartbeatRequestData()
+                .setGroupId(groupId)
+                .setMemberId(memberId)
+                .setMemberEpoch(0)
+                .setRebalanceTimeoutMs(5000)
+                .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
+                .setServerAssignor("range")
+                .setTopicPartitions(Collections.emptyList()));
+
+        assertResponseEquals(
+            new ConsumerGroupHeartbeatResponseData()
+                .setMemberId(memberId)
+                .setMemberEpoch(1)
+                .setHeartbeatIntervalMs(5000)
+                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
+                    .setTopicPartitions(Collections.singletonList(
+                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
+                            .setTopicId(fooTopicId)
+                            .setPartitions(Arrays.asList(0, 1))))),
+            result.response()
+        );
+
+        // Otherwise, a partial response should be sent back.
+        result = context.consumerGroupHeartbeat(
+            new ConsumerGroupHeartbeatRequestData()
+                .setGroupId(groupId)
+                .setMemberId(memberId)
+                .setMemberEpoch(result.response().memberEpoch()));
+
+        assertResponseEquals(
+            new ConsumerGroupHeartbeatResponseData()
+                .setMemberId(memberId)
+                .setMemberEpoch(1)
+                .setHeartbeatIntervalMs(5000),
+            result.response()
+        );
+
+        // A full response should be sent back when the member sends
+        // a full request again.
+        result = context.consumerGroupHeartbeat(
+            new ConsumerGroupHeartbeatRequestData()
+                .setGroupId(groupId)
+                .setMemberId(memberId)
+                .setMemberEpoch(result.response().memberEpoch())
+                .setRebalanceTimeoutMs(5000)
+                .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
+                .setServerAssignor("range")
+                .setTopicPartitions(Collections.emptyList()));
+
+        assertResponseEquals(
+            new ConsumerGroupHeartbeatResponseData()
+                .setMemberId(memberId)
+                .setMemberEpoch(1)
+                .setHeartbeatIntervalMs(5000)
+                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
+                    .setTopicPartitions(Collections.singletonList(
+                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
+                            .setTopicId(fooTopicId)
+                            .setPartitions(Arrays.asList(0, 1))))),
+            result.response()
+        );
+    }
+
     @Test
     public void testReconciliationProcess() {
         String groupId = "fooup";
@@ -1904,16 +2000,7 @@ public void testReconciliationProcess() {
             new ConsumerGroupHeartbeatResponseData()
                 .setMemberId(memberId1)
                 .setMemberEpoch(11)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0, 1)),
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(barTopicId)
-                            .setPartitions(Collections.singletonList(0))
-                    ))),
+                .setHeartbeatIntervalMs(5000),
             result.response()
         );
 
@@ -3057,12 +3144,7 @@ public void testRebalanceTimeoutLifecycle() {
             new ConsumerGroupHeartbeatResponseData()
                 .setMemberId(memberId1)
                 .setMemberEpoch(2)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0, 1))))),
+                .setHeartbeatIntervalMs(5000),
             result.response()
         );
 

From 12a1d85362bfc200714bb2d76e6da7f5af1f82dd Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Wed, 20 Mar 2024 10:36:25 +0800
Subject: [PATCH 174/258] KAFKA-12187 replace assertTrue(obj instanceof X) with
 assertInstanceOf (#15512)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../clients/CommonClientConfigsTest.java      |   5 +-
 .../clients/admin/KafkaAdminClientTest.java   |  68 ++++------
 .../AbortTransactionHandlerTest.java          |   4 +-
 .../internals/CoordinatorStrategyTest.java    |  10 +-
 ...DeleteConsumerGroupOffsetsHandlerTest.java |   4 +-
 .../DeleteConsumerGroupsHandlerTest.java      |   4 +-
 .../DescribeConsumerGroupsHandlerTest.java    |   4 +-
 .../DescribeProducersHandlerTest.java         |  10 +-
 .../DescribeTransactionsHandlerTest.java      |   4 +-
 .../internals/FenceProducersHandlerTest.java  |   4 +-
 .../ListConsumerGroupOffsetsHandlerTest.java  |   5 +-
 .../PartitionLeaderStrategyTest.java          |  10 +-
 ...veMembersFromConsumerGroupHandlerTest.java |   4 +-
 .../ConsumerPartitionAssignorTest.java        |  13 +-
 .../clients/consumer/KafkaConsumerTest.java   |   7 +-
 .../internals/AbstractCoordinatorTest.java    |  42 +++----
 .../internals/AsyncKafkaConsumerTest.java     |   2 +-
 .../internals/CommitRequestManagerTest.java   |   3 +-
 .../internals/ConsumerCoordinatorTest.java    |  31 ++---
 .../internals/ConsumerNetworkClientTest.java  |  13 +-
 .../CoordinatorRequestManagerTest.java        |   3 +-
 .../consumer/internals/OffsetFetcherTest.java |   3 +-
 .../internals/OffsetsRequestManagerTest.java  |   6 +-
 .../TopicMetadataRequestManagerTest.java      |   5 +-
 .../consumer/internals/WakeupTriggerTest.java |  12 +-
 .../clients/producer/KafkaProducerTest.java   |  15 +--
 .../clients/producer/MockProducerTest.java    |   3 +-
 .../producer/internals/SenderTest.java        |  27 ++--
 .../internals/TransactionManagerTest.java     | 119 ++++++++----------
 .../apache/kafka/common/KafkaFutureTest.java  |   5 +-
 .../kafka/common/TopicPartitionTest.java      |   6 +-
 .../common/network/ChannelBuildersTest.java   |   3 +-
 .../common/network/SslTransportLayerTest.java |   5 +-
 .../common/record/DefaultRecordBatchTest.java |   5 +-
 .../common/requests/RequestContextTest.java   |   3 +-
 .../SaslAuthenticatorFailureDelayTest.java    |   3 +-
 .../authenticator/SaslAuthenticatorTest.java  |  17 +--
 .../OAuthBearerLoginCallbackHandlerTest.java  |   5 +-
 .../internals/secured/OAuthBearerTest.java    |  16 +--
 ...UnsecuredValidatorCallbackHandlerTest.java |  12 +-
 .../common/security/ssl/SslFactoryTest.java   |   5 +-
 .../serialization/ListDeserializerTest.java   |  10 +-
 .../serialization/ListSerializerTest.java     |  10 +-
 .../ClientTelemetryReporterTest.java          |  11 +-
 .../java/org/apache/kafka/test/TestUtils.java |   3 +-
 .../apache/kafka/connect/data/ValuesTest.java |  29 ++---
 .../kafka/connect/mirror/MirrorUtilsTest.java |   6 +-
 .../ConnectWorkerIntegrationTest.java         |   3 +-
 .../ExactlyOnceSourceIntegrationTest.java     |   3 +-
 .../connect/runtime/WorkerConnectorTest.java  |   3 +-
 .../kafka/connect/runtime/WorkerTest.java     |   5 +-
 .../distributed/DistributedHerderTest.java    |  31 ++---
 .../runtime/errors/ErrorReporterTest.java     |   3 +-
 .../runtime/isolation/PluginsTest.java        |  14 +--
 .../standalone/StandaloneHerderTest.java      |  17 +--
 .../storage/KafkaOffsetBackingStoreTest.java  |   3 +-
 .../kafka/connect/util/TopicAdminTest.java    |   7 +-
 .../kafka/connect/transforms/FlattenTest.java |   8 +-
 ...ribeTopicPartitionsRequestHandlerTest.java |   3 +-
 .../ReplicationControlManagerTest.java        |   5 +-
 .../raft/KafkaRaftClientSnapshotTest.java     |   3 +-
 .../kafka/raft/RaftClientTestContext.java     |  32 ++---
 .../util/InterBrokerSendThreadTest.java       |   7 +-
 .../kafka/streams/StreamsConfigTest.java      |  11 +-
 .../SessionWindowedDeserializerTest.java      |   7 +-
 .../SessionWindowedSerializerTest.java        |   6 +-
 .../kstream/TimeWindowedDeserializerTest.java |   6 +-
 .../kstream/TimeWindowedSerializerTest.java   |   6 +-
 .../streams/kstream/WindowedSerdesTest.java   |  18 +--
 .../internals/InternalStreamsBuilderTest.java |  13 +-
 .../internals/KStreamFlatTransformTest.java   |   4 +-
 .../KStreamFlatTransformValuesTest.java       |   4 +-
 .../internals/DefaultStateUpdaterTest.java    |   3 +-
 .../InternalTopologyBuilderTest.java          |  11 +-
 .../internals/ProcessorStateManagerTest.java  |   3 +-
 .../processor/internals/StreamThreadTest.java |  46 +++----
 .../StandbyTaskAssignorFactoryTest.java       |  10 +-
 .../tools/LeaderElectionCommandErrorTest.java |   3 +-
 .../tools/LeaderElectionCommandTest.java      |   3 +-
 .../tools/MetadataQuorumCommandTest.java      |  13 +-
 .../apache/kafka/tools/TopicCommandTest.java  |   5 +-
 .../consumer/ConsoleConsumerOptionsTest.java  |   9 +-
 82 files changed, 450 insertions(+), 472 deletions(-)

diff --git a/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java b/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java
index c38f3e6a030b9..d53630f121153 100644
--- a/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/CommonClientConfigsTest.java
@@ -35,6 +35,7 @@
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
 import static org.apache.kafka.common.config.ConfigDef.ValidString.in;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -136,7 +137,7 @@ public void testMetricsReporters() {
         TestConfig config = new TestConfig(Collections.emptyMap());
         List<MetricsReporter> reporters = CommonClientConfigs.metricsReporters("clientId", config);
         assertEquals(1, reporters.size());
-        assertTrue(reporters.get(0) instanceof JmxReporter);
+        assertInstanceOf(JmxReporter.class, reporters.get(0));
 
         config = new TestConfig(Collections.singletonMap(CommonClientConfigs.AUTO_INCLUDE_JMX_REPORTER_CONFIG, "false"));
         reporters = CommonClientConfigs.metricsReporters("clientId", config);
@@ -145,7 +146,7 @@ public void testMetricsReporters() {
         config = new TestConfig(Collections.singletonMap(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, JmxReporter.class.getName()));
         reporters = CommonClientConfigs.metricsReporters("clientId", config);
         assertEquals(1, reporters.size());
-        assertTrue(reporters.get(0) instanceof JmxReporter);
+        assertInstanceOf(JmxReporter.class, reporters.get(0));
 
         Map<String, String> props = new HashMap<>();
         props.put(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, JmxReporter.class.getName() + "," + MyJmxReporter.class.getName());
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
index 43d391a220efc..71e802365fd0a 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
@@ -271,6 +271,7 @@
 import static org.apache.kafka.common.message.ListPartitionReassignmentsResponseData.OngoingTopicReassignment;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -294,7 +295,7 @@ public void testDefaultApiTimeoutAndRequestTimeoutConflicts() {
         final AdminClientConfig config = newConfMap(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "500");
         KafkaException exception = assertThrows(KafkaException.class,
             () -> KafkaAdminClient.createInternal(config, null));
-        assertTrue(exception.getCause() instanceof ConfigException);
+        assertInstanceOf(ConfigException.class, exception.getCause());
     }
 
     @Test
@@ -499,7 +500,7 @@ public void testAdminClientFailureWhenClosed() {
         ExecutionException e = assertThrows(ExecutionException.class, () -> env.adminClient().createTopics(
                 singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
                 new CreateTopicsOptions().timeoutMs(10000)).all().get());
-        assertTrue(e.getCause() instanceof IllegalStateException,
+        assertInstanceOf(IllegalStateException.class, e.getCause(),
                 "Expected an IllegalStateException error, but got " + Utils.stackTrace(e));
     }
 
@@ -1456,38 +1457,38 @@ private void callAdminClientApisAndExpectAnAuthenticationError(AdminClientUnitTe
         ExecutionException e = assertThrows(ExecutionException.class, () -> env.adminClient().createTopics(
             singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
             new CreateTopicsOptions().timeoutMs(10000)).all().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         Map<String, NewPartitions> counts = new HashMap<>();
         counts.put("my_topic", NewPartitions.increaseTo(3));
         counts.put("other_topic", NewPartitions.increaseTo(3, asList(asList(2), asList(3))));
         e = assertThrows(ExecutionException.class, () -> env.adminClient().createPartitions(counts).all().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         e = assertThrows(ExecutionException.class, () -> env.adminClient().createAcls(asList(ACL1, ACL2)).all().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         e = assertThrows(ExecutionException.class, () -> env.adminClient().describeAcls(FILTER1).values().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         e = assertThrows(ExecutionException.class, () -> env.adminClient().deleteAcls(asList(FILTER1, FILTER2)).all().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         e = assertThrows(ExecutionException.class, () -> env.adminClient().describeConfigs(
             singleton(new ConfigResource(ConfigResource.Type.BROKER, "0"))).all().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
     }
 
     private void callClientQuotasApisAndExpectAnAuthenticationError(AdminClientUnitTestEnv env) {
         ExecutionException e = assertThrows(ExecutionException.class,
             () -> env.adminClient().describeClientQuotas(ClientQuotaFilter.all()).entities().get());
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
 
         ClientQuotaEntity entity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.USER, "user"));
@@ -1495,7 +1496,7 @@ private void callClientQuotasApisAndExpectAnAuthenticationError(AdminClientUnitT
         e = assertThrows(ExecutionException.class,
             () -> env.adminClient().alterClientQuotas(asList(alteration)).all().get());
 
-        assertTrue(e.getCause() instanceof AuthenticationException,
+        assertInstanceOf(AuthenticationException.class, e.getCause(),
             "Expected an authentication error, but got " + Utils.stackTrace(e));
     }
 
@@ -1877,7 +1878,7 @@ public void testDescribeLogDirs() throws ExecutionException, InterruptedExceptio
                     env.cluster().nodeById(0));
             final DescribeLogDirsResult errorResult = env.adminClient().describeLogDirs(brokers);
             ExecutionException exception = assertThrows(ExecutionException.class, () -> errorResult.allDescriptions().get());
-            assertTrue(exception.getCause() instanceof ClusterAuthorizationException);
+            assertInstanceOf(ClusterAuthorizationException.class, exception.getCause());
 
             // Empty results with an error with version >= 3
             env.kafkaClient().prepareResponseFrom(
@@ -1885,7 +1886,7 @@ public void testDescribeLogDirs() throws ExecutionException, InterruptedExceptio
                     env.cluster().nodeById(0));
             final DescribeLogDirsResult errorResult2 = env.adminClient().describeLogDirs(brokers);
             exception = assertThrows(ExecutionException.class, () -> errorResult2.allDescriptions().get());
-            assertTrue(exception.getCause() instanceof UnknownServerException);
+            assertInstanceOf(UnknownServerException.class, exception.getCause());
         }
     }
 
@@ -1941,7 +1942,7 @@ public void testDescribeLogDirsWithVolumeBytes() throws ExecutionException, Inte
                     env.cluster().nodeById(0));
             final DescribeLogDirsResult errorResult = env.adminClient().describeLogDirs(brokers);
             ExecutionException exception = assertThrows(ExecutionException.class, () -> errorResult.allDescriptions().get());
-            assertTrue(exception.getCause() instanceof ClusterAuthorizationException);
+            assertInstanceOf(ClusterAuthorizationException.class, exception.getCause());
 
             // Empty results with an error with version >= 3
             env.kafkaClient().prepareResponseFrom(
@@ -1949,7 +1950,7 @@ public void testDescribeLogDirsWithVolumeBytes() throws ExecutionException, Inte
                     env.cluster().nodeById(0));
             final DescribeLogDirsResult errorResult2 = env.adminClient().describeLogDirs(brokers);
             exception = assertThrows(ExecutionException.class, () -> errorResult2.allDescriptions().get());
-            assertTrue(exception.getCause() instanceof UnknownServerException);
+            assertInstanceOf(UnknownServerException.class, exception.getCause());
         }
     }
 
@@ -2168,14 +2169,9 @@ public void testCreatePartitions() throws Exception {
             KafkaFuture<Void> myTopicResult = values.get("my_topic");
             myTopicResult.get();
             KafkaFuture<Void> otherTopicResult = values.get("other_topic");
-            try {
-                otherTopicResult.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e0) {
-                assertTrue(e0.getCause() instanceof InvalidTopicException);
-                InvalidTopicException e = (InvalidTopicException) e0.getCause();
-                assertEquals("some detailed reason", e.getMessage());
-            }
+            assertEquals("some detailed reason",
+                assertInstanceOf(InvalidTopicException.class,
+                    assertThrows(ExecutionException.class, otherTopicResult::get).getCause()).getMessage());
         }
     }
 
@@ -2434,30 +2430,18 @@ public void testDeleteRecords() throws Exception {
 
             // "offset out of range" failure on records deletion for partition 1
             KafkaFuture<DeletedRecords> myTopicPartition1Result = values.get(myTopicPartition1);
-            try {
-                myTopicPartition1Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e0) {
-                assertTrue(e0.getCause() instanceof OffsetOutOfRangeException);
-            }
+            assertInstanceOf(OffsetOutOfRangeException.class,
+                assertThrows(ExecutionException.class, myTopicPartition1Result::get).getCause());
 
             // not authorized to delete records for partition 2
             KafkaFuture<DeletedRecords> myTopicPartition2Result = values.get(myTopicPartition2);
-            try {
-                myTopicPartition2Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e1) {
-                assertTrue(e1.getCause() instanceof TopicAuthorizationException);
-            }
+            assertInstanceOf(TopicAuthorizationException.class,
+                assertThrows(ExecutionException.class, myTopicPartition2Result::get).getCause());
 
             // the response does not contain a result for partition 3
             KafkaFuture<DeletedRecords> myTopicPartition3Result = values.get(myTopicPartition3);
-            try {
-                myTopicPartition3Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e1) {
-                assertTrue(e1.getCause() instanceof ApiException);
-            }
+            assertInstanceOf(ApiException.class,
+                assertThrows(ExecutionException.class, myTopicPartition3Result::get).getCause());
         }
     }
 
@@ -4736,8 +4720,8 @@ public void testRemoveMembersFromGroup() throws Exception {
                     new RemoveMembersFromConsumerGroupOptions()
             );
             ExecutionException exception = assertThrows(ExecutionException.class, () -> partialFailureResults.all().get());
-            assertTrue(exception.getCause() instanceof KafkaException);
-            assertTrue(exception.getCause().getCause() instanceof UnknownMemberIdException);
+            assertInstanceOf(KafkaException.class, exception.getCause());
+            assertInstanceOf(UnknownMemberIdException.class, exception.getCause().getCause());
 
             // Return with success for "removeAll" scenario
             // 1 prepare response for AdminClient.describeConsumerGroups
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandlerTest.java
index 68c850583d6ae..e0e4e9b202523 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandlerTest.java
@@ -38,9 +38,9 @@
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class AbortTransactionHandlerTest {
     private final LogContext logContext = new LogContext();
@@ -214,7 +214,7 @@ private void assertFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(topicPartition), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(topicPartition)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(topicPartition));
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/CoordinatorStrategyTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/CoordinatorStrategyTest.java
index dd83c6be19999..fc52e9e6717ed 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/CoordinatorStrategyTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/CoordinatorStrategyTest.java
@@ -37,8 +37,8 @@
 import static java.util.Collections.singletonMap;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class CoordinatorStrategyTest {
 
@@ -217,7 +217,7 @@ public void testFatalErrorOldLookupResponses() {
         assertFatalOldLookup(group, Errors.UNKNOWN_SERVER_ERROR);
 
         Throwable throwable = assertFatalOldLookup(group, Errors.GROUP_AUTHORIZATION_FAILED);
-        assertTrue(throwable instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, throwable);
         GroupAuthorizationException exception = (GroupAuthorizationException) throwable;
         assertEquals("foo", exception.groupId());
     }
@@ -233,7 +233,7 @@ public Throwable assertFatalOldLookup(
         assertEquals(singleton(key), result.failedKeys.keySet());
 
         Throwable throwable = result.failedKeys.get(key);
-        assertTrue(error.exception().getClass().isInstance(throwable));
+        assertInstanceOf(error.exception().getClass(), throwable);
         return throwable;
     }
 
@@ -244,7 +244,7 @@ public void testFatalErrorLookupResponses() {
         assertFatalLookup(group, Errors.UNKNOWN_SERVER_ERROR);
 
         Throwable throwable = assertFatalLookup(group, Errors.GROUP_AUTHORIZATION_FAILED);
-        assertTrue(throwable instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, throwable);
         GroupAuthorizationException exception = (GroupAuthorizationException) throwable;
         assertEquals("foo", exception.groupId());
     }
@@ -264,7 +264,7 @@ public Throwable assertFatalLookup(
         assertEquals(singleton(key), result.failedKeys.keySet());
 
         Throwable throwable = result.failedKeys.get(key);
-        assertTrue(error.exception().getClass().isInstance(throwable));
+        assertInstanceOf(error.exception().getClass(), throwable);
         return throwable;
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupOffsetsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupOffsetsHandlerTest.java
index 629ca89801559..fef8a55074131 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupOffsetsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupOffsetsHandlerTest.java
@@ -21,7 +21,7 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 import java.util.Arrays;
 import java.util.Collection;
@@ -189,7 +189,7 @@ private void assertGroupFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(key), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(key)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(key));
     }
 
     private void assertPartitionFailed(
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupsHandlerTest.java
index 3e7cead6f5d40..1d1b152afba05 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DeleteConsumerGroupsHandlerTest.java
@@ -21,7 +21,7 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
@@ -125,6 +125,6 @@ private void assertFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(key), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(key)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(key));
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeConsumerGroupsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeConsumerGroupsHandlerTest.java
index 0ece97640492c..7179e13a4fc78 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeConsumerGroupsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeConsumerGroupsHandlerTest.java
@@ -22,8 +22,8 @@
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -369,7 +369,7 @@ private void assertFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(key), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(key)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(key));
     }
 
     private void assertRequestAndKeys(
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandlerTest.java
index 0f39b4dd01633..6ecf328819240 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandlerTest.java
@@ -52,8 +52,8 @@
 import static java.util.Collections.singletonMap;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class DescribeProducersHandlerTest {
     private DescribeProducersHandler newHandler(
@@ -137,7 +137,7 @@ public void testBuildRequest() {
     public void testAuthorizationFailure() {
         TopicPartition topicPartition = new TopicPartition("foo", 5);
         Throwable exception = assertFatalError(topicPartition, Errors.TOPIC_AUTHORIZATION_FAILED);
-        assertTrue(exception instanceof TopicAuthorizationException);
+        assertInstanceOf(TopicAuthorizationException.class, exception);
         TopicAuthorizationException authException = (TopicAuthorizationException) exception;
         assertEquals(mkSet("foo"), authException.unauthorizedTopics());
     }
@@ -146,7 +146,7 @@ public void testAuthorizationFailure() {
     public void testInvalidTopic() {
         TopicPartition topicPartition = new TopicPartition("foo", 5);
         Throwable exception = assertFatalError(topicPartition, Errors.INVALID_TOPIC_EXCEPTION);
-        assertTrue(exception instanceof InvalidTopicException);
+        assertInstanceOf(InvalidTopicException.class, exception);
         InvalidTopicException invalidTopicException = (InvalidTopicException) exception;
         assertEquals(mkSet("foo"), invalidTopicException.invalidTopics());
     }
@@ -155,7 +155,7 @@ public void testInvalidTopic() {
     public void testUnexpectedError() {
         TopicPartition topicPartition = new TopicPartition("foo", 5);
         Throwable exception = assertFatalError(topicPartition, Errors.UNKNOWN_SERVER_ERROR);
-        assertTrue(exception instanceof UnknownServerException);
+        assertInstanceOf(UnknownServerException.class, exception);
     }
 
     @Test
@@ -185,7 +185,7 @@ public void testFatalNotLeaderErrorIfStaticMapped() {
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(mkSet(topicPartition), result.failedKeys.keySet());
         Throwable exception = result.failedKeys.get(topicPartition);
-        assertTrue(exception instanceof NotLeaderOrFollowerException);
+        assertInstanceOf(NotLeaderOrFollowerException.class, exception);
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeTransactionsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeTransactionsHandlerTest.java
index 7ffda2b00e014..b744795cd98ec 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeTransactionsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/DescribeTransactionsHandlerTest.java
@@ -37,7 +37,7 @@
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class DescribeTransactionsHandlerTest {
     private final LogContext logContext = new LogContext();
@@ -107,7 +107,7 @@ private void assertFatalError(
         assertEquals(mkSet(key), result.failedKeys.keySet());
 
         Throwable throwable = result.failedKeys.get(key);
-        assertTrue(error.exception().getClass().isInstance(throwable));
+        assertInstanceOf(error.exception().getClass(), throwable);
     }
 
     private void assertRetriableError(
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/FenceProducersHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/FenceProducersHandlerTest.java
index c4151ebb0e15e..8b20f84c9889f 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/FenceProducersHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/FenceProducersHandlerTest.java
@@ -34,7 +34,7 @@
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class FenceProducersHandlerTest {
     private final LogContext logContext = new LogContext();
@@ -96,7 +96,7 @@ private void assertFatalError(
         assertEquals(mkSet(key), result.failedKeys.keySet());
 
         Throwable throwable = result.failedKeys.get(key);
-        assertTrue(error.exception().getClass().isInstance(throwable));
+        assertInstanceOf(error.exception().getClass(), throwable);
     }
 
     private void assertRetriableError(
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListConsumerGroupOffsetsHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListConsumerGroupOffsetsHandlerTest.java
index 95fabb3fc2a2f..6484b842dd817 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListConsumerGroupOffsetsHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/ListConsumerGroupOffsetsHandlerTest.java
@@ -21,6 +21,7 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.Arrays;
@@ -403,7 +404,7 @@ private void assertFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(key), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(key)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(key));
     }
 
     private void assertFailedForMultipleGroups(
@@ -415,7 +416,7 @@ private void assertFailedForMultipleGroups(
         for (String g : groupToExceptionMap.keySet()) {
             CoordinatorKey key = CoordinatorKey.byGroupId(g);
             assertTrue(result.failedKeys.containsKey(key));
-            assertTrue(groupToExceptionMap.get(g).isInstance(result.failedKeys.get(key)));
+            assertInstanceOf(groupToExceptionMap.get(g), result.failedKeys.get(key));
         }
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyTest.java
index f65b97b44554d..d17443b8be071 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyTest.java
@@ -44,7 +44,7 @@
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class PartitionLeaderStrategyTest {
 
@@ -78,7 +78,7 @@ public void testBuildLookupRequest() {
     public void testTopicAuthorizationFailure() {
         TopicPartition topicPartition = new TopicPartition("foo", 0);
         Throwable exception = assertFatalTopicError(topicPartition, Errors.TOPIC_AUTHORIZATION_FAILED);
-        assertTrue(exception instanceof TopicAuthorizationException);
+        assertInstanceOf(TopicAuthorizationException.class, exception);
         TopicAuthorizationException authException = (TopicAuthorizationException) exception;
         assertEquals(mkSet("foo"), authException.unauthorizedTopics());
     }
@@ -87,7 +87,7 @@ public void testTopicAuthorizationFailure() {
     public void testInvalidTopicError() {
         TopicPartition topicPartition = new TopicPartition("foo", 0);
         Throwable exception = assertFatalTopicError(topicPartition, Errors.INVALID_TOPIC_EXCEPTION);
-        assertTrue(exception instanceof InvalidTopicException);
+        assertInstanceOf(InvalidTopicException.class, exception);
         InvalidTopicException invalidTopicException = (InvalidTopicException) exception;
         assertEquals(mkSet("foo"), invalidTopicException.invalidTopics());
     }
@@ -96,7 +96,7 @@ public void testInvalidTopicError() {
     public void testUnexpectedTopicError() {
         TopicPartition topicPartition = new TopicPartition("foo", 0);
         Throwable exception = assertFatalTopicError(topicPartition, Errors.UNKNOWN_SERVER_ERROR);
-        assertTrue(exception instanceof UnknownServerException);
+        assertInstanceOf(UnknownServerException.class, exception);
     }
 
     @Test
@@ -121,7 +121,7 @@ public void testRetriablePartitionErrors() {
     public void testUnexpectedPartitionError() {
         TopicPartition topicPartition = new TopicPartition("foo", 0);
         Throwable exception = assertFatalPartitionError(topicPartition, Errors.UNKNOWN_SERVER_ERROR);
-        assertTrue(exception instanceof UnknownServerException);
+        assertInstanceOf(UnknownServerException.class, exception);
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/RemoveMembersFromConsumerGroupHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/RemoveMembersFromConsumerGroupHandlerTest.java
index 3ecd1f10ee2f6..1ad66bf7822e1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/internals/RemoveMembersFromConsumerGroupHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/RemoveMembersFromConsumerGroupHandlerTest.java
@@ -21,7 +21,7 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -164,7 +164,7 @@ private void assertFailed(
         assertEquals(emptySet(), result.completedKeys.keySet());
         assertEquals(emptyList(), result.unmappedKeys);
         assertEquals(singleton(key), result.failedKeys.keySet());
-        assertTrue(expectedExceptionType.isInstance(result.failedKeys.get(key)));
+        assertInstanceOf(expectedExceptionType, result.failedKeys.get(key));
     }
 
     private void assertMemberFailed(
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignorTest.java
index 6cbda7c29989e..3cab19a6a10c0 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerPartitionAssignorTest.java
@@ -35,7 +35,6 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class ConsumerPartitionAssignorTest {
 
@@ -45,7 +44,7 @@ public void shouldInstantiateAssignor() {
                 Collections.singletonList(StickyAssignor.class.getName()),
                 Collections.emptyMap()
         );
-        assertTrue(assignors.get(0) instanceof StickyAssignor);
+        assertInstanceOf(StickyAssignor.class, assignors.get(0));
     }
 
     @Test
@@ -54,8 +53,8 @@ public void shouldInstantiateListOfAssignors() {
                 Arrays.asList(StickyAssignor.class.getName(), CooperativeStickyAssignor.class.getName()),
                 Collections.emptyMap()
         );
-        assertTrue(assignors.get(0) instanceof StickyAssignor);
-        assertTrue(assignors.get(1) instanceof CooperativeStickyAssignor);
+        assertInstanceOf(StickyAssignor.class, assignors.get(0));
+        assertInstanceOf(CooperativeStickyAssignor.class, assignors.get(1));
     }
 
     @Test
@@ -80,7 +79,7 @@ public void shouldInstantiateFromClassType() {
                 initConsumerConfigWithClassTypes(Collections.singletonList(StickyAssignor.class))
                 .getList(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG);
         List<ConsumerPartitionAssignor> assignors = getAssignorInstances(classTypes, Collections.emptyMap());
-        assertTrue(assignors.get(0) instanceof StickyAssignor);
+        assertInstanceOf(StickyAssignor.class, assignors.get(0));
     }
 
     @Test
@@ -91,8 +90,8 @@ public void shouldInstantiateFromListOfClassTypes() {
 
         List<ConsumerPartitionAssignor> assignors = getAssignorInstances(classTypes, Collections.emptyMap());
 
-        assertTrue(assignors.get(0) instanceof StickyAssignor);
-        assertTrue(assignors.get(1) instanceof CooperativeStickyAssignor);
+        assertInstanceOf(StickyAssignor.class, assignors.get(0));
+        assertInstanceOf(CooperativeStickyAssignor.class, assignors.get(1));
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
index 7dec7305d4db4..8cce3fb847eab 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
@@ -145,6 +145,7 @@
 import static org.apache.kafka.common.utils.Utils.propsToMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNotSame;
@@ -259,7 +260,7 @@ public void testExplicitlyOnlyEnableJmxReporter(GroupProtocol groupProtocol) {
         props.setProperty(ConsumerConfig.ENABLE_METRICS_PUSH_CONFIG, "false");
         consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
         assertEquals(1, consumer.metricsRegistry().reporters().size());
-        assertTrue(consumer.metricsRegistry().reporters().get(0) instanceof JmxReporter);
+        assertInstanceOf(JmxReporter.class, consumer.metricsRegistry().reporters().get(0));
     }
 
     @ParameterizedTest
@@ -272,7 +273,7 @@ public void testExplicitlyOnlyEnableClientTelemetryReporter(GroupProtocol groupP
         props.setProperty(ConsumerConfig.AUTO_INCLUDE_JMX_REPORTER_CONFIG, "false");
         consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
         assertEquals(1, consumer.metricsRegistry().reporters().size());
-        assertTrue(consumer.metricsRegistry().reporters().get(0) instanceof ClientTelemetryReporter);
+        assertInstanceOf(ClientTelemetryReporter.class, consumer.metricsRegistry().reporters().get(0));
     }
 
     // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
@@ -2071,7 +2072,7 @@ private void consumerCloseTest(GroupProtocol groupProtocol,
                 TestUtils.waitForCondition(
                     () -> closeException.get() != null, "InterruptException did not occur within timeout.");
 
-                assertTrue(closeException.get() instanceof InterruptException, "Expected exception not thrown " + closeException);
+                assertInstanceOf(InterruptException.class, closeException.get(), "Expected exception not thrown " + closeException);
             } else {
                 future.get(closeTimeoutMs, TimeUnit.MILLISECONDS); // Should succeed without TimeoutException or ExecutionException
                 assertNull(closeException.get(), "Unexpected exception during close");
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java
index bfcf6547677bc..1f13aea21fd42 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractCoordinatorTest.java
@@ -78,6 +78,7 @@
 import static java.util.Collections.emptyMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -341,7 +342,7 @@ public void testGroupMaxSizeExceptionIsFatal() {
 
         RequestFuture<ByteBuffer> future = coordinator.sendJoinGroupRequest();
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
-        assertTrue(future.exception().getClass().isInstance(Errors.GROUP_MAX_SIZE_REACHED.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.GROUP_MAX_SIZE_REACHED.exception());
         assertFalse(future.isRetriable());
     }
 
@@ -359,7 +360,7 @@ public void testJoinGroupRequestTimeout() {
 
         mockTime.sleep(REBALANCE_TIMEOUT_MS - REQUEST_TIMEOUT_MS + AbstractCoordinator.JOIN_GROUP_TIMEOUT_LAPSE);
         assertTrue(consumerClient.poll(future, mockTime.timer(0)));
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
     }
 
     @Test
@@ -377,7 +378,7 @@ public void testJoinGroupRequestTimeoutLowerBoundedByDefaultRequestTimeout() {
 
         mockTime.sleep(expectedRequestDeadline - mockTime.milliseconds() + 1);
         assertTrue(consumerClient.poll(future, mockTime.timer(0)));
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
     }
 
     @Test
@@ -773,7 +774,7 @@ public void testJoinGroupUnknownMemberResponseWithOldGeneration() throws Interru
         mockClient.respond(joinGroupFollowerResponse(currGen.generationId + 1, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID));
 
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
-        assertTrue(future.exception().getClass().isInstance(Errors.UNKNOWN_MEMBER_ID.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.UNKNOWN_MEMBER_ID.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -813,7 +814,7 @@ public void testSyncGroupUnknownMemberResponseWithOldGeneration() throws Interru
 
         mockClient.respond(syncGroupResponse(Errors.UNKNOWN_MEMBER_ID));
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
-        assertTrue(future.exception().getClass().isInstance(Errors.UNKNOWN_MEMBER_ID.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.UNKNOWN_MEMBER_ID.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -853,7 +854,7 @@ public void testSyncGroupIllegalGenerationResponseWithOldGeneration() throws Int
 
         mockClient.respond(syncGroupResponse(Errors.ILLEGAL_GENERATION));
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
-        assertTrue(future.exception().getClass().isInstance(Errors.ILLEGAL_GENERATION.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.ILLEGAL_GENERATION.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -1027,18 +1028,17 @@ public void testHeartbeatRequestWithFencedInstanceIdException() throws Interrupt
         mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
         mockClient.prepareResponse(heartbeatResponse(Errors.FENCED_INSTANCE_ID));
 
-        try {
-            coordinator.ensureActiveGroup();
-            mockTime.sleep(HEARTBEAT_INTERVAL_MS);
-            long startMs = System.currentTimeMillis();
-            while (System.currentTimeMillis() - startMs < 1000) {
-                Thread.sleep(10);
-                coordinator.pollHeartbeat(mockTime.milliseconds());
-            }
-            fail("Expected pollHeartbeat to raise fenced instance id exception in 1 second");
-        } catch (RuntimeException exception) {
-            assertTrue(exception instanceof FencedInstanceIdException);
-        }
+        assertThrows(FencedInstanceIdException.class,
+            () -> {
+                coordinator.ensureActiveGroup();
+                mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+                long startMs = System.currentTimeMillis();
+                while (System.currentTimeMillis() - startMs < 1000) {
+                    Thread.sleep(10);
+                    coordinator.pollHeartbeat(mockTime.milliseconds());
+                }
+            },
+            "Expected pollHeartbeat to raise fenced instance id exception in 1 second");
     }
 
     @Test
@@ -1069,7 +1069,7 @@ public void testJoinGroupRequestWithRebalanceInProgress() {
         RequestFuture<ByteBuffer> future = coordinator.sendJoinGroupRequest();
 
         assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
-        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.REBALANCE_IN_PROGRESS.exception());
         assertEquals(Errors.REBALANCE_IN_PROGRESS.message(), future.exception().getMessage());
         assertTrue(coordinator.rejoinNeededOrPending());
 
@@ -1157,7 +1157,7 @@ public void testHandleMultipleMembersLeaveGroupResponse() {
             leaveGroupResponse(Arrays.asList(memberResponse, memberResponse));
         RequestFuture<Void> leaveGroupFuture = setupLeaveGroup(response);
         assertNotNull(leaveGroupFuture);
-        assertTrue(leaveGroupFuture.exception() instanceof IllegalStateException);
+        assertInstanceOf(IllegalStateException.class, leaveGroupFuture.exception());
     }
 
     @Test
@@ -1178,7 +1178,7 @@ public void testHandleLeaveGroupResponseWithException() {
             leaveGroupResponse(Collections.singletonList(memberResponse));
         RequestFuture<Void> leaveGroupFuture = setupLeaveGroup(response);
         assertNotNull(leaveGroupFuture);
-        assertTrue(leaveGroupFuture.exception() instanceof UnknownMemberIdException);
+        assertInstanceOf(UnknownMemberIdException.class, leaveGroupFuture.exception());
     }
 
     private RequestFuture<Void> setupLeaveGroup(LeaveGroupResponse leaveGroupResponse) {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index 5777aa245ab88..fa413eed36a03 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -729,7 +729,7 @@ public void testCompleteQuietly() {
         assertDoesNotThrow(() -> consumer.completeQuietly(() -> {
             throw new KafkaException("Test exception");
         }, "test", exception));
-        assertTrue(exception.get() instanceof KafkaException);
+        assertInstanceOf(KafkaException.class, exception.get());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
index 7e18924e7b136..d22d618904232 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
@@ -78,6 +78,7 @@
 import static org.apache.kafka.test.TestUtils.assertFutureThrows;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -1315,7 +1316,7 @@ private ClientResponse buildOffsetFetchClientResponse(
             final Errors error,
             final boolean disconnected) {
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof OffsetFetchRequest);
+        assertInstanceOf(OffsetFetchRequest.class, abstractRequest);
         OffsetFetchRequest offsetFetchRequest = (OffsetFetchRequest) abstractRequest;
         OffsetFetchResponse response =
                 new OffsetFetchResponse(error, topicPartitionData);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java
index ba0d3bacef4d9..954ed1c11e09b 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java
@@ -127,6 +127,7 @@
 import static org.apache.kafka.test.TestUtils.toSet;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -642,7 +643,7 @@ public void testManyInFlightAsyncCommitsWithCoordinatorDisconnect() {
             coordinator.commitOffsetsAsync(offsets, (offsets1, exception) -> {
                 responses.incrementAndGet();
                 Throwable cause = exception.getCause();
-                assertTrue(cause instanceof DisconnectException,
+                assertInstanceOf(DisconnectException.class, cause,
                     "Unexpected exception cause type: " + (cause == null ? null : cause.getClass()));
             });
         }
@@ -689,7 +690,7 @@ public void onSuccess(ClientResponse value, RequestFuture<Object> future) {}
 
                     @Override
                     public void onFailure(RuntimeException e, RequestFuture<Object> future) {
-                        assertTrue(e instanceof DisconnectException, "Unexpected exception type: " + e.getClass());
+                        assertInstanceOf(DisconnectException.class, e, "Unexpected exception type: " + e.getClass());
                         assertTrue(coordinator.coordinatorUnknown());
                         asyncCallbackInvoked.set(true);
                     }
@@ -952,7 +953,7 @@ public void testCoordinatorDisconnect() {
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
         assertTrue(coordinator.coordinatorUnknown());
     }
 
@@ -2362,8 +2363,8 @@ private void testInFlightRequestsFailedAfterCoordinatorMarkedDead(Errors error)
         coordinator.invokeCompletedOffsetCommitCallbacks();
 
         assertTrue(coordinator.coordinatorUnknown());
-        assertTrue(firstCommitCallback.exception instanceof RetriableCommitFailedException);
-        assertTrue(secondCommitCallback.exception instanceof RetriableCommitFailedException);
+        assertInstanceOf(RetriableCommitFailedException.class, firstCommitCallback.exception);
+        assertInstanceOf(RetriableCommitFailedException.class, secondCommitCallback.exception);
         assertEquals(coordinator.inFlightAsyncCommits.get(), 0);
     }
 
@@ -2597,7 +2598,7 @@ public void testCommitOffsetAsyncFailedWithDefaultCallback() {
         assertEquals(coordinator.inFlightAsyncCommits.get(), 0);
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertEquals(invokedBeforeTest + 1, mockOffsetCommitCallback.invoked);
-        assertTrue(mockOffsetCommitCallback.exception instanceof RetriableCommitFailedException);
+        assertInstanceOf(RetriableCommitFailedException.class, mockOffsetCommitCallback.exception);
     }
 
     @Test
@@ -2614,7 +2615,7 @@ public void testCommitOffsetAsyncCoordinatorNotAvailable() {
 
         assertTrue(coordinator.coordinatorUnknown());
         assertEquals(1, cb.invoked);
-        assertTrue(cb.exception instanceof RetriableCommitFailedException);
+        assertInstanceOf(RetriableCommitFailedException.class, cb.exception);
     }
 
     @Test
@@ -2631,7 +2632,7 @@ public void testCommitOffsetAsyncNotCoordinator() {
 
         assertTrue(coordinator.coordinatorUnknown());
         assertEquals(1, cb.invoked);
-        assertTrue(cb.exception instanceof RetriableCommitFailedException);
+        assertInstanceOf(RetriableCommitFailedException.class, cb.exception);
     }
 
     @Test
@@ -2648,7 +2649,7 @@ public void testCommitOffsetAsyncDisconnected() {
 
         assertTrue(coordinator.coordinatorUnknown());
         assertEquals(1, cb.invoked);
-        assertTrue(cb.exception instanceof RetriableCommitFailedException);
+        assertInstanceOf(RetriableCommitFailedException.class, cb.exception);
     }
 
     @Test
@@ -2794,7 +2795,7 @@ public void testCommitOffsetIllegalGenerationWithNewGeneration() {
         coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.REBALANCE_IN_PROGRESS.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -2842,7 +2843,7 @@ public void testCommitOffsetIllegalGenerationWithResetGeneration() {
         coordinator.setNewGeneration(AbstractCoordinator.Generation.NO_GENERATION);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+        assertInstanceOf(future.exception().getClass(), new CommitFailedException());
 
         // the generation should not be reset
         assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
@@ -2872,7 +2873,7 @@ public void testCommitOffsetUnknownMemberWithNewGeneration() {
         coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.REBALANCE_IN_PROGRESS.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -2897,7 +2898,7 @@ public void testCommitOffsetUnknownMemberWithResetGeneration() {
         coordinator.setNewGeneration(AbstractCoordinator.Generation.NO_GENERATION);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+        assertInstanceOf(future.exception().getClass(), new CommitFailedException());
 
         // the generation should be reset
         assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
@@ -2947,7 +2948,7 @@ public void testCommitOffsetFencedInstanceWithRebalancingGeneration() {
         coordinator.setNewGeneration(newGen);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+        assertInstanceOf(future.exception().getClass(), Errors.REBALANCE_IN_PROGRESS.exception());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
@@ -2976,7 +2977,7 @@ public void testCommitOffsetFencedInstanceWithNewGeneration() {
         coordinator.setNewGeneration(newGen);
 
         assertTrue(consumerClient.poll(future, time.timer(30000)));
-        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+        assertInstanceOf(future.exception().getClass(), new CommitFailedException());
 
         // the generation should not be reset
         assertEquals(newGen, coordinator.generation());
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClientTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClientTest.java
index 186675491db4f..a07ea42caa556 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerNetworkClientTest.java
@@ -49,6 +49,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -94,7 +95,7 @@ public void sendWithinBackoffPeriodAfterAuthenticationFailure() {
         final RequestFuture<ClientResponse> future = consumerClient.send(node, heartbeat());
         consumerClient.poll(future);
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof AuthenticationException, "Expected only an authentication error.");
+        assertInstanceOf(AuthenticationException.class, future.exception(), "Expected only an authentication error.");
 
         time.sleep(30); // wait less than the backoff period
         assertTrue(client.connectionFailed(node));
@@ -102,7 +103,7 @@ public void sendWithinBackoffPeriodAfterAuthenticationFailure() {
         final RequestFuture<ClientResponse> future2 = consumerClient.send(node, heartbeat());
         consumerClient.poll(future2);
         assertTrue(future2.failed());
-        assertTrue(future2.exception() instanceof AuthenticationException, "Expected only an authentication error.");
+        assertInstanceOf(AuthenticationException.class, future2.exception(), "Expected only an authentication error.");
     }
 
     @Test
@@ -127,7 +128,7 @@ public void testDisconnectWithUnsentRequests() {
         consumerClient.disconnectAsync(node);
         consumerClient.pollNoWakeup();
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
     }
 
     @Test
@@ -139,7 +140,7 @@ public void testDisconnectWithInFlightRequests() {
         consumerClient.disconnectAsync(node);
         consumerClient.pollNoWakeup();
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
     }
 
     @Test
@@ -159,7 +160,7 @@ public void testTimeoutUnsentRequest() {
 
         assertFalse(consumerClient.hasPendingRequests());
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof TimeoutException);
+        assertInstanceOf(TimeoutException.class, future.exception());
     }
 
     @Test
@@ -227,7 +228,7 @@ public void testDisconnectWakesUpPoll() throws Exception {
         consumerClient.disconnectAsync(node);
         t.join();
         assertTrue(future.failed());
-        assertTrue(future.exception() instanceof DisconnectException);
+        assertInstanceOf(DisconnectException.class, future.exception());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
index afee3fff39675..d4496522c07b8 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CoordinatorRequestManagerTest.java
@@ -37,6 +37,7 @@
 import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.argThat;
@@ -202,7 +203,7 @@ private ClientResponse buildResponse(
         Errors error
     ) {
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof FindCoordinatorRequest);
+        assertInstanceOf(FindCoordinatorRequest.class, abstractRequest);
         FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) abstractRequest;
 
         FindCoordinatorResponse findCoordinatorResponse =
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetFetcherTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetFetcherTest.java
index a92c4ede8a823..596b549dd5524 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetFetcherTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetFetcherTest.java
@@ -87,6 +87,7 @@
 import static org.apache.kafka.test.TestUtils.assertOptional;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -462,7 +463,7 @@ private boolean listOffsetMatchesExpectedReset(
         OffsetResetStrategy strategy,
         AbstractRequest request
     ) {
-        assertTrue(request instanceof ListOffsetsRequest);
+        assertInstanceOf(ListOffsetsRequest.class, request);
 
         ListOffsetsRequest req = (ListOffsetsRequest) request;
         assertEquals(singleton(tp.topic()), req.data().topics().stream()
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
index 59ed7d440e264..582efdfe76b4c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java
@@ -835,7 +835,7 @@ private ClientResponse buildOffsetsForLeaderEpochResponse(
             final int endOffset) {
 
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof OffsetsForLeaderEpochRequest);
+        assertInstanceOf(OffsetsForLeaderEpochRequest.class, abstractRequest);
         OffsetsForLeaderEpochRequest offsetsForLeaderEpochRequest = (OffsetsForLeaderEpochRequest) abstractRequest;
         OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
         partitions.forEach(tp -> {
@@ -870,7 +870,7 @@ private ClientResponse buildOffsetsForLeaderEpochResponseWithErrors(
             final Map<TopicPartition, Errors> partitionErrors) {
 
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof OffsetsForLeaderEpochRequest);
+        assertInstanceOf(OffsetsForLeaderEpochRequest.class, abstractRequest);
         OffsetsForLeaderEpochRequest offsetsForLeaderEpochRequest = (OffsetsForLeaderEpochRequest) abstractRequest;
         OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
         partitionErrors.keySet().forEach(tp -> {
@@ -931,7 +931,7 @@ private ClientResponse buildClientResponse(
             final boolean disconnected,
             final AuthenticationException authenticationException) {
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof ListOffsetsRequest);
+        assertInstanceOf(ListOffsetsRequest.class, abstractRequest);
         ListOffsetsRequest offsetFetchRequest = (ListOffsetsRequest) abstractRequest;
         ListOffsetsResponse response = buildListOffsetsResponse(topicResponses);
         return new ClientResponse(
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/TopicMetadataRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/TopicMetadataRequestManagerTest.java
index c7b23150602c4..3f2b2c3d983d2 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/TopicMetadataRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/TopicMetadataRequestManagerTest.java
@@ -56,6 +56,7 @@
 import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.spy;
 
@@ -222,7 +223,7 @@ private ClientResponse buildTopicMetadataClientResponse(
             final String topic,
             final Errors error) {
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof MetadataRequest);
+        assertInstanceOf(MetadataRequest.class, abstractRequest);
         MetadataRequest metadataRequest = (MetadataRequest) abstractRequest;
         Cluster cluster = mockCluster(3, 0);
         List<MetadataResponse.TopicMetadata> topics = new ArrayList<>();
@@ -248,7 +249,7 @@ private ClientResponse buildAllTopicsMetadataClientResponse(
         final NetworkClientDelegate.UnsentRequest request,
         final Errors error) {
         AbstractRequest abstractRequest = request.requestBuilder().build();
-        assertTrue(abstractRequest instanceof MetadataRequest);
+        assertInstanceOf(MetadataRequest.class, abstractRequest);
         MetadataRequest metadataRequest = (MetadataRequest) abstractRequest;
         Cluster cluster = mockCluster(3, 0);
         List<MetadataResponse.TopicMetadata> topics = new ArrayList<>();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
index 3e1518814e7b8..a1f15d3d73e5c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/WakeupTriggerTest.java
@@ -33,7 +33,6 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -174,14 +173,7 @@ public void testDisableWakeupPreservedByClearTask() {
 
     private void assertWakeupExceptionIsThrown(final CompletableFuture<?> future) {
         assertTrue(future.isCompletedExceptionally());
-        try {
-            future.get(DEFAULT_TIMEOUT_MS, TimeUnit.MILLISECONDS);
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof WakeupException);
-            return;
-        } catch (Exception e) {
-            fail("The task should throw an ExecutionException but got:" + e);
-        }
-        fail("The task should throw an ExecutionException");
+        assertInstanceOf(WakeupException.class,
+            assertThrows(ExecutionException.class, () -> future.get(DEFAULT_TIMEOUT_MS, TimeUnit.MILLISECONDS)).getCause());
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java
index 3cb6fc0782d12..781aaaaf31405 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java
@@ -127,6 +127,7 @@
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -491,7 +492,7 @@ public void testExplicitlyOnlyEnableJmxReporter() {
         props.setProperty(ProducerConfig.ENABLE_METRICS_PUSH_CONFIG, "false");
         KafkaProducer<String, String> producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
         assertEquals(1, producer.metrics.reporters().size());
-        assertTrue(producer.metrics.reporters().get(0) instanceof JmxReporter);
+        assertInstanceOf(JmxReporter.class, producer.metrics.reporters().get(0));
         producer.close();
     }
 
@@ -503,7 +504,7 @@ public void testExplicitlyOnlyEnableClientTelemetryReporter() {
         props.setProperty(ProducerConfig.AUTO_INCLUDE_JMX_REPORTER_CONFIG, "false");
         KafkaProducer<String, String> producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
         assertEquals(1, producer.metrics.reporters().size());
-        assertTrue(producer.metrics.reporters().get(0) instanceof ClientTelemetryReporter);
+        assertInstanceOf(ClientTelemetryReporter.class, producer.metrics.reporters().get(0));
         producer.close();
     }
 
@@ -692,7 +693,7 @@ public void shouldCloseProperlyAndThrowIfInterrupted() throws Exception {
             TestUtils.waitForCondition(() -> closeException.get() != null,
                     "InterruptException did not occur within timeout.");
 
-            assertTrue(closeException.get() instanceof InterruptException, "Expected exception not thrown " + closeException);
+            assertInstanceOf(InterruptException.class, closeException.get(), "Expected exception not thrown " + closeException);
         } finally {
             executor.shutdownNow();
         }
@@ -848,9 +849,7 @@ public void testMetadataTimeoutWithMissingTopic(boolean isIdempotenceEnabled) th
         verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
         verify(metadata, times(5)).fetch();
         try {
-            future.get();
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof TimeoutException);
+            assertInstanceOf(TimeoutException.class, assertThrows(ExecutionException.class, future::get).getCause());
         } finally {
             producer.close(Duration.ofMillis(0));
         }
@@ -917,9 +916,7 @@ public void testMetadataTimeoutWithPartitionOutOfRange(boolean isIdempotenceEnab
         verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
         verify(metadata, times(5)).fetch();
         try {
-            future.get();
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof TimeoutException);
+            assertInstanceOf(TimeoutException.class, assertThrows(ExecutionException.class, future::get).getCause());
         } finally {
             producer.close(Duration.ofMillis(0));
         }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
index 9141fd51cb65a..61950cce34a9c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
@@ -43,6 +43,7 @@
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -262,7 +263,7 @@ public void shouldThrowOnSendIfProducerGotFenced() {
         producer.initTransactions();
         producer.fenceProducer();
         Throwable e = assertThrows(KafkaException.class, () -> producer.send(null));
-        assertTrue(e.getCause() instanceof ProducerFencedException, "The root cause of the exception should be ProducerFenced");
+        assertInstanceOf(ProducerFencedException.class, e.getCause(), "The root cause of the exception should be ProducerFenced");
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
index 9af3aff8bbf3f..29f0b7b1a8faa 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
@@ -114,6 +114,7 @@
 import static org.apache.kafka.clients.producer.internals.ProducerTestUtils.runUntil;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -733,7 +734,7 @@ public void testClusterAuthorizationExceptionInInitProducerIdRequest() throws Ex
         prepareAndReceiveInitProducerId(producerId, Errors.CLUSTER_AUTHORIZATION_FAILED);
         assertFalse(transactionManager.hasProducerId());
         assertTrue(transactionManager.hasError());
-        assertTrue(transactionManager.lastError() instanceof ClusterAuthorizationException);
+        assertInstanceOf(ClusterAuthorizationException.class, transactionManager.lastError());
         assertEquals(-1, transactionManager.producerIdAndEpoch().epoch);
 
         assertSendFailure(ClusterAuthorizationException.class);
@@ -768,11 +769,7 @@ public void testCanRetryWithoutIdempotence() throws Exception {
         }, produceResponse(tp0, -1L, Errors.TOPIC_AUTHORIZATION_FAILED, 0));
         sender.runOnce();
         assertTrue(future.isDone());
-        try {
-            future.get();
-        } catch (Exception e) {
-            assertTrue(e.getCause() instanceof TopicAuthorizationException);
-        }
+        assertInstanceOf(TopicAuthorizationException.class, assertThrows(Exception.class, future::get).getCause());
     }
 
     @Test
@@ -2539,12 +2536,10 @@ public void testInflightBatchesExpireOnDeliveryTimeout() throws InterruptedExcep
         time.sleep(deliveryTimeoutMs);
         sender.runOnce();  // receive first response
         assertEquals(0, sender.inFlightBatches(tp0).size(), "Expect zero in-flight batch in accumulator");
-        try {
-            request.get();
-            fail("The expired batch should throw a TimeoutException");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof TimeoutException);
-        }
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, request::get).getCause(),
+            "The expired batch should throw a TimeoutException");
     }
 
     @Test
@@ -2578,10 +2573,10 @@ public void testRecordErrorPropagatedToApplication() throws InterruptedException
             KafkaException exception = TestUtils.assertFutureThrows(future, KafkaException.class);
             Integer index = futureEntry.getKey();
             if (index == 0 || index == 2) {
-                assertTrue(exception instanceof InvalidRecordException);
+                assertInstanceOf(InvalidRecordException.class, exception);
                 assertEquals(index.toString(), exception.getMessage());
             } else if (index == 3) {
-                assertTrue(exception instanceof InvalidRecordException);
+                assertInstanceOf(InvalidRecordException.class, exception);
                 assertEquals(Errors.INVALID_RECORD.message(), exception.getMessage());
             } else {
                 assertEquals(KafkaException.class, exception.getClass());
@@ -2722,10 +2717,10 @@ public void testExpiredBatchesInMultiplePartitions() throws Exception {
         assertEquals(0, sender.inFlightBatches(tp0).size(), "Expect zero in-flight batch in accumulator");
 
         ExecutionException e = assertThrows(ExecutionException.class, request1::get);
-        assertTrue(e.getCause() instanceof TimeoutException);
+        assertInstanceOf(TimeoutException.class, e.getCause());
 
         e = assertThrows(ExecutionException.class, request2::get);
-        assertTrue(e.getCause() instanceof TimeoutException);
+        assertInstanceOf(TimeoutException.class, e.getCause());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index b099ebfc41932..bf9d77a395d2a 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -109,6 +109,7 @@
 import static java.util.Collections.singletonMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -883,7 +884,7 @@ public void testUnsupportedFindCoordinator() {
 
         runUntil(transactionManager::hasFatalError);
         assertTrue(transactionManager.hasFatalError());
-        assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
+        assertInstanceOf(UnsupportedVersionException.class, transactionManager.lastError());
     }
 
     @Test
@@ -902,7 +903,7 @@ public void testUnsupportedInitTransactions() {
 
         runUntil(transactionManager::hasFatalError);
         assertTrue(transactionManager.hasFatalError());
-        assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
+        assertInstanceOf(UnsupportedVersionException.class, transactionManager.lastError());
     }
 
     @Test
@@ -920,10 +921,10 @@ public void testUnsupportedForMessageFormatInTxnOffsetCommit() {
         prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT));
         runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.lastError() instanceof UnsupportedForMessageFormatException);
+        assertInstanceOf(UnsupportedForMessageFormatException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof UnsupportedForMessageFormatException);
+        assertInstanceOf(UnsupportedForMessageFormatException.class, sendOffsetsResult.error());
         assertFatalError(UnsupportedForMessageFormatException.class);
     }
 
@@ -954,10 +955,10 @@ public void testFencedInstanceIdInTxnOffsetCommitByGroupMetadata() {
         }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.FENCED_INSTANCE_ID)));
 
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof FencedInstanceIdException);
+        assertInstanceOf(FencedInstanceIdException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof FencedInstanceIdException);
+        assertInstanceOf(FencedInstanceIdException.class, sendOffsetsResult.error());
         assertAbortableError(FencedInstanceIdException.class);
     }
 
@@ -987,10 +988,10 @@ public void testUnknownMemberIdInTxnOffsetCommitByGroupMetadata() {
         }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.UNKNOWN_MEMBER_ID)));
 
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof CommitFailedException);
+        assertInstanceOf(CommitFailedException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof CommitFailedException);
+        assertInstanceOf(CommitFailedException.class, sendOffsetsResult.error());
         assertAbortableError(CommitFailedException.class);
     }
 
@@ -1022,10 +1023,10 @@ public void testIllegalGenerationInTxnOffsetCommitByGroupMetadata() {
         }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.ILLEGAL_GENERATION)));
 
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof CommitFailedException);
+        assertInstanceOf(CommitFailedException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof CommitFailedException);
+        assertInstanceOf(CommitFailedException.class, sendOffsetsResult.error());
         assertAbortableError(CommitFailedException.class);
     }
 
@@ -1128,7 +1129,7 @@ public void testTransactionalIdAuthorizationFailureInFindCoordinator() {
         runUntil(transactionManager::hasError);
 
         assertTrue(transactionManager.hasFatalError());
-        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, transactionManager.lastError());
         assertFalse(initPidResult.isSuccessful());
         assertThrows(TransactionalIdAuthorizationException.class, initPidResult::await);
         assertFatalError(TransactionalIdAuthorizationException.class);
@@ -1162,11 +1163,11 @@ public void testGroupAuthorizationFailureInFindCoordinator() {
 
         prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED, false, CoordinatorType.GROUP, consumerGroupId);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, transactionManager.lastError());
 
         runUntil(sendOffsetsResult::isCompleted);
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, sendOffsetsResult.error());
 
         GroupAuthorizationException exception = (GroupAuthorizationException) sendOffsetsResult.error();
         assertEquals(consumerGroupId, exception.groupId());
@@ -1191,10 +1192,10 @@ public void testGroupAuthorizationFailureInTxnOffsetCommit() {
         prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp1, Errors.GROUP_AUTHORIZATION_FAILED));
 
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof GroupAuthorizationException);
+        assertInstanceOf(GroupAuthorizationException.class, sendOffsetsResult.error());
         assertFalse(transactionManager.hasPendingOffsetCommits());
 
         GroupAuthorizationException exception = (GroupAuthorizationException) sendOffsetsResult.error();
@@ -1215,10 +1216,10 @@ public void testTransactionalIdAuthorizationFailureInAddOffsetsToTxn() {
 
         prepareAddOffsetsToTxnResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, consumerGroupId, producerId, epoch);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, sendOffsetsResult.error());
 
         assertFatalError(TransactionalIdAuthorizationException.class);
     }
@@ -1235,10 +1236,10 @@ public void testInvalidTxnStateFailureInAddOffsetsToTxn() {
 
         prepareAddOffsetsToTxnResponse(Errors.INVALID_TXN_STATE, consumerGroupId, producerId, epoch);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof InvalidTxnStateException);
+        assertInstanceOf(InvalidTxnStateException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof InvalidTxnStateException);
+        assertInstanceOf(InvalidTxnStateException.class, sendOffsetsResult.error());
 
         assertFatalError(InvalidTxnStateException.class);
     }
@@ -1260,10 +1261,10 @@ public void testTransactionalIdAuthorizationFailureInTxnOffsetCommit() {
         prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED));
         runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, transactionManager.lastError());
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, sendOffsetsResult.error());
 
         assertFatalError(TransactionalIdAuthorizationException.class);
     }
@@ -1289,7 +1290,7 @@ public void testTopicAuthorizationFailureInAddPartitions() throws InterruptedExc
         prepareAddPartitionsToTxn(errors);
         runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.lastError() instanceof TopicAuthorizationException);
+        assertInstanceOf(TopicAuthorizationException.class, transactionManager.lastError());
         assertFalse(transactionManager.isPartitionPendingAdd(tp0));
         assertFalse(transactionManager.isPartitionPendingAdd(tp1));
         assertFalse(transactionManager.isPartitionAdded(tp0));
@@ -1351,7 +1352,7 @@ public void testCommitWithTopicAuthorizationFailureInAddPartitionsInFlight() thr
         assertTrue(commitResult.isCompleted());
         TestUtils.assertFutureThrows(firstPartitionAppend, KafkaException.class);
         TestUtils.assertFutureThrows(secondPartitionAppend, KafkaException.class);
-        assertTrue(commitResult.error() instanceof TopicAuthorizationException);
+        assertInstanceOf(TopicAuthorizationException.class, commitResult.error());
     }
 
     @Test
@@ -1626,7 +1627,7 @@ public void testTransactionalIdAuthorizationFailureInAddPartitions() {
 
         prepareAddPartitionsToTxn(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
+        assertInstanceOf(TransactionalIdAuthorizationException.class, transactionManager.lastError());
 
         assertFatalError(TransactionalIdAuthorizationException.class);
     }
@@ -1642,7 +1643,7 @@ public void testInvalidTxnStateInAddPartitions() {
 
         prepareAddPartitionsToTxn(tp, Errors.INVALID_TXN_STATE);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof InvalidTxnStateException);
+        assertInstanceOf(InvalidTxnStateException.class, transactionManager.lastError());
 
         assertFatalError(InvalidTxnStateException.class);
     }
@@ -1866,7 +1867,7 @@ private void verifyProducerFenced(Future<RecordMetadata> responseFuture) throws
             responseFuture.get();
             fail("Expected to get a ExecutionException from the response");
         } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof ProducerFencedException);
+            assertInstanceOf(ProducerFencedException.class, e.getCause());
         }
 
         // make sure the exception was thrown directly from the follow-up calls.
@@ -1932,13 +1933,13 @@ public void testInvalidProducerEpochFromProduce() throws InterruptedException {
 
         // First we will get an EndTxn for abort.
         assertNotNull(handler);
-        assertTrue(handler.requestBuilder() instanceof EndTxnRequest.Builder);
+        assertInstanceOf(EndTxnRequest.Builder.class, handler.requestBuilder());
 
         handler = transactionManager.nextRequest(false);
 
         // Second we will see an InitPid for handling InvalidProducerEpoch.
         assertNotNull(handler);
-        assertTrue(handler.requestBuilder() instanceof InitProducerIdRequest.Builder);
+        assertInstanceOf(InitProducerIdRequest.Builder.class, handler.requestBuilder());
     }
 
     @Test
@@ -2433,7 +2434,7 @@ public void testSendOffsetWithGroupMetadataFailAsAutoDowngradeTxnCommitNotEnable
 
         assertTrue(addOffsetsResult.isCompleted());
         assertFalse(addOffsetsResult.isSuccessful());
-        assertTrue(addOffsetsResult.error() instanceof UnsupportedVersionException);
+        assertInstanceOf(UnsupportedVersionException.class, addOffsetsResult.error());
         assertFatalError(UnsupportedVersionException.class);
     }
 
@@ -2614,13 +2615,11 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
 
         runUntil(responseFuture::isDone);
 
-        try {
-            // make sure the produce was expired.
-            responseFuture.get();
-            fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof  TimeoutException);
-        }
+        // make sure the produce was expired.
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, responseFuture::get).getCause(),
+            "Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
         assertTrue(transactionManager.hasAbortableError());
     }
 
@@ -2664,21 +2663,17 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
         runUntil(firstBatchResponse::isDone);
         runUntil(secondBatchResponse::isDone);
 
-        try {
-            // make sure the produce was expired.
-            firstBatchResponse.get();
-            fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof  TimeoutException);
-        }
+        // make sure the produce was expired.
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, firstBatchResponse::get).getCause(),
+            "Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
+        // make sure the produce was expired.
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, secondBatchResponse::get).getCause(),
+            "Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
 
-        try {
-            // make sure the produce was expired.
-            secondBatchResponse.get();
-            fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof  TimeoutException);
-        }
         assertTrue(transactionManager.hasAbortableError());
     }
 
@@ -2713,13 +2708,11 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException {
 
         runUntil(responseFuture::isDone);  // We should try to flush the produce, but expire it instead without sending anything.
 
-        try {
-            // make sure the produce was expired.
-            responseFuture.get();
-            fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof  TimeoutException);
-        }
+        // make sure the produce was expired.
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, responseFuture::get).getCause(),
+            "Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
         runUntil(commitResult::isCompleted);  // the commit shouldn't be completed without being sent since the produce request failed.
         assertFalse(commitResult.isSuccessful());  // the commit shouldn't succeed since the produce request failed.
         assertThrows(TimeoutException.class, commitResult::await);
@@ -2784,13 +2777,11 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
 
         runUntil(responseFuture::isDone);  // We should try to flush the produce, but expire it instead without sending anything.
 
-        try {
-            // make sure the produce was expired.
-            responseFuture.get();
-            fail("Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
-        } catch (ExecutionException e) {
-            assertTrue(e.getCause() instanceof  TimeoutException);
-        }
+        // make sure the produce was expired.
+        assertInstanceOf(
+            TimeoutException.class,
+            assertThrows(ExecutionException.class, responseFuture::get).getCause(),
+            "Expected to get a TimeoutException since the queued ProducerBatch should have been expired");
         runUntil(commitResult::isCompleted);
         assertFalse(commitResult.isSuccessful());  // the commit should have been dropped.
 
diff --git a/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java b/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
index f9dc5e75fc57a..81d2530650dfa 100644
--- a/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
@@ -34,6 +34,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.Supplier;
 
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -316,7 +317,7 @@ public void testThenApplyOnSucceededFutureAndFunctionThrowsCompletionException()
         assertIsFailed(dependantFuture);
         awaitAndAssertResult(future, 21, null);
         Throwable cause = awaitAndAssertFailure(dependantFuture, CompletionException.class, "java.lang.RuntimeException: We require more vespene gas");
-        assertTrue(cause.getCause() instanceof RuntimeException);
+        assertInstanceOf(RuntimeException.class, cause.getCause());
         assertEquals(cause.getCause().getMessage(), "We require more vespene gas");
     }
 
@@ -438,7 +439,7 @@ public void testWhenCompleteOnCancelledFuture() {
         assertFalse(dependantFuture.isDone());
         assertTrue(future.cancel(true));
         assertTrue(ran[0]);
-        assertTrue(err[0] instanceof CancellationException);
+        assertInstanceOf(CancellationException.class, err[0]);
     }
 
     private static class CompleterThread<T> extends Thread {
diff --git a/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java b/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java
index fae7ce7daa1d6..ede6918500afe 100644
--- a/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java
@@ -21,7 +21,7 @@
 
 import java.io.IOException;
 
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 /**
@@ -48,7 +48,7 @@ public void testSerializationRoundtrip() throws IOException, ClassNotFoundExcept
 
         //deserialize the byteArray and check if the values are same as original
         Object deserializedObject = Serializer.deserialize(byteArray);
-        assertTrue(deserializedObject instanceof TopicPartition);
+        assertInstanceOf(TopicPartition.class, deserializedObject);
         checkValues((TopicPartition) deserializedObject);
     }
 
@@ -57,7 +57,7 @@ public void testTopiPartitionSerializationCompatibility() throws IOException, Cl
         // assert serialized TopicPartition object in file (serializedData/topicPartitionSerializedfile) is
         // deserializable into TopicPartition and is compatible
         Object deserializedObject = Serializer.deserialize(fileName);
-        assertTrue(deserializedObject instanceof TopicPartition);
+        assertInstanceOf(TopicPartition.class, deserializedObject);
         checkValues((TopicPartition) deserializedObject);
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java b/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
index f1d367bcc40bb..36c9d6f6e2fef 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
@@ -30,6 +30,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -40,7 +41,7 @@ public void testCreateConfigurableKafkaPrincipalBuilder() {
         Map<String, Object> configs = new HashMap<>();
         configs.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, ConfigurableKafkaPrincipalBuilder.class);
         KafkaPrincipalBuilder builder = ChannelBuilders.createPrincipalBuilder(configs, null, null);
-        assertTrue(builder instanceof ConfigurableKafkaPrincipalBuilder);
+        assertInstanceOf(ConfigurableKafkaPrincipalBuilder.class, builder);
         assertTrue(((ConfigurableKafkaPrincipalBuilder) builder).configured);
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java b/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
index 8b00bcdb955b4..21ae83e16acd8 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslTransportLayerTest.java
@@ -73,6 +73,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
@@ -1064,7 +1065,7 @@ false, securityProtocol, config, null, null, time, new LogContext(),
 
         CertStores newServerCertStores = certBuilder(true, "server", args.useInlinePem).addHostName("localhost").build();
         Map<String, Object> newKeystoreConfigs = newServerCertStores.keyStoreProps();
-        assertTrue(serverChannelBuilder instanceof ListenerReconfigurable, "SslChannelBuilder not reconfigurable");
+        assertInstanceOf(ListenerReconfigurable.class, serverChannelBuilder, "SslChannelBuilder not reconfigurable");
         ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
         assertEquals(listenerName, reconfigurableBuilder.listenerName());
         reconfigurableBuilder.validateReconfiguration(newKeystoreConfigs);
@@ -1192,7 +1193,7 @@ false, securityProtocol, config, null, null, time, new LogContext(),
         CertStores newClientCertStores = certBuilder(true, "client", args.useInlinePem).addHostName("localhost").build();
         args.sslClientConfigs = args.getTrustingConfig(newClientCertStores, args.serverCertStores);
         Map<String, Object> newTruststoreConfigs = newClientCertStores.trustStoreProps();
-        assertTrue(serverChannelBuilder instanceof ListenerReconfigurable, "SslChannelBuilder not reconfigurable");
+        assertInstanceOf(ListenerReconfigurable.class, serverChannelBuilder, "SslChannelBuilder not reconfigurable");
         ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
         assertEquals(listenerName, reconfigurableBuilder.listenerName());
         reconfigurableBuilder.validateReconfiguration(newTruststoreConfigs);
diff --git a/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordBatchTest.java b/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordBatchTest.java
index e92ccb05199b2..f9d3ff3d57fd8 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordBatchTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/DefaultRecordBatchTest.java
@@ -44,6 +44,7 @@
 import static org.apache.kafka.common.record.DefaultRecordBatch.RECORDS_OFFSET;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -425,12 +426,12 @@ public void testSkipKeyValueIteratorCorrectness(CompressionType compressionType)
 
             if (CompressionType.NONE == compressionType) {
                 // assert that for uncompressed data stream record iterator is not used
-                assertTrue(skipKeyValueIterator instanceof DefaultRecordBatch.RecordIterator);
+                assertInstanceOf(DefaultRecordBatch.RecordIterator.class, skipKeyValueIterator);
                 // superficial validation for correctness. Deep validation is already performed in other tests
                 assertEquals(Utils.toList(records.records()).size(), Utils.toList(skipKeyValueIterator).size());
             } else {
                 // assert that a streaming iterator is used for compressed records
-                assertTrue(skipKeyValueIterator instanceof DefaultRecordBatch.StreamRecordIterator);
+                assertInstanceOf(DefaultRecordBatch.StreamRecordIterator.class, skipKeyValueIterator);
                 // assert correctness for compressed records
                 assertIterableEquals(Arrays.asList(
                         new PartialDefaultRecord(9, (byte) 0, 0L, 1L, -1, 1, 1),
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
index 254dea0430ece..5caeb7730e7d2 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
@@ -40,6 +40,7 @@
 import java.util.Collections;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -62,7 +63,7 @@ public void testSerdeUnsupportedApiVersionRequest() throws Exception {
         requestBuffer.flip();
 
         RequestAndSize requestAndSize = context.parseRequest(requestBuffer);
-        assertTrue(requestAndSize.request instanceof ApiVersionsRequest);
+        assertInstanceOf(ApiVersionsRequest.class, requestAndSize.request);
         ApiVersionsRequest request = (ApiVersionsRequest) requestAndSize.request;
         assertTrue(request.hasUnsupportedRequestVersion());
 
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorFailureDelayTest.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorFailureDelayTest.java
index db6ba89463057..f80170063ecce 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorFailureDelayTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorFailureDelayTest.java
@@ -45,6 +45,7 @@
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public abstract class SaslAuthenticatorFailureDelayTest {
@@ -227,7 +228,7 @@ private void createAndCheckClientAuthenticationFailure(SecurityProtocol security
                                                            String mechanism, String expectedErrorMessage) throws Exception {
         ChannelState finalState = createAndCheckClientConnectionFailure(securityProtocol, node);
         Exception exception = finalState.exception();
-        assertTrue(exception instanceof SaslAuthenticationException, "Invalid exception class " + exception.getClass());
+        assertInstanceOf(SaslAuthenticationException.class, exception, "Invalid exception class " + exception.getClass());
         if (expectedErrorMessage == null)
             expectedErrorMessage = "Authentication failed during authentication due to invalid credentials with SASL mechanism " + mechanism;
         assertEquals(expectedErrorMessage, exception.getMessage());
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorTest.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorTest.java
index 98b66827e5605..75684e4db33cc 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorTest.java
@@ -140,6 +140,7 @@
 import static org.apache.kafka.common.protocol.ApiKeys.LIST_OFFSETS;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -1204,7 +1205,7 @@ public void testClientLoginCallbackOverride() throws Exception {
         try {
             createClientConnection(securityProtocol, "invalid");
         } catch (Exception e) {
-            assertTrue(e.getCause() instanceof LoginException, "Unexpected exception " + e.getCause());
+            assertInstanceOf(LoginException.class, e.getCause(), "Unexpected exception " + e.getCause());
         }
     }
 
@@ -1805,12 +1806,12 @@ public void testValidSaslOauthBearerMechanismWithoutServerTokens() throws Except
         // Server with extensions, but without a token should fail to start up since it could indicate a configuration error
         saslServerConfigs.put("listener.name.sasl_ssl.oauthbearer." + SaslConfigs.SASL_JAAS_CONFIG,
                 TestJaasConfig.jaasConfigProperty("OAUTHBEARER", Collections.singletonMap("unsecuredLoginExtension_test", "something")));
-        try {
-            createEchoServer(securityProtocol);
-            fail("Server created with invalid login config containing extensions without a token");
-        } catch (Throwable e) {
-            assertTrue(e.getCause() instanceof LoginException, "Unexpected exception " + Utils.stackTrace(e));
-        }
+
+        Throwable throwable = assertThrows(
+            Throwable.class,
+            () -> createEchoServer(securityProtocol),
+            "Server created with invalid login config containing extensions without a token");
+        assertInstanceOf(LoginException.class, throwable.getCause(), "Unexpected exception " + Utils.stackTrace(throwable));
     }
 
     /**
@@ -2221,7 +2222,7 @@ private void createAndCheckClientAuthenticationFailure(SecurityProtocol security
             String mechanism, String expectedErrorMessage) throws Exception {
         ChannelState finalState = createAndCheckClientConnectionFailure(securityProtocol, node);
         Exception exception = finalState.exception();
-        assertTrue(exception instanceof SaslAuthenticationException, "Invalid exception class " + exception.getClass());
+        assertInstanceOf(SaslAuthenticationException.class, exception, "Invalid exception class " + exception.getClass());
         String expectedExceptionMessage = expectedErrorMessage != null ? expectedErrorMessage :
                 "Authentication failed during authentication due to invalid credentials with SASL mechanism " + mechanism;
         assertEquals(expectedExceptionMessage, exception.getMessage());
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginCallbackHandlerTest.java
index e7b839e4cfc3d..c362c211e4bc6 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginCallbackHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginCallbackHandlerTest.java
@@ -21,6 +21,7 @@
 import static org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginCallbackHandler.CLIENT_ID_CONFIG;
 import static org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginCallbackHandler.CLIENT_SECRET_CONFIG;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -213,7 +214,7 @@ public void testConfigureWithAccessTokenFile() throws Exception {
         Map<String, ?> configs = getSaslConfigs(SASL_OAUTHBEARER_TOKEN_ENDPOINT_URL, accessTokenFile.toURI().toString());
         Map<String, Object> jaasConfigs = Collections.emptyMap();
         configureHandler(handler, configs, jaasConfigs);
-        assertTrue(handler.getAccessTokenRetriever() instanceof FileTokenRetriever);
+        assertInstanceOf(FileTokenRetriever.class, handler.getAccessTokenRetriever());
     }
 
     @Test
@@ -224,7 +225,7 @@ public void testConfigureWithAccessClientCredentials() {
         jaasConfigs.put(CLIENT_ID_CONFIG, "an ID");
         jaasConfigs.put(CLIENT_SECRET_CONFIG, "a secret");
         configureHandler(handler, configs, jaasConfigs);
-        assertTrue(handler.getAccessTokenRetriever() instanceof HttpAccessTokenRetriever);
+        assertInstanceOf(HttpAccessTokenRetriever.class, handler.getAccessTokenRetriever());
     }
 
     private void testInvalidAccessToken(String accessToken, String expectedMessageSubstring) throws Exception {
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/secured/OAuthBearerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/secured/OAuthBearerTest.java
index 4cad8675079c6..ba06650f8ee41 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/secured/OAuthBearerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/secured/OAuthBearerTest.java
@@ -17,6 +17,7 @@
 
 package org.apache.kafka.common.security.oauthbearer.internals.secured;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
@@ -39,6 +40,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
 import javax.security.auth.login.AppConfigurationEntry;
+
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
@@ -65,19 +67,7 @@ public abstract class OAuthBearerTest {
     protected void assertThrowsWithMessage(Class<? extends Exception> clazz,
         Executable executable,
         String substring) {
-        boolean failed = false;
-
-        try {
-            executable.execute();
-        } catch (Throwable t) {
-            failed = true;
-            assertTrue(clazz.isInstance(t), String.format("Test failed by exception %s, but expected %s", t.getClass(), clazz));
-
-            assertErrorMessageContains(t.getMessage(), substring);
-        }
-
-        if (!failed)
-            fail("Expected test to fail with " + clazz + " that contains the string " + substring);
+        assertErrorMessageContains(assertThrows(clazz, executable).getMessage(), substring);
     }
 
     protected void assertErrorMessageContains(String actual, String expectedSubstring) {
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java
index 0153abcc8380d..d7d6013a45717 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java
@@ -17,8 +17,8 @@
 package org.apache.kafka.common.security.oauthbearer.internals.unsecured;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
@@ -77,8 +77,8 @@ public void validToken() {
                     + (includeOptionalIssuedAtClaim ? comma(ISSUED_AT_CLAIM_TEXT) : "") + "}";
             Object validationResult = validationResult(UNSECURED_JWT_HEADER_JSON, claimsJson,
                     MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED);
-            assertTrue(validationResult instanceof OAuthBearerValidatorCallback);
-            assertTrue(((OAuthBearerValidatorCallback) validationResult).token() instanceof OAuthBearerUnsecuredJws);
+            assertInstanceOf(OAuthBearerValidatorCallback.class, validationResult);
+            assertInstanceOf(OAuthBearerUnsecuredJws.class, ((OAuthBearerValidatorCallback) validationResult).token());
         }
     }
 
@@ -103,8 +103,8 @@ public void includesRequiredScope() {
         String claimsJson = "{" + SUB_CLAIM_TEXT + comma(EXPIRATION_TIME_CLAIM_TEXT) + comma(SCOPE_CLAIM_TEXT) + "}";
         Object validationResult = validationResult(UNSECURED_JWT_HEADER_JSON, claimsJson,
                 MODULE_OPTIONS_MAP_REQUIRE_EXISTING_SCOPE);
-        assertTrue(validationResult instanceof OAuthBearerValidatorCallback);
-        assertTrue(((OAuthBearerValidatorCallback) validationResult).token() instanceof OAuthBearerUnsecuredJws);
+        assertInstanceOf(OAuthBearerValidatorCallback.class, validationResult);
+        assertInstanceOf(OAuthBearerUnsecuredJws.class, ((OAuthBearerValidatorCallback) validationResult).token());
     }
 
     @Test
@@ -123,7 +123,7 @@ private static void confirmFailsValidation(String headerJson, String claimsJson,
             Map<String, String> moduleOptionsMap, String optionalFailureScope) throws OAuthBearerConfigException,
             OAuthBearerIllegalTokenException {
         Object validationResultObj = validationResult(headerJson, claimsJson, moduleOptionsMap);
-        assertTrue(validationResultObj instanceof OAuthBearerValidatorCallback);
+        assertInstanceOf(OAuthBearerValidatorCallback.class, validationResultObj);
         OAuthBearerValidatorCallback callback = (OAuthBearerValidatorCallback) validationResultObj;
         assertNull(callback.token());
         assertNull(callback.errorOpenIDConfiguration());
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/SslFactoryTest.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/SslFactoryTest.java
index 9f02e847335ff..bcdb5c471ab7c 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/ssl/SslFactoryTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/SslFactoryTest.java
@@ -51,6 +51,7 @@
 import static org.apache.kafka.common.security.ssl.SslFactory.CertificateEntries.ensureCompatible;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNotSame;
@@ -474,7 +475,7 @@ public void testClientSpecifiedSslEngineFactoryUsed() throws Exception {
         clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
         SslFactory sslFactory = new SslFactory(Mode.CLIENT);
         sslFactory.configure(clientSslConfig);
-        assertTrue(sslFactory.sslEngineFactory() instanceof TestSslUtils.TestSslEngineFactory,
+        assertInstanceOf(TestSslUtils.TestSslEngineFactory.class, sslFactory.sslEngineFactory(),
             "SslEngineFactory must be of expected type");
     }
 
@@ -507,7 +508,7 @@ public void testServerSpecifiedSslEngineFactoryUsed() throws Exception {
         serverSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
         SslFactory sslFactory = new SslFactory(Mode.SERVER);
         sslFactory.configure(serverSslConfig);
-        assertTrue(sslFactory.sslEngineFactory() instanceof TestSslUtils.TestSslEngineFactory,
+        assertInstanceOf(TestSslUtils.TestSslEngineFactory.class, sslFactory.sslEngineFactory(),
             "SslEngineFactory must be of expected type");
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/serialization/ListDeserializerTest.java b/clients/src/test/java/org/apache/kafka/common/serialization/ListDeserializerTest.java
index aff01e3fe8a89..fd9ee8395c532 100644
--- a/clients/src/test/java/org/apache/kafka/common/serialization/ListDeserializerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/serialization/ListDeserializerTest.java
@@ -17,9 +17,9 @@
 package org.apache.kafka.common.serialization;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.KafkaException;
@@ -45,7 +45,7 @@ public void testListKeyDeserializerNoArgConstructorsWithClassNames() {
         listDeserializer.configure(props, true);
         final Deserializer<?> inner = listDeserializer.innerDeserializer();
         assertNotNull(inner, "Inner deserializer should be not null");
-        assertTrue(inner instanceof StringDeserializer, "Inner deserializer type should be StringDeserializer");
+        assertInstanceOf(StringDeserializer.class, inner, "Inner deserializer type should be StringDeserializer");
     }
 
     @Test
@@ -55,7 +55,7 @@ public void testListValueDeserializerNoArgConstructorsWithClassNames() {
         listDeserializer.configure(props, false);
         final Deserializer<?> inner = listDeserializer.innerDeserializer();
         assertNotNull(inner, "Inner deserializer should be not null");
-        assertTrue(inner instanceof IntegerDeserializer, "Inner deserializer type should be IntegerDeserializer");
+        assertInstanceOf(IntegerDeserializer.class, inner, "Inner deserializer type should be IntegerDeserializer");
     }
 
     @Test
@@ -65,7 +65,7 @@ public void testListKeyDeserializerNoArgConstructorsWithClassObjects() {
         listDeserializer.configure(props, true);
         final Deserializer<?> inner = listDeserializer.innerDeserializer();
         assertNotNull(inner, "Inner deserializer should be not null");
-        assertTrue(inner instanceof StringDeserializer, "Inner deserializer type should be StringDeserializer");
+        assertInstanceOf(StringDeserializer.class, inner, "Inner deserializer type should be StringDeserializer");
     }
 
     @Test
@@ -75,7 +75,7 @@ public void testListValueDeserializerNoArgConstructorsWithClassObjects() {
         listDeserializer.configure(props, false);
         final Deserializer<?> inner = listDeserializer.innerDeserializer();
         assertNotNull(inner, "Inner deserializer should be not null");
-        assertTrue(inner instanceof StringDeserializer, "Inner deserializer type should be StringDeserializer");
+        assertInstanceOf(StringDeserializer.class, inner, "Inner deserializer type should be StringDeserializer");
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/common/serialization/ListSerializerTest.java b/clients/src/test/java/org/apache/kafka/common/serialization/ListSerializerTest.java
index a8ab191cad203..94473858fef96 100644
--- a/clients/src/test/java/org/apache/kafka/common/serialization/ListSerializerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/serialization/ListSerializerTest.java
@@ -17,9 +17,9 @@
 package org.apache.kafka.common.serialization;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.KafkaException;
@@ -43,7 +43,7 @@ public void testListKeySerializerNoArgConstructorsWithClassName() {
         listSerializer.configure(props, true);
         final Serializer<?> inner = listSerializer.getInnerSerializer();
         assertNotNull(inner, "Inner serializer should be not null");
-        assertTrue(inner instanceof StringSerializer, "Inner serializer type should be StringSerializer");
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
@@ -52,7 +52,7 @@ public void testListValueSerializerNoArgConstructorsWithClassName() {
         listSerializer.configure(props, false);
         final Serializer<?> inner = listSerializer.getInnerSerializer();
         assertNotNull(inner, "Inner serializer should be not null");
-        assertTrue(inner instanceof StringSerializer, "Inner serializer type should be StringSerializer");
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
@@ -61,7 +61,7 @@ public void testListKeySerializerNoArgConstructorsWithClassObject() {
         listSerializer.configure(props, true);
         final Serializer<?> inner = listSerializer.getInnerSerializer();
         assertNotNull(inner, "Inner serializer should be not null");
-        assertTrue(inner instanceof StringSerializer, "Inner serializer type should be StringSerializer");
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
@@ -70,7 +70,7 @@ public void testListValueSerializerNoArgConstructorsWithClassObject() {
         listSerializer.configure(props, false);
         final Serializer<?> inner = listSerializer.getInnerSerializer();
         assertNotNull(inner, "Inner serializer should be not null");
-        assertTrue(inner instanceof StringSerializer, "Inner serializer type should be StringSerializer");
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporterTest.java b/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporterTest.java
index 12672a30f6f11..9652ebfd37cc2 100644
--- a/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporterTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryReporterTest.java
@@ -59,6 +59,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -247,7 +248,7 @@ public void testCreateRequestSubscriptionNeeded() {
         Optional<AbstractRequest.Builder<?>> requestOptional = telemetrySender.createRequest();
         assertNotNull(requestOptional);
         assertTrue(requestOptional.isPresent());
-        assertTrue(requestOptional.get().build() instanceof GetTelemetrySubscriptionsRequest);
+        assertInstanceOf(GetTelemetrySubscriptionsRequest.class, requestOptional.get().build());
         GetTelemetrySubscriptionsRequest request = (GetTelemetrySubscriptionsRequest) requestOptional.get().build();
 
         GetTelemetrySubscriptionsRequest expectedResult = new GetTelemetrySubscriptionsRequest.Builder(
@@ -266,7 +267,7 @@ public void testCreateRequestSubscriptionNeededAfterExistingSubscription() {
         Optional<AbstractRequest.Builder<?>> requestOptional = telemetrySender.createRequest();
         assertNotNull(requestOptional);
         assertTrue(requestOptional.isPresent());
-        assertTrue(requestOptional.get().build() instanceof GetTelemetrySubscriptionsRequest);
+        assertInstanceOf(GetTelemetrySubscriptionsRequest.class, requestOptional.get().build());
         GetTelemetrySubscriptionsRequest request = (GetTelemetrySubscriptionsRequest) requestOptional.get().build();
 
         GetTelemetrySubscriptionsRequest expectedResult = new GetTelemetrySubscriptionsRequest.Builder(
@@ -290,7 +291,7 @@ public void testCreateRequestPushNeeded() {
         Optional<AbstractRequest.Builder<?>> requestOptional = telemetrySender.createRequest();
         assertNotNull(requestOptional);
         assertTrue(requestOptional.isPresent());
-        assertTrue(requestOptional.get().build() instanceof PushTelemetryRequest);
+        assertInstanceOf(PushTelemetryRequest.class, requestOptional.get().build());
         PushTelemetryRequest request = (PushTelemetryRequest) requestOptional.get().build();
 
         PushTelemetryRequest expectedResult = new PushTelemetryRequest.Builder(
@@ -373,7 +374,7 @@ public void testCreateRequestPushCompression(CompressionType compressionType) {
         Optional<AbstractRequest.Builder<?>> requestOptional = telemetrySender.createRequest();
         assertNotNull(requestOptional);
         assertTrue(requestOptional.isPresent());
-        assertTrue(requestOptional.get().build() instanceof PushTelemetryRequest);
+        assertInstanceOf(PushTelemetryRequest.class, requestOptional.get().build());
         PushTelemetryRequest request = (PushTelemetryRequest) requestOptional.get().build();
 
         assertEquals(subscription.clientInstanceId(), request.data().clientInstanceId());
@@ -401,7 +402,7 @@ public void testCreateRequestPushCompressionException() {
             Optional<AbstractRequest.Builder<?>> requestOptional = telemetrySender.createRequest();
             assertNotNull(requestOptional);
             assertTrue(requestOptional.isPresent());
-            assertTrue(requestOptional.get().build() instanceof PushTelemetryRequest);
+            assertInstanceOf(PushTelemetryRequest.class, requestOptional.get().build());
             PushTelemetryRequest request = (PushTelemetryRequest) requestOptional.get().build();
 
             assertEquals(subscription.clientInstanceId(), request.data().clientInstanceId());
diff --git a/clients/src/test/java/org/apache/kafka/test/TestUtils.java b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
index 2052240d38f78..ec4cc412d8621 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -74,6 +74,7 @@
 
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -556,7 +557,7 @@ public static Set<TopicPartition> generateRandomTopicPartitions(int numTopic, in
      */
     public static <T extends Throwable> T assertFutureThrows(Future<?> future, Class<T> exceptionCauseClass) {
         ExecutionException exception = assertThrows(ExecutionException.class, future::get);
-        assertTrue(exceptionCauseClass.isInstance(exception.getCause()),
+        assertInstanceOf(exceptionCauseClass, exception.getCause(),
             "Unexpected exception cause " + exception.getCause());
         return exceptionCauseClass.cast(exception.getCause());
     }
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java b/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
index abb6ea4221460..3aad588fc220b 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
@@ -38,6 +38,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -829,14 +830,14 @@ public void shouldParseBigIntegerAsDecimalWithZeroScale() {
             String.valueOf(value)
         );
         assertEquals(Decimal.schema(0), schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof BigDecimal);
+        assertInstanceOf(BigDecimal.class, schemaAndValue.value());
         assertEquals(value, ((BigDecimal) schemaAndValue.value()).unscaledValue());
         value = BigInteger.valueOf(Long.MIN_VALUE).subtract(new BigInteger("1"));
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Decimal.schema(0), schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof BigDecimal);
+        assertInstanceOf(BigDecimal.class, schemaAndValue.value());
         assertEquals(value, ((BigDecimal) schemaAndValue.value()).unscaledValue());
     }
 
@@ -847,14 +848,14 @@ public void shouldParseByteAsInt8() {
             String.valueOf(value)
         );
         assertEquals(Schema.INT8_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Byte);
+        assertInstanceOf(Byte.class, schemaAndValue.value());
         assertEquals(value.byteValue(), ((Byte) schemaAndValue.value()).byteValue());
         value = Byte.MIN_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.INT8_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Byte);
+        assertInstanceOf(Byte.class, schemaAndValue.value());
         assertEquals(value.byteValue(), ((Byte) schemaAndValue.value()).byteValue());
     }
 
@@ -865,14 +866,14 @@ public void shouldParseShortAsInt16() {
             String.valueOf(value)
         );
         assertEquals(Schema.INT16_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Short);
+        assertInstanceOf(Short.class, schemaAndValue.value());
         assertEquals(value.shortValue(), ((Short) schemaAndValue.value()).shortValue());
         value = Short.MIN_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.INT16_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Short);
+        assertInstanceOf(Short.class, schemaAndValue.value());
         assertEquals(value.shortValue(), ((Short) schemaAndValue.value()).shortValue());
     }
 
@@ -883,14 +884,14 @@ public void shouldParseIntegerAsInt32() {
             String.valueOf(value)
         );
         assertEquals(Schema.INT32_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Integer);
+        assertInstanceOf(Integer.class, schemaAndValue.value());
         assertEquals(value.intValue(), ((Integer) schemaAndValue.value()).intValue());
         value = Integer.MIN_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.INT32_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Integer);
+        assertInstanceOf(Integer.class, schemaAndValue.value());
         assertEquals(value.intValue(), ((Integer) schemaAndValue.value()).intValue());
     }
 
@@ -901,14 +902,14 @@ public void shouldParseLongAsInt64() {
             String.valueOf(value)
         );
         assertEquals(Schema.INT64_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Long);
+        assertInstanceOf(Long.class, schemaAndValue.value());
         assertEquals(value.longValue(), ((Long) schemaAndValue.value()).longValue());
         value = Long.MIN_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.INT64_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Long);
+        assertInstanceOf(Long.class, schemaAndValue.value());
         assertEquals(value.longValue(), ((Long) schemaAndValue.value()).longValue());
     }
 
@@ -919,14 +920,14 @@ public void shouldParseFloatAsFloat32() {
             String.valueOf(value)
         );
         assertEquals(Schema.FLOAT32_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Float);
+        assertInstanceOf(Float.class, schemaAndValue.value());
         assertEquals(value, (Float) schemaAndValue.value(), 0);
         value = -Float.MAX_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.FLOAT32_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Float);
+        assertInstanceOf(Float.class, schemaAndValue.value());
         assertEquals(value, (Float) schemaAndValue.value(), 0);
     }
 
@@ -937,14 +938,14 @@ public void shouldParseDoubleAsFloat64() {
             String.valueOf(value)
         );
         assertEquals(Schema.FLOAT64_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Double);
+        assertInstanceOf(Double.class, schemaAndValue.value());
         assertEquals(value, (Double) schemaAndValue.value(), 0);
         value = -Double.MAX_VALUE;
         schemaAndValue = Values.parseString(
             String.valueOf(value)
         );
         assertEquals(Schema.FLOAT64_SCHEMA, schemaAndValue.schema());
-        assertTrue(schemaAndValue.value() instanceof Double);
+        assertInstanceOf(Double.class, schemaAndValue.value());
         assertEquals(value, (Double) schemaAndValue.value(), 0);
     }
 
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorUtilsTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorUtilsTest.java
index a73997ad57c82..99cdb610759f6 100644
--- a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorUtilsTest.java
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorUtilsTest.java
@@ -32,8 +32,8 @@
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -121,7 +121,7 @@ public void testCreateCompactedTopicFailsWithInvalidConfigurationException() thr
         when(admin.createTopics(any(), any())).thenReturn(ctr);
         Throwable ce = assertThrows(ConnectException.class, () -> MirrorUtils.createCompactedTopic(TOPIC, (short) 1, (short) 1, admin), "Should have exception thrown");
 
-        assertTrue(ce.getCause() instanceof InvalidConfigurationException);
+        assertInstanceOf(InvalidConfigurationException.class, ce.getCause());
         verify(future).get();
         verify(ctr).values();
         verify(admin).createTopics(any(), any());
@@ -135,7 +135,7 @@ public void testCreateCompactedTopicFailsWithTimeoutException() throws Exception
         when(admin.createTopics(any(), any())).thenReturn(ctr);
         Throwable ce = assertThrows(ConnectException.class, () -> MirrorUtils.createCompactedTopic(TOPIC, (short) 1, (short) 1, admin), "Should have exception thrown");
 
-        assertTrue(ce.getCause() instanceof TimeoutException);
+        assertInstanceOf(TimeoutException.class, ce.getCause());
         verify(future).get();
         verify(ctr).values();
         verify(admin).createTopics(any(), any());
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java
index d651ef87090d3..35bfeab7c3dc9 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java
@@ -78,6 +78,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Test simple operations on the workers of a Connect cluster.
@@ -860,7 +861,7 @@ private void assertTimeoutException(Runnable operation, String expectedStageDesc
                         return false;
                     } catch (Throwable t) {
                         latestError.set(t);
-                        assertTrue(t instanceof ConnectRestException);
+                        assertInstanceOf(ConnectRestException.class, t);
                         ConnectRestException restException = (ConnectRestException) t;
 
                         assertEquals(INTERNAL_SERVER_ERROR.getStatusCode(), restException.statusCode());
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExactlyOnceSourceIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExactlyOnceSourceIntegrationTest.java
index 09ea3f5ae4bbf..04dd4c7de67c4 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExactlyOnceSourceIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExactlyOnceSourceIntegrationTest.java
@@ -104,6 +104,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 @Category(IntegrationTest.class)
 public class ExactlyOnceSourceIntegrationTest {
@@ -1078,7 +1079,7 @@ private Map<Integer, List<Long>> parseOffsetForTasks(ConsumerRecords<byte[], byt
     @SuppressWarnings("unchecked")
     private static <T> T assertAndCast(Object o, Class<T> klass, String objectDescription) {
         String className = o == null ? "null" : o.getClass().getName();
-        assertTrue(objectDescription + " should be " + klass.getName() + "; was " + className + " instead", klass.isInstance(o));
+        assertInstanceOf(klass, o, objectDescription + " should be " + klass.getName() + "; was " + className + " instead");
         return (T) o;
     }
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
index 9ced632b90676..f91e37ef9f9f9 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
@@ -51,6 +51,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
@@ -529,7 +530,7 @@ public void testFailConnectorThatIsNeitherSourceNorSink() {
         ArgumentCaptor<Throwable> exceptionCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(listener).onFailure(eq(CONNECTOR), exceptionCapture.capture());
         Throwable e = exceptionCapture.getValue();
-        assertTrue(e instanceof ConnectException);
+        assertInstanceOf(ConnectException.class, e);
         assertTrue(e.getMessage().contains("must be a subclass of"));
     }
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
index c48abb0130a7f..c09fa093fb32e 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
@@ -154,6 +154,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyCollection;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -2123,7 +2124,7 @@ public void testNormalizeSourceConnectorOffsets() throws Exception {
                 Collections.singletonMap("position", 20)
         );
 
-        assertTrue(offsets.values().iterator().next().get("position") instanceof Integer);
+        assertInstanceOf(Integer.class, offsets.values().iterator().next().get("position"));
 
         mockInternalConverters();
 
@@ -2135,7 +2136,7 @@ public void testNormalizeSourceConnectorOffsets() throws Exception {
         assertEquals(1, normalizedOffsets.size());
 
         // The integer value 20 gets deserialized as a long value by the JsonConverter
-        assertTrue(normalizedOffsets.values().iterator().next().get("position") instanceof Long);
+        assertInstanceOf(Long.class, normalizedOffsets.values().iterator().next().get("position"));
     }
 
     @Test
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedHerderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedHerderTest.java
index 6c5ce0fdb53fa..e680f41d2099c 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedHerderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedHerderTest.java
@@ -126,6 +126,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.AdditionalMatchers.leq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -868,7 +869,7 @@ public void testCreateConnectorFailedValidation() throws Exception {
 
         ArgumentCaptor<Throwable> error = ArgumentCaptor.forClass(Throwable.class);
         verify(putConnectorCallback).onCompletion(error.capture(), isNull());
-        assertTrue(error.getValue() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, error.getValue());
         verifyNoMoreInteractions(worker, member, configBackingStore, statusBackingStore, putConnectorCallback);
 
         assertEquals(
@@ -1174,7 +1175,7 @@ public void testRestartConnectorAndTasksUnknownConnector() throws Exception {
         herder.restartConnectorAndTasks(restartRequest, callback);
         herder.tick();
         ExecutionException ee = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(ee.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, ee.getCause());
         assertTrue(ee.getMessage().contains("Unknown connector:"));
 
         verifyNoMoreInteractions(worker, member, configBackingStore, statusBackingStore);
@@ -1200,7 +1201,7 @@ public void testRestartConnectorAndTasksNotLeader() throws Exception {
         herder.restartConnectorAndTasks(restartRequest, callback);
         herder.tick();
         ExecutionException ee = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(ee.getCause() instanceof NotLeaderException);
+        assertInstanceOf(NotLeaderException.class, ee.getCause());
 
         verifyNoMoreInteractions(worker, member, configBackingStore, statusBackingStore);
     }
@@ -1228,7 +1229,7 @@ public void testRestartConnectorAndTasksUnknownStatus() throws Exception {
         herder.restartConnectorAndTasks(restartRequest, callback);
         herder.tick();
         ExecutionException ee = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(ee.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, ee.getCause());
         assertTrue(ee.getMessage().contains("Status for connector"));
 
         verifyNoMoreInteractions(worker, member, configBackingStore, statusBackingStore);
@@ -1837,7 +1838,7 @@ public void testStopConnectorNotLeader() throws Exception {
                 ExecutionException.class,
                 () -> cb.get(0, TimeUnit.SECONDS)
         );
-        assertTrue(e.getCause() instanceof NotLeaderException);
+        assertInstanceOf(NotLeaderException.class, e.getCause());
 
         verifyNoMoreInteractions(worker, member, configBackingStore, statusBackingStore);
     }
@@ -2470,7 +2471,7 @@ public void testPutTaskConfigsMissingRequiredSignature() {
 
         ArgumentCaptor<Throwable> errorCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(taskConfigCb).onCompletion(errorCapture.capture(), isNull());
-        assertTrue(errorCapture.getValue() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, errorCapture.getValue());
 
         verifyNoMoreInteractions(member, taskConfigCb);
     }
@@ -2487,7 +2488,7 @@ public void testPutTaskConfigsDisallowedSignatureAlgorithm() {
 
         ArgumentCaptor<Throwable> errorCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(taskConfigCb).onCompletion(errorCapture.capture(), isNull());
-        assertTrue(errorCapture.getValue() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, errorCapture.getValue());
 
         verifyNoMoreInteractions(member, taskConfigCb);
     }
@@ -2513,7 +2514,7 @@ public void testPutTaskConfigsInvalidSignature() {
 
         ArgumentCaptor<Throwable> errorCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(taskConfigCb).onCompletion(errorCapture.capture(), isNull());
-        assertTrue(errorCapture.getValue() instanceof ConnectRestException);
+        assertInstanceOf(ConnectRestException.class, errorCapture.getValue());
         assertEquals(FORBIDDEN.getStatusCode(), ((ConnectRestException) errorCapture.getValue()).statusCode());
 
         verifyNoMoreInteractions(member, taskConfigCb);
@@ -2531,7 +2532,7 @@ public void putTaskConfigsWorkerStillStarting() {
 
         ArgumentCaptor<Throwable> errorCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(taskConfigCb).onCompletion(errorCapture.capture(), isNull());
-        assertTrue(errorCapture.getValue() instanceof ConnectRestException);
+        assertInstanceOf(ConnectRestException.class, errorCapture.getValue());
         assertEquals(SERVICE_UNAVAILABLE.getStatusCode(), ((ConnectRestException) errorCapture.getValue()).statusCode());
 
         verifyNoMoreInteractions(member, taskConfigCb);
@@ -2667,7 +2668,7 @@ public void testFenceZombiesInvalidSignature() {
 
         ArgumentCaptor<Throwable> errorCapture = ArgumentCaptor.forClass(Throwable.class);
         verify(taskConfigCb).onCompletion(errorCapture.capture(), isNull());
-        assertTrue(errorCapture.getValue() instanceof ConnectRestException);
+        assertInstanceOf(ConnectRestException.class, errorCapture.getValue());
         assertEquals(FORBIDDEN.getStatusCode(), ((ConnectRestException) errorCapture.getValue()).statusCode());
 
         verifyNoMoreInteractions(member);
@@ -2723,7 +2724,7 @@ private void testTaskRequestedZombieFencingForwardingToLeader(boolean succeed) t
         if (!succeed) {
             ExecutionException fencingException =
                     assertThrows(ExecutionException.class, () -> fencing.get(10, TimeUnit.SECONDS));
-            assertTrue(fencingException.getCause() instanceof ConnectException);
+            assertInstanceOf(ConnectException.class, fencingException.getCause());
         } else {
             fencing.get(10, TimeUnit.SECONDS);
         }
@@ -2959,7 +2960,7 @@ public void testExternalZombieFencingRequestAsynchronousFailure() throws Excepti
         herderFencingCallbacks.getAllValues().forEach(cb -> cb.accept(null, fencingException));
 
         ExecutionException exception = assertThrows(ExecutionException.class, () -> fencing.get(10, TimeUnit.SECONDS));
-        assertTrue(exception.getCause() instanceof ConnectException);
+        assertInstanceOf(ConnectException.class, exception.getCause());
 
         stopBackgroundHerder();
 
@@ -3657,7 +3658,7 @@ public void testModifyConnectorOffsetsUnknownConnector() throws Exception {
         herder.modifyConnectorOffsets("connector-does-not-exist", new HashMap<>(), callback);
         herder.tick();
         ExecutionException e = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, e.getCause());
     }
 
     @Test
@@ -3675,7 +3676,7 @@ public void testModifyOffsetsConnectorNotInStoppedState() throws Exception {
         herder.modifyConnectorOffsets(CONN1, null, callback);
         herder.tick();
         ExecutionException e = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, e.getCause());
     }
 
     @Test
@@ -3693,7 +3694,7 @@ public void testModifyOffsetsNotLeader() throws Exception {
         herder.modifyConnectorOffsets(CONN1, new HashMap<>(), callback);
         herder.tick();
         ExecutionException e = assertThrows(ExecutionException.class, () -> callback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof NotLeaderException);
+        assertInstanceOf(NotLeaderException.class, e.getCause());
     }
 
     @Test
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/ErrorReporterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/ErrorReporterTest.java
index cd7bd90c926de..216e8c86c16c6 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/ErrorReporterTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/ErrorReporterTest.java
@@ -59,6 +59,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -220,7 +221,7 @@ public void testLogReportAndReturnFuture() {
             "partition=5, offset=100}.", msg);
 
         Future<RecordMetadata> future = logReporter.report(context);
-        assertTrue(future instanceof CompletableFuture);
+        assertInstanceOf(CompletableFuture.class, future);
     }
 
     @Test
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
index 72fe97c2cb67a..4b60223cbbcc7 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
@@ -71,6 +71,7 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class PluginsTest {
 
@@ -151,7 +152,7 @@ public void shouldInstantiateAndConfigureExplicitlySetHeaderConverterWithCurrent
                                                                      WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
                                                                      ClassLoaderUsage.CURRENT_CLASSLOADER);
         assertNotNull(headerConverter);
-        assertTrue(headerConverter instanceof TestHeaderConverter);
+        assertInstanceOf(TestHeaderConverter.class, headerConverter);
         this.headerConverter = (TestHeaderConverter) headerConverter;
 
         // Validate extra configs got passed through to overridden converters
@@ -162,7 +163,7 @@ public void shouldInstantiateAndConfigureExplicitlySetHeaderConverterWithCurrent
                                                      WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
                                                      ClassLoaderUsage.PLUGINS);
         assertNotNull(headerConverter);
-        assertTrue(headerConverter instanceof TestHeaderConverter);
+        assertInstanceOf(TestHeaderConverter.class, headerConverter);
         this.headerConverter = (TestHeaderConverter) headerConverter;
 
         // Validate extra configs got passed through to overridden converters
@@ -207,7 +208,7 @@ public void shouldInstantiateAndConfigureDefaultHeaderConverter() {
                                                      WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
                                                      ClassLoaderUsage.PLUGINS);
         assertNotNull(headerConverter);
-        assertTrue(headerConverter instanceof SimpleHeaderConverter);
+        assertInstanceOf(SimpleHeaderConverter.class, headerConverter);
     }
 
     @Test
@@ -657,13 +658,6 @@ public static void assertPluginClassLoaderAlwaysActive(Object plugin) {
         }
     }
 
-    public static void assertInstanceOf(Class<?> expected, Object actual, String message) {
-        assertTrue(
-            "Expected an instance of " + expected.getSimpleName() + ", found " + actual + " instead: " + message,
-            expected.isInstance(actual)
-        );
-    }
-
     protected void instantiateAndConfigureConverter(String configPropName, ClassLoaderUsage classLoaderUsage) {
         converter = (TestConverter) plugins.newConverter(config, configPropName, classLoaderUsage);
         assertNotNull(converter);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
index 993f5368ff7d4..41091bbdf467d 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
@@ -90,6 +90,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -297,7 +298,7 @@ public void testDestroyConnector() throws Exception {
                 ExecutionException.class,
                 () -> failedDeleteCallback.get(1000L, TimeUnit.MILLISECONDS)
         );
-        assertTrue(e.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, e.getCause());
     }
 
     @Test
@@ -482,7 +483,7 @@ public void testRestartConnectorAndTasksUnknownConnector() {
         RestartRequest restartRequest = new RestartRequest("UnknownConnector", false, true);
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
         ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(ee.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, ee.getCause());
     }
 
     @Test
@@ -504,7 +505,7 @@ public void testRestartConnectorAndTasksNoStatus() throws Exception {
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
         ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(ee.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, ee.getCause());
         assertTrue(ee.getMessage().contains("Status for connector"));
     }
 
@@ -856,7 +857,7 @@ public void testCorruptConfig() {
         );
         assertNotNull(e.getCause());
         Throwable cause = e.getCause();
-        assertTrue(cause instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, cause);
         assertEquals(
                 cause.getMessage(),
                 "Connector configuration is invalid and contains the following 1 error(s):\n" +
@@ -910,12 +911,12 @@ public void testModifyConnectorOffsetsUnknownConnector() {
                 Collections.singletonMap(Collections.singletonMap("partitionKey", "partitionValue"), Collections.singletonMap("offsetKey", "offsetValue")),
                 alterOffsetsCallback);
         ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, e.getCause());
 
         FutureCallback<Message> resetOffsetsCallback = new FutureCallback<>();
         herder.resetConnectorOffsets("unknown-connector", resetOffsetsCallback);
         e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof NotFoundException);
+        assertInstanceOf(NotFoundException.class, e.getCause());
     }
 
     @Test
@@ -938,12 +939,12 @@ public void testModifyConnectorOffsetsConnectorNotInStoppedState() {
                 Collections.singletonMap(Collections.singletonMap("partitionKey", "partitionValue"), Collections.singletonMap("offsetKey", "offsetValue")),
                 alterOffsetsCallback);
         ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, e.getCause());
 
         FutureCallback<Message> resetOffsetsCallback = new FutureCallback<>();
         herder.resetConnectorOffsets(CONNECTOR_NAME, resetOffsetsCallback);
         e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
-        assertTrue(e.getCause() instanceof BadRequestException);
+        assertInstanceOf(BadRequestException.class, e.getCause());
     }
 
     @Test
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
index 78f995253b337..a0e4d569f403e 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStoreTest.java
@@ -65,6 +65,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.AdditionalMatchers.aryEq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.isNull;
@@ -404,7 +405,7 @@ public void testSetFailure() {
         callback0.getValue().onCompletion(null, null);
         ExecutionException e = assertThrows(ExecutionException.class, () -> setFuture.get(10000, TimeUnit.MILLISECONDS));
         assertNotNull(e.getCause());
-        assertTrue(e.getCause() instanceof KafkaException);
+        assertInstanceOf(KafkaException.class, e.getCause());
 
         store.stop();
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java
index 9f78977967252..07ccbde9de0af 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java
@@ -76,6 +76,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class TopicAdminTest {
 
@@ -110,7 +111,7 @@ public void throwsWithApiVersionMismatchOnDescribe() {
             env.kafkaClient().prepareResponse(describeTopicResponseWithUnsupportedVersion(newTopic));
             TopicAdmin admin = new TopicAdmin(env.adminClient());
             Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
-            assertTrue(e.getCause() instanceof UnsupportedVersionException);
+            assertInstanceOf(UnsupportedVersionException.class, e.getCause());
         }
     }
 
@@ -136,7 +137,7 @@ public void throwsWithClusterAuthorizationFailureOnDescribe() {
             env.kafkaClient().prepareResponse(describeTopicResponseWithClusterAuthorizationException(newTopic));
             TopicAdmin admin = new TopicAdmin(env.adminClient());
             Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
-            assertTrue(e.getCause() instanceof ClusterAuthorizationException);
+            assertInstanceOf(ClusterAuthorizationException.class, e.getCause());
         }
     }
 
@@ -162,7 +163,7 @@ public void throwsWithTopicAuthorizationFailureOnDescribe() {
             env.kafkaClient().prepareResponse(describeTopicResponseWithTopicAuthorizationException(newTopic));
             TopicAdmin admin = new TopicAdmin(env.adminClient());
             Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
-            assertTrue(e.getCause() instanceof TopicAuthorizationException);
+            assertInstanceOf(TopicAuthorizationException.class, e.getCause());
         }
     }
 
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/FlattenTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/FlattenTest.java
index b5729c1929b69..1017518efab63 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/FlattenTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/FlattenTest.java
@@ -34,10 +34,10 @@
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class FlattenTest {
     private final Flatten<SourceRecord> xformKey = new Flatten.Key<>();
@@ -145,7 +145,7 @@ public void testNestedMapWithDelimiter() {
                 null, twoLevelNestedMap));
 
         assertNull(transformed.valueSchema());
-        assertTrue(transformed.value() instanceof Map);
+        assertInstanceOf(Map.class, transformed.value());
         @SuppressWarnings("unchecked")
         Map<String, Object> transformedMap = (Map<String, Object>) transformed.value();
         assertEquals(9, transformedMap.size());
@@ -241,7 +241,7 @@ public void testOptionalFieldMap() {
                 null, oneLevelNestedMap));
 
         assertNull(transformed.valueSchema());
-        assertTrue(transformed.value() instanceof Map);
+        assertInstanceOf(Map.class, transformed.value());
         @SuppressWarnings("unchecked")
         Map<String, Object> transformedMap = (Map<String, Object>) transformed.value();
 
@@ -257,7 +257,7 @@ public void testKey() {
         SourceRecord transformed = xformKey.apply(src);
 
         assertNull(transformed.keySchema());
-        assertTrue(transformed.key() instanceof Map);
+        assertInstanceOf(Map.class, transformed.key());
         @SuppressWarnings("unchecked")
         Map<String, Object> transformedMap = (Map<String, Object>) transformed.key();
         assertEquals(12, transformedMap.get("A.B"));
diff --git a/core/src/test/java/kafka/server/handlers/DescribeTopicPartitionsRequestHandlerTest.java b/core/src/test/java/kafka/server/handlers/DescribeTopicPartitionsRequestHandlerTest.java
index 2596054ffee26..bbab021c95f27 100644
--- a/core/src/test/java/kafka/server/handlers/DescribeTopicPartitionsRequestHandlerTest.java
+++ b/core/src/test/java/kafka/server/handlers/DescribeTopicPartitionsRequestHandlerTest.java
@@ -76,6 +76,7 @@
 import java.util.stream.Collectors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -469,7 +470,7 @@ void testDescribeTopicPartitionsRequestWithEdgeCases() {
         try {
             handler.handleDescribeTopicPartitionsRequest(buildRequest(describeTopicPartitionsRequest, plaintextListener));
         } catch (Exception e) {
-            assertTrue(e instanceof InvalidRequestException, e.getMessage());
+            assertInstanceOf(InvalidRequestException.class, e, e.getMessage());
         }
     }
 
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
index e9885fa8e94b5..7e04818627050 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
@@ -151,6 +151,7 @@
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -323,7 +324,7 @@ void deleteTopic(ControllerRequestContext context, Uuid topicId) throws Exceptio
             assertEquals(1, result.records().size());
 
             ApiMessageAndVersion removeRecordAndVersion = result.records().get(0);
-            assertTrue(removeRecordAndVersion.message() instanceof RemoveTopicRecord);
+            assertInstanceOf(RemoveTopicRecord.class, removeRecordAndVersion.message());
 
             RemoveTopicRecord removeRecord = (RemoveTopicRecord) removeRecordAndVersion.message();
             assertEquals(topicId, removeRecord.topicId());
@@ -2213,7 +2214,7 @@ public void testElectUncleanLeaders_WithoutElr(boolean electAllPartitions) throw
         assertEquals(1, result.records().size());
 
         ApiMessageAndVersion record = result.records().get(0);
-        assertTrue(record.message() instanceof PartitionChangeRecord);
+        assertInstanceOf(PartitionChangeRecord.class, record.message());
 
         PartitionChangeRecord partitionChangeRecord = (PartitionChangeRecord) record.message();
         assertEquals(0, partitionChangeRecord.partitionId());
diff --git a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java
index 44d39880aa0d4..07d8b5a0d6e6f 100644
--- a/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java
+++ b/raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientSnapshotTest.java
@@ -52,6 +52,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -1938,7 +1939,7 @@ private static Optional<FetchSnapshotRequestData.PartitionSnapshot> assertFetchS
         int replicaId,
         int maxBytes
     ) {
-        assertTrue(request.data() instanceof FetchSnapshotRequestData);
+        assertInstanceOf(FetchSnapshotRequestData.class, request.data());
 
         FetchSnapshotRequestData data = (FetchSnapshotRequestData) request.data();
 
diff --git a/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java b/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
index ea4078b7c5663..97eb91dfb4f4f 100644
--- a/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
+++ b/raft/src/test/java/org/apache/kafka/raft/RaftClientTestContext.java
@@ -83,6 +83,7 @@
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -452,8 +453,9 @@ DescribeQuorumResponseData collectDescribeQuorumResponse() {
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.DESCRIBE_QUORUM);
         assertEquals(1, sentMessages.size());
         RaftResponse.Outbound raftMessage = sentMessages.get(0);
-        assertTrue(
-            raftMessage.data() instanceof DescribeQuorumResponseData,
+        assertInstanceOf(
+            DescribeQuorumResponseData.class,
+            raftMessage.data(),
             "Unexpected request type " + raftMessage.data());
         return (DescribeQuorumResponseData) raftMessage.data();
     }
@@ -493,7 +495,7 @@ void assertSentVoteResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.VOTE);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof VoteResponseData);
+        assertInstanceOf(VoteResponseData.class, raftMessage.data());
         VoteResponseData response = (VoteResponseData) raftMessage.data();
 
         assertEquals(error, Errors.forCode(response.errorCode()));
@@ -508,7 +510,7 @@ void assertSentVoteResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.VOTE);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof VoteResponseData);
+        assertInstanceOf(VoteResponseData.class, raftMessage.data());
         VoteResponseData response = (VoteResponseData) raftMessage.data();
         assertTrue(hasValidTopicPartition(response, metadataPartition));
 
@@ -585,7 +587,7 @@ void assertSentBeginQuorumEpochResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.BEGIN_QUORUM_EPOCH);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof BeginQuorumEpochResponseData);
+        assertInstanceOf(BeginQuorumEpochResponseData.class, raftMessage.data());
         BeginQuorumEpochResponseData response = (BeginQuorumEpochResponseData) raftMessage.data();
         assertEquals(responseError, Errors.forCode(response.errorCode()));
     }
@@ -598,7 +600,7 @@ void assertSentBeginQuorumEpochResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.BEGIN_QUORUM_EPOCH);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof BeginQuorumEpochResponseData);
+        assertInstanceOf(BeginQuorumEpochResponseData.class, raftMessage.data());
         BeginQuorumEpochResponseData response = (BeginQuorumEpochResponseData) raftMessage.data();
         assertEquals(Errors.NONE, Errors.forCode(response.errorCode()));
 
@@ -623,7 +625,7 @@ void assertSentEndQuorumEpochResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.END_QUORUM_EPOCH);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof EndQuorumEpochResponseData);
+        assertInstanceOf(EndQuorumEpochResponseData.class, raftMessage.data());
         EndQuorumEpochResponseData response = (EndQuorumEpochResponseData) raftMessage.data();
         assertEquals(responseError, Errors.forCode(response.errorCode()));
     }
@@ -636,7 +638,7 @@ void assertSentEndQuorumEpochResponse(
         List<RaftResponse.Outbound> sentMessages = drainSentResponses(ApiKeys.END_QUORUM_EPOCH);
         assertEquals(1, sentMessages.size());
         RaftMessage raftMessage = sentMessages.get(0);
-        assertTrue(raftMessage.data() instanceof EndQuorumEpochResponseData);
+        assertInstanceOf(EndQuorumEpochResponseData.class, raftMessage.data());
         EndQuorumEpochResponseData response = (EndQuorumEpochResponseData) raftMessage.data();
         assertEquals(Errors.NONE, Errors.forCode(response.errorCode()));
 
@@ -736,7 +738,7 @@ void assertSentFetchSnapshotResponse(Errors responseError) {
         assertEquals(1, sentMessages.size());
 
         RaftMessage message = sentMessages.get(0);
-        assertTrue(message.data() instanceof FetchSnapshotResponseData);
+        assertInstanceOf(FetchSnapshotResponseData.class, message.data());
 
         FetchSnapshotResponseData response = (FetchSnapshotResponseData) message.data();
         assertEquals(responseError, Errors.forCode(response.errorCode()));
@@ -747,7 +749,7 @@ Optional<FetchSnapshotResponseData.PartitionSnapshot> assertSentFetchSnapshotRes
         assertEquals(1, sentMessages.size());
 
         RaftMessage message = sentMessages.get(0);
-        assertTrue(message.data() instanceof FetchSnapshotResponseData);
+        assertInstanceOf(FetchSnapshotResponseData.class, message.data());
 
         FetchSnapshotResponseData response = (FetchSnapshotResponseData) message.data();
         assertEquals(Errors.NONE, Errors.forCode(response.errorCode()));
@@ -802,7 +804,7 @@ void discoverLeaderAsObserver(
     private List<RaftRequest.Outbound> collectBeginEpochRequests(int epoch) {
         List<RaftRequest.Outbound> requests = new ArrayList<>();
         for (RaftRequest.Outbound raftRequest : channel.drainSentRequests(Optional.of(ApiKeys.BEGIN_QUORUM_EPOCH))) {
-            assertTrue(raftRequest.data() instanceof BeginQuorumEpochRequestData);
+            assertInstanceOf(BeginQuorumEpochRequestData.class, raftRequest.data());
             BeginQuorumEpochRequestData request = (BeginQuorumEpochRequestData) raftRequest.data();
 
             BeginQuorumEpochRequestData.PartitionData partitionRequest =
@@ -967,10 +969,10 @@ void assertFetchRequestData(
         long fetchOffset,
         int lastFetchedEpoch
     ) {
-        assertTrue(
-            message.data() instanceof FetchRequestData,
-            "unexpected request type " + message.data()
-        );
+        assertInstanceOf(
+            FetchRequestData.class,
+            message.data(),
+            "unexpected request type " + message.data());
         FetchRequestData request = (FetchRequestData) message.data();
         assertEquals(KafkaRaftClient.MAX_FETCH_SIZE_BYTES, request.maxBytes());
         assertEquals(fetchMaxWaitMs, request.maxWaitMs());
diff --git a/server-common/src/test/java/org/apache/kafka/server/util/InterBrokerSendThreadTest.java b/server-common/src/test/java/org/apache/kafka/server/util/InterBrokerSendThreadTest.java
index f810cc6e63929..90b55ed4a1776 100644
--- a/server-common/src/test/java/org/apache/kafka/server/util/InterBrokerSendThreadTest.java
+++ b/server-common/src/test/java/org/apache/kafka/server/util/InterBrokerSendThreadTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.server.util;
 
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -137,7 +138,7 @@ public void testDisconnectWithoutShutdownShouldCauseException() {
 
         Throwable thrown = throwable.get();
         assertNotNull(thrown);
-        assertTrue(thrown instanceof FatalExitError);
+        assertInstanceOf(FatalExitError.class, thrown);
     }
 
     @Test
@@ -317,9 +318,9 @@ public void testInterruption(boolean isShuttingDown) throws InterruptedException
         final InterBrokerSendThread thread =
             new TestInterBrokerSendThread(networkClient, t -> {
                 if (isShuttingDown)
-                    assertTrue(t instanceof InterruptedException);
+                    assertInstanceOf(InterruptedException.class, t);
                 else
-                    assertTrue(t instanceof FatalExitError);
+                    assertInstanceOf(FatalExitError.class, t);
                 exception.getAndSet(t);
             });
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java b/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
index 3f1bfef63194e..e7353002b22b6 100644
--- a/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/StreamsConfigTest.java
@@ -87,6 +87,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class StreamsConfigTest {
     @Rule
@@ -870,16 +871,16 @@ public void shouldUseNewConfigsWhenPresent() {
         props.put(StreamsConfig.DEFAULT_TIMESTAMP_EXTRACTOR_CLASS_CONFIG, MockTimestampExtractor.class);
 
         final StreamsConfig config = new StreamsConfig(props);
-        assertTrue(config.defaultKeySerde() instanceof Serdes.LongSerde);
-        assertTrue(config.defaultValueSerde() instanceof Serdes.LongSerde);
-        assertTrue(config.defaultTimestampExtractor() instanceof MockTimestampExtractor);
+        assertInstanceOf(Serdes.LongSerde.class, config.defaultKeySerde());
+        assertInstanceOf(Serdes.LongSerde.class, config.defaultValueSerde());
+        assertInstanceOf(MockTimestampExtractor.class, config.defaultTimestampExtractor());
     }
 
     @Test
     public void shouldUseCorrectDefaultsWhenNoneSpecified() {
         final StreamsConfig config = new StreamsConfig(getStreamsConfig());
 
-        assertTrue(config.defaultTimestampExtractor() instanceof FailOnInvalidTimestamp);
+        assertInstanceOf(FailOnInvalidTimestamp.class, config.defaultTimestampExtractor());
         assertThrows(ConfigException.class, config::defaultKeySerde);
         assertThrows(ConfigException.class, config::defaultValueSerde);
     }
@@ -1460,7 +1461,7 @@ public void shouldReturnRackAwareAssignmentNonOverlapCost() {
     @Test
     public void shouldReturnDefaultClientSupplier() {
         final KafkaClientSupplier supplier = streamsConfig.getKafkaClientSupplier();
-        assertTrue(supplier instanceof DefaultKafkaClientSupplier);
+        assertInstanceOf(DefaultKafkaClientSupplier.class, supplier);
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedDeserializerTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedDeserializerTest.java
index cc8f9642455b8..944f72d932e59 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedDeserializerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedDeserializerTest.java
@@ -28,9 +28,8 @@
 import java.util.Map;
 
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertThrows;
-
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class SessionWindowedDeserializerTest {
     private final SessionWindowedDeserializer<?> sessionWindowedDeserializer = new SessionWindowedDeserializer<>(new StringDeserializer());
@@ -41,7 +40,7 @@ public void testSessionWindowedDeserializerConstructor() {
         sessionWindowedDeserializer.configure(props, true);
         final Deserializer<?> inner = sessionWindowedDeserializer.innerDeserializer();
         assertNotNull("Inner deserializer should be not null", inner);
-        assertTrue("Inner deserializer type should be StringDeserializer", inner instanceof StringDeserializer);
+        assertInstanceOf(StringDeserializer.class, inner, "Inner deserializer type should be StringDeserializer");
     }
 
     @Test
@@ -49,7 +48,7 @@ public void shouldSetWindowedInnerClassDeserialiserThroughConfig() {
         props.put(StreamsConfig.WINDOWED_INNER_CLASS_SERDE, Serdes.ByteArraySerde.class.getName());
         final SessionWindowedDeserializer<?> deserializer = new SessionWindowedDeserializer<>();
         deserializer.configure(props, false);
-        assertTrue(deserializer.innerDeserializer() instanceof ByteArrayDeserializer);
+        assertInstanceOf(ByteArrayDeserializer.class, deserializer.innerDeserializer());
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedSerializerTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedSerializerTest.java
index 2a560ed5ac065..78b67643b28e0 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedSerializerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/SessionWindowedSerializerTest.java
@@ -28,8 +28,8 @@
 import java.util.Map;
 
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class SessionWindowedSerializerTest {
     private final SessionWindowedSerializer<?> sessionWindowedSerializer = new SessionWindowedSerializer<>(Serdes.String().serializer());
@@ -40,7 +40,7 @@ public void testSessionWindowedSerializerConstructor() {
         sessionWindowedSerializer.configure(props, true);
         final Serializer<?> inner = sessionWindowedSerializer.innerSerializer();
         assertNotNull("Inner serializer should be not null", inner);
-        assertTrue("Inner serializer type should be StringSerializer", inner instanceof StringSerializer);
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
@@ -48,7 +48,7 @@ public void shouldSetWindowedInnerClassSerialiserThroughConfig() {
         props.put(StreamsConfig.WINDOWED_INNER_CLASS_SERDE, Serdes.ByteArraySerde.class.getName());
         final SessionWindowedSerializer<?> serializer = new SessionWindowedSerializer<>();
         serializer.configure(props, false);
-        assertTrue(serializer.innerSerializer() instanceof ByteArraySerializer);
+        assertInstanceOf(ByteArraySerializer.class, serializer.innerSerializer());
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedDeserializerTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedDeserializerTest.java
index a0357636d92ec..f3b1b8323bd07 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedDeserializerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedDeserializerTest.java
@@ -31,7 +31,7 @@
 import static org.hamcrest.core.Is.is;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class TimeWindowedDeserializerTest {
     private final long windowSize = 5000000;
@@ -43,7 +43,7 @@ public void testTimeWindowedDeserializerConstructor() {
         timeWindowedDeserializer.configure(props, true);
         final Deserializer<?> inner = timeWindowedDeserializer.innerDeserializer();
         assertNotNull("Inner deserializer should be not null", inner);
-        assertTrue("Inner deserializer type should be StringDeserializer", inner instanceof StringDeserializer);
+        assertInstanceOf(StringDeserializer.class, inner, "Inner deserializer type should be StringDeserializer");
         assertThat(timeWindowedDeserializer.getWindowSize(), is(5000000L));
     }
 
@@ -54,7 +54,7 @@ public void shouldSetWindowSizeAndWindowedInnerDeserialiserThroughConfigs() {
         final TimeWindowedDeserializer<?> deserializer = new TimeWindowedDeserializer<>();
         deserializer.configure(props, false);
         assertThat(deserializer.getWindowSize(), is(500L));
-        assertTrue(deserializer.innerDeserializer() instanceof ByteArrayDeserializer);
+        assertInstanceOf(ByteArrayDeserializer.class, deserializer.innerDeserializer());
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedSerializerTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedSerializerTest.java
index b5e975443091b..5b221f052b118 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedSerializerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/TimeWindowedSerializerTest.java
@@ -29,7 +29,7 @@
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class TimeWindowedSerializerTest {
     private final TimeWindowedSerializer<?> timeWindowedSerializer = new TimeWindowedSerializer<>(Serdes.String().serializer());
@@ -40,7 +40,7 @@ public void testTimeWindowedSerializerConstructor() {
         timeWindowedSerializer.configure(props, true);
         final Serializer<?> inner = timeWindowedSerializer.innerSerializer();
         assertNotNull("Inner serializer should be not null", inner);
-        assertTrue("Inner serializer type should be StringSerializer", inner instanceof StringSerializer);
+        assertInstanceOf(StringSerializer.class, inner, "Inner serializer type should be StringSerializer");
     }
 
     @Test
@@ -48,7 +48,7 @@ public void shouldSetWindowedInnerClassSerialiserThroughConfig() {
         props.put(StreamsConfig.WINDOWED_INNER_CLASS_SERDE, Serdes.ByteArraySerde.class.getName());
         final TimeWindowedSerializer<?> serializer = new TimeWindowedSerializer<>();
         serializer.configure(props, false);
-        assertTrue(serializer.innerSerializer() instanceof ByteArraySerializer);
+        assertInstanceOf(ByteArraySerializer.class, serializer.innerSerializer());
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/WindowedSerdesTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/WindowedSerdesTest.java
index f12d0db660a1f..4c7e97be40752 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/WindowedSerdesTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/WindowedSerdesTest.java
@@ -27,7 +27,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class WindowedSerdesTest {
 
@@ -36,19 +36,19 @@ public class WindowedSerdesTest {
     @Test
     public void shouldWrapForTimeWindowedSerde() {
         final Serde<Windowed<String>> serde = WindowedSerdes.timeWindowedSerdeFrom(String.class, Long.MAX_VALUE);
-        assertTrue(serde.serializer() instanceof TimeWindowedSerializer);
-        assertTrue(serde.deserializer() instanceof TimeWindowedDeserializer);
-        assertTrue(((TimeWindowedSerializer) serde.serializer()).innerSerializer() instanceof StringSerializer);
-        assertTrue(((TimeWindowedDeserializer) serde.deserializer()).innerDeserializer() instanceof StringDeserializer);
+        assertInstanceOf(TimeWindowedSerializer.class, serde.serializer());
+        assertInstanceOf(TimeWindowedDeserializer.class, serde.deserializer());
+        assertInstanceOf(StringSerializer.class, ((TimeWindowedSerializer) serde.serializer()).innerSerializer());
+        assertInstanceOf(StringDeserializer.class, ((TimeWindowedDeserializer) serde.deserializer()).innerDeserializer());
     }
 
     @Test
     public void shouldWrapForSessionWindowedSerde() {
         final Serde<Windowed<String>> serde = WindowedSerdes.sessionWindowedSerdeFrom(String.class);
-        assertTrue(serde.serializer() instanceof SessionWindowedSerializer);
-        assertTrue(serde.deserializer() instanceof SessionWindowedDeserializer);
-        assertTrue(((SessionWindowedSerializer) serde.serializer()).innerSerializer() instanceof StringSerializer);
-        assertTrue(((SessionWindowedDeserializer) serde.deserializer()).innerDeserializer() instanceof StringDeserializer);
+        assertInstanceOf(SessionWindowedSerializer.class, serde.serializer());
+        assertInstanceOf(SessionWindowedDeserializer.class, serde.deserializer());
+        assertInstanceOf(StringSerializer.class, ((SessionWindowedSerializer) serde.serializer()).innerSerializer());
+        assertInstanceOf(StringDeserializer.class, ((SessionWindowedDeserializer) serde.deserializer()).innerDeserializer());
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilderTest.java
index 1106e914cf362..3517dfc651d53 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilderTest.java
@@ -74,6 +74,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class InternalStreamsBuilderTest {
 
@@ -1186,40 +1187,40 @@ public void shouldSetUseVersionedSemanticsOnTableForeignSelfJoin() {
 
     private void verifyVersionedSemantics(final TableFilterNode<?, ?> filterNode, final boolean expectedValue) {
         final ProcessorSupplier<?, ?, ?, ?> processorSupplier = filterNode.processorParameters().processorSupplier();
-        assertTrue(processorSupplier instanceof KTableFilter);
+        assertInstanceOf(KTableFilter.class, processorSupplier);
         final KTableFilter<?, ?> tableFilter = (KTableFilter<?, ?>) processorSupplier;
         assertEquals(expectedValue, tableFilter.isUseVersionedSemantics());
     }
 
     private void verifyVersionedSemantics(final TableRepartitionMapNode<?, ?> repartitionMapNode, final boolean expectedValue) {
         final ProcessorSupplier<?, ?, ?, ?> processorSupplier = repartitionMapNode.processorParameters().processorSupplier();
-        assertTrue(processorSupplier instanceof KTableRepartitionMap);
+        assertInstanceOf(KTableRepartitionMap.class, processorSupplier);
         final KTableRepartitionMap<?, ?, ?, ?> repartitionMap = (KTableRepartitionMap<?, ?, ?, ?>) processorSupplier;
         assertEquals(expectedValue, repartitionMap.isUseVersionedSemantics());
     }
 
     private void verifyVersionedSemantics(final KTableKTableJoinNode<?, ?, ?, ?> joinNode, final boolean expectedValueLeft, final boolean expectedValueRight) {
         final ProcessorSupplier<?, ?, ?, ?> thisProcessorSupplier = joinNode.thisProcessorParameters().processorSupplier();
-        assertTrue(thisProcessorSupplier instanceof KTableKTableAbstractJoin);
+        assertInstanceOf(KTableKTableAbstractJoin.class, thisProcessorSupplier);
         final KTableKTableAbstractJoin<?, ?, ?, ?> thisJoin = (KTableKTableAbstractJoin<?, ?, ?, ?>) thisProcessorSupplier;
         assertEquals(expectedValueLeft, thisJoin.isUseVersionedSemantics());
 
         final ProcessorSupplier<?, ?, ?, ?> otherProcessorSupplier = joinNode.otherProcessorParameters().processorSupplier();
-        assertTrue(otherProcessorSupplier instanceof KTableKTableAbstractJoin);
+        assertInstanceOf(KTableKTableAbstractJoin.class, otherProcessorSupplier);
         final KTableKTableAbstractJoin<?, ?, ?, ?> otherJoin = (KTableKTableAbstractJoin<?, ?, ?, ?>) otherProcessorSupplier;
         assertEquals(expectedValueRight, otherJoin.isUseVersionedSemantics());
     }
 
     private void verifyVersionedSemantics(final ForeignJoinSubscriptionSendNode<?, ?> joinThisNode, final boolean expectedValue) {
         final ProcessorSupplier<?, ?, ?, ?> thisProcessorSupplier = joinThisNode.processorParameters().processorSupplier();
-        assertTrue(thisProcessorSupplier instanceof SubscriptionSendProcessorSupplier);
+        assertInstanceOf(SubscriptionSendProcessorSupplier.class, thisProcessorSupplier);
         final SubscriptionSendProcessorSupplier<?, ?, ?> joinThis = (SubscriptionSendProcessorSupplier<?, ?, ?>) thisProcessorSupplier;
         assertEquals(expectedValue, joinThis.isUseVersionedSemantics());
     }
 
     private void verifyVersionedSemantics(final ForeignTableJoinNode<?, ?> joinOtherNode, final boolean expectedValue) {
         final ProcessorSupplier<?, ?, ?, ?> otherProcessorSupplier = joinOtherNode.processorParameters().processorSupplier();
-        assertTrue(otherProcessorSupplier instanceof ForeignTableJoinProcessorSupplier);
+        assertInstanceOf(ForeignTableJoinProcessorSupplier.class, otherProcessorSupplier);
         final ForeignTableJoinProcessorSupplier<?, ?, ?> joinThis = (ForeignTableJoinProcessorSupplier<?, ?, ?>) otherProcessorSupplier;
         assertEquals(expectedValue, joinThis.isUseVersionedSemantics());
     }
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformTest.java
index 2a3042e2fc4c9..9ea08636e50b4 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformTest.java
@@ -34,7 +34,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -129,6 +129,6 @@ public void shouldGetFlatTransformProcessor() {
 
         final Processor<Number, Number, Integer, Integer> processor = processorSupplier.get();
 
-        assertTrue(processor instanceof KStreamFlatTransformProcessor);
+        assertInstanceOf(KStreamFlatTransformProcessor.class, processor);
     }
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformValuesTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformValuesTest.java
index 988eab92bf334..4b9e2f82d95ef 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformValuesTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamFlatTransformValuesTest.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.streams.kstream.internals;
 
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -129,6 +129,6 @@ public void shouldGetFlatTransformValuesProcessor() {
 
         final Processor<Integer, Integer, Integer, String> processor = processorSupplier.get();
 
-        assertTrue(processor instanceof KStreamFlatTransformValuesProcessor);
+        assertInstanceOf(KStreamFlatTransformValuesProcessor.class, processor);
     }
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdaterTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdaterTest.java
index 4d4728edab0ea..4f6a659aca702 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdaterTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdaterTest.java
@@ -61,6 +61,7 @@
 import static org.hamcrest.Matchers.not;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.anyBoolean;
@@ -1676,7 +1677,7 @@ private void verifyGetTasks(final Set<StreamTask> expectedActiveTasks,
         final Set<Task> tasks = stateUpdater.getTasks();
 
         assertEquals(expectedActiveTasks.size() + expectedStandbyTasks.size(), tasks.size());
-        tasks.forEach(task -> assertTrue(task instanceof ReadOnlyTask));
+        tasks.forEach(task -> assertInstanceOf(ReadOnlyTask.class, task));
         final Set<TaskId> actualTaskIds = tasks.stream().map(Task::id).collect(Collectors.toSet());
         final Set<Task> expectedTasks = new HashSet<>(expectedActiveTasks);
         expectedTasks.addAll(expectedStandbyTasks);
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
index a63e7e884a774..63f00b8cb7f5f 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
@@ -78,6 +78,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 public class InternalTopologyBuilderTest {
 
@@ -891,14 +892,14 @@ public void shouldAddInternalTopicConfigForWindowStores() {
         assertEquals(TopicConfig.CLEANUP_POLICY_COMPACT + "," + TopicConfig.CLEANUP_POLICY_DELETE, properties1.get(TopicConfig.CLEANUP_POLICY_CONFIG));
         assertEquals("40000", properties1.get(TopicConfig.RETENTION_MS_CONFIG));
         assertEquals("appId-store1-changelog", topicConfig1.name());
-        assertTrue(topicConfig1 instanceof WindowedChangelogTopicConfig);
+        assertInstanceOf(WindowedChangelogTopicConfig.class, topicConfig1);
         final InternalTopicConfig topicConfig2 = topicsInfo.stateChangelogTopics.get("appId-store2-changelog");
         final Map<String, String> properties2 = topicConfig2.getProperties(Collections.emptyMap(), 10000);
         assertEquals(3, properties2.size());
         assertEquals(TopicConfig.CLEANUP_POLICY_COMPACT + "," + TopicConfig.CLEANUP_POLICY_DELETE, properties2.get(TopicConfig.CLEANUP_POLICY_CONFIG));
         assertEquals("40000", properties2.get(TopicConfig.RETENTION_MS_CONFIG));
         assertEquals("appId-store2-changelog", topicConfig2.name());
-        assertTrue(topicConfig2 instanceof WindowedChangelogTopicConfig);
+        assertInstanceOf(WindowedChangelogTopicConfig.class, topicConfig2);
     }
 
     @Test
@@ -923,7 +924,7 @@ public void shouldAddInternalTopicConfigForVersionedStores() {
         assertEquals(TopicConfig.CLEANUP_POLICY_COMPACT, properties.get(TopicConfig.CLEANUP_POLICY_CONFIG));
         assertEquals(Long.toString(60_000L + 24 * 60 * 60 * 1000L), properties.get(TopicConfig.MIN_COMPACTION_LAG_MS_CONFIG));
         assertEquals("appId-vstore-changelog", topicConfig.name());
-        assertTrue(topicConfig instanceof VersionedChangelogTopicConfig);
+        assertInstanceOf(VersionedChangelogTopicConfig.class, topicConfig);
     }
 
     @Test
@@ -940,7 +941,7 @@ public void shouldAddInternalTopicConfigForNonWindowNonVersionedStores() {
         assertEquals(2, properties.size());
         assertEquals(TopicConfig.CLEANUP_POLICY_COMPACT, properties.get(TopicConfig.CLEANUP_POLICY_CONFIG));
         assertEquals("appId-testStore-changelog", topicConfig.name());
-        assertTrue(topicConfig instanceof UnwindowedUnversionedChangelogTopicConfig);
+        assertInstanceOf(UnwindowedUnversionedChangelogTopicConfig.class, topicConfig);
     }
 
     @Test
@@ -956,7 +957,7 @@ public void shouldAddInternalTopicConfigForRepartitionTopics() {
         assertEquals(String.valueOf(-1), properties.get(TopicConfig.RETENTION_MS_CONFIG));
         assertEquals(TopicConfig.CLEANUP_POLICY_DELETE, properties.get(TopicConfig.CLEANUP_POLICY_CONFIG));
         assertEquals("appId-foo", topicConfig.name());
-        assertTrue(topicConfig instanceof RepartitionTopicConfig);
+        assertInstanceOf(RepartitionTopicConfig.class, topicConfig);
     }
 
     @Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorStateManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorStateManagerTest.java
index fef0fc055bb0d..8de868f0e828a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorStateManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorStateManagerTest.java
@@ -81,6 +81,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -996,7 +997,7 @@ public void shouldThrowIllegalStateIfInitializingOffsetsForCorruptedTasks() {
             stateMgr.markChangelogAsCorrupted(mkSet(persistentStorePartition));
 
             final ProcessorStateException thrown = assertThrows(ProcessorStateException.class, () -> stateMgr.initializeStoreOffsetsFromCheckpoint(true));
-            assertTrue(thrown.getCause() instanceof IllegalStateException);
+            assertInstanceOf(IllegalStateException.class, thrown.getCause());
         } finally {
             stateMgr.close();
         }
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
index f0feebb73f739..44644ba1577b6 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
@@ -29,7 +29,6 @@
 import org.apache.kafka.clients.producer.MockProducer;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
@@ -161,7 +160,6 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -1634,18 +1632,16 @@ public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerWasFencedWhilePr
         mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1L);
         consumer.addRecord(new ConsumerRecord<>(topic1, 1, 1, new byte[0], new byte[0]));
 
-        try {
-            if (processingThreadsEnabled) {
-                runUntilTimeoutOrException(this::runOnce);
-            } else {
-                runOnce();
-            }
-            fail("Should have thrown TaskMigratedException");
-        } catch (final KafkaException expected) {
-            assertTrue(String.format("Expected TaskMigratedException but got %s", expected), expected instanceof TaskMigratedException);
-            assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
-                thread.readOnlyActiveTasks().stream().anyMatch(task -> task.id().equals(task1)));
-        }
+        assertThrows(TaskMigratedException.class,
+            () -> {
+                if (processingThreadsEnabled) {
+                    runUntilTimeoutOrException(this::runOnce);
+                } else {
+                    runOnce();
+                }
+            });
+        assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
+            thread.readOnlyActiveTasks().stream().anyMatch(task -> task.id().equals(task1)));
 
         assertThat(producer.commitCount(), equalTo(1L));
     }
@@ -1847,18 +1843,16 @@ public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerGotFencedInCommi
         producer.commitTransactionException = new ProducerFencedException("Producer is fenced");
         mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1L);
         consumer.addRecord(new ConsumerRecord<>(topic1, 1, 1, new byte[0], new byte[0]));
-        try {
-            if (processingThreadsEnabled) {
-                runUntilTimeoutOrException(this::runOnce);
-            } else {
-                runOnce();
-            }
-            fail("Should have thrown TaskMigratedException");
-        } catch (final KafkaException expected) {
-            assertTrue(expected instanceof TaskMigratedException);
-            assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
-                thread.readOnlyActiveTasks().stream().anyMatch(task -> task.id().equals(task1)));
-        }
+        assertThrows(TaskMigratedException.class,
+            () -> {
+                if (processingThreadsEnabled) {
+                    runUntilTimeoutOrException(this::runOnce);
+                } else {
+                    runOnce();
+                }
+            });
+        assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
+            thread.readOnlyActiveTasks().stream().anyMatch(task -> task.id().equals(task1)));
 
         assertThat(producer.commitCount(), equalTo(0L));
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StandbyTaskAssignorFactoryTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StandbyTaskAssignorFactoryTest.java
index 4c4b35f83a2e6..ef3f0fbf4c241 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StandbyTaskAssignorFactoryTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StandbyTaskAssignorFactoryTest.java
@@ -31,7 +31,7 @@
 import org.mockito.quality.Strictness;
 
 import static java.util.Collections.singletonList;
-import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -80,7 +80,7 @@ public void setUp() {
     @Test
     public void shouldReturnClientTagAwareStandbyTaskAssignorWhenRackAwareAssignmentTagsIsSet() {
         final StandbyTaskAssignor standbyTaskAssignor = StandbyTaskAssignorFactory.create(newAssignmentConfigs(singletonList("az")), rackAwareTaskAssignor);
-        assertTrue(standbyTaskAssignor instanceof ClientTagAwareStandbyTaskAssignor);
+        assertInstanceOf(ClientTagAwareStandbyTaskAssignor.class, standbyTaskAssignor);
         if (state != State.NULL) {
             verify(rackAwareTaskAssignor, never()).racksForProcess();
             verify(rackAwareTaskAssignor, never()).validClientRack();
@@ -91,15 +91,15 @@ public void shouldReturnClientTagAwareStandbyTaskAssignorWhenRackAwareAssignment
     public void shouldReturnDefaultOrRackAwareStandbyTaskAssignorWhenRackAwareAssignmentTagsIsEmpty() {
         final StandbyTaskAssignor standbyTaskAssignor = StandbyTaskAssignorFactory.create(newAssignmentConfigs(Collections.emptyList()), rackAwareTaskAssignor);
         if (state == State.ENABLED) {
-            assertTrue(standbyTaskAssignor instanceof ClientTagAwareStandbyTaskAssignor);
+            assertInstanceOf(ClientTagAwareStandbyTaskAssignor.class, standbyTaskAssignor);
             verify(rackAwareTaskAssignor, times(1)).racksForProcess();
             verify(rackAwareTaskAssignor, times(1)).validClientRack();
         } else if (state == State.DISABLED) {
-            assertTrue(standbyTaskAssignor instanceof DefaultStandbyTaskAssignor);
+            assertInstanceOf(DefaultStandbyTaskAssignor.class, standbyTaskAssignor);
             verify(rackAwareTaskAssignor, never()).racksForProcess();
             verify(rackAwareTaskAssignor, times(1)).validClientRack();
         } else {
-            assertTrue(standbyTaskAssignor instanceof DefaultStandbyTaskAssignor);
+            assertInstanceOf(DefaultStandbyTaskAssignor.class, standbyTaskAssignor);
         }
     }
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
index fef75bfd1043b..f293c465c4008 100644
--- a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
@@ -22,6 +22,7 @@
 
 import java.time.Duration;
 
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -84,6 +85,6 @@ public void testInvalidBroker() {
             "--election-type", "unclean",
             "--all-topic-partitions"
         ));
-        assertTrue(e.getCause() instanceof TimeoutException);
+        assertInstanceOf(TimeoutException.class, e.getCause());
     }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
index f68b55ed3fec1..cc99fd02118f4 100644
--- a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
@@ -53,6 +53,7 @@
 import java.util.concurrent.ExecutionException;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -250,7 +251,7 @@ public void testTopicDoesNotExist() {
             "--topic", "unknown-topic-name",
             "--partition", "0"
         ));
-        assertTrue(e.getSuppressed()[0] instanceof UnknownTopicOrPartitionException);
+        assertInstanceOf(UnknownTopicOrPartitionException.class, e.getSuppressed()[0]);
     }
 
     @ClusterTest
diff --git a/tools/src/test/java/org/apache/kafka/tools/MetadataQuorumCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/MetadataQuorumCommandTest.java
index 4002d84b48fd5..aa858f468ac0d 100644
--- a/tools/src/test/java/org/apache/kafka/tools/MetadataQuorumCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/MetadataQuorumCommandTest.java
@@ -37,6 +37,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -158,19 +159,15 @@ public void testCommandConfig() throws IOException {
 
     @ClusterTest(clusterType = Type.ZK, brokers = 1)
     public void testDescribeQuorumInZkMode() {
-        assertTrue(
-            assertThrows(
+        assertInstanceOf(UnsupportedVersionException.class, assertThrows(
                 ExecutionException.class,
                 () -> MetadataQuorumCommand.execute("--bootstrap-server", cluster.bootstrapServers(), "describe", "--status")
-            ).getCause() instanceof UnsupportedVersionException
-        );
+        ).getCause());
 
-        assertTrue(
-            assertThrows(
+        assertInstanceOf(UnsupportedVersionException.class, assertThrows(
                 ExecutionException.class,
                 () -> MetadataQuorumCommand.execute("--bootstrap-server", cluster.bootstrapServers(), "describe", "--replication")
-            ).getCause() instanceof UnsupportedVersionException
-        );
+        ).getCause());
 
     }
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/TopicCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/TopicCommandTest.java
index 6831b0bae4c31..7541d0a1121fc 100644
--- a/tools/src/test/java/org/apache/kafka/tools/TopicCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/TopicCommandTest.java
@@ -50,6 +50,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -233,7 +234,7 @@ public void testDeleteTopicDoesNotRetryThrottlingQuotaExceededException() {
                 "--delete", "--topic", topicName
             })));
 
-        assertTrue(exception.getCause() instanceof ThrottlingQuotaExceededException);
+        assertInstanceOf(ThrottlingQuotaExceededException.class, exception.getCause());
 
         verify(adminClient).deleteTopics(
             argThat((Collection<String> topics) -> topics.equals(Arrays.asList(topicName))),
@@ -262,7 +263,7 @@ public void testCreatePartitionsDoesNotRetryThrottlingQuotaExceededException() {
                 "--alter", "--topic", topicName, "--partitions", "3",
                 "--bootstrap-server", bootstrapServer
             })));
-        assertTrue(exception.getCause() instanceof ThrottlingQuotaExceededException);
+        assertInstanceOf(ThrottlingQuotaExceededException.class, exception.getCause());
 
         verify(adminClient, times(1)).createPartitions(
             argThat(newPartitions -> newPartitions.get(topicName).totalCount() == 3),
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
index 3242b642cdb82..c7d3a82232b76 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptionsTest.java
@@ -32,6 +32,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -392,11 +393,11 @@ public void testCustomPropertyShouldBePassedToConfigureMethod() throws Exception
 
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
-        assertTrue(config.formatter() instanceof DefaultMessageFormatter);
+        assertInstanceOf(DefaultMessageFormatter.class, config.formatter());
         assertTrue(config.formatterArgs().containsKey("key.deserializer.my-props"));
         DefaultMessageFormatter formatter = (DefaultMessageFormatter) config.formatter();
         assertTrue(formatter.keyDeserializer().isPresent());
-        assertTrue(formatter.keyDeserializer().get() instanceof MockDeserializer);
+        assertInstanceOf(MockDeserializer.class, formatter.keyDeserializer().get());
         MockDeserializer keyDeserializer = (MockDeserializer) formatter.keyDeserializer().get();
         assertEquals(1, keyDeserializer.configs.size());
         assertEquals("abc", keyDeserializer.configs.get("my-props"));
@@ -419,11 +420,11 @@ public void testCustomConfigShouldBePassedToConfigureMethod() throws Exception {
 
         ConsoleConsumerOptions config = new ConsoleConsumerOptions(args);
 
-        assertTrue(config.formatter() instanceof DefaultMessageFormatter);
+        assertInstanceOf(DefaultMessageFormatter.class, config.formatter());
         assertTrue(config.formatterArgs().containsKey("key.deserializer.my-props"));
         DefaultMessageFormatter formatter = (DefaultMessageFormatter) config.formatter();
         assertTrue(formatter.keyDeserializer().isPresent());
-        assertTrue(formatter.keyDeserializer().get() instanceof MockDeserializer);
+        assertInstanceOf(MockDeserializer.class, formatter.keyDeserializer().get());
         MockDeserializer keyDeserializer = (MockDeserializer) formatter.keyDeserializer().get();
         assertEquals(1, keyDeserializer.configs.size());
         assertEquals("abc", keyDeserializer.configs.get("my-props"));

From 8c0fafba58d1893b7db086ef845ed9dd0b440b94 Mon Sep 17 00:00:00 2001
From: Manikumar Reddy <manikumar.reddy@gmail.com>
Date: Mon, 18 Mar 2024 16:21:20 +0530
Subject: [PATCH 175/258] MINOR: Update upgrade docs to refer 3.6.2 version

---
 docs/upgrade.html | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/upgrade.html b/docs/upgrade.html
index 33009f062e27a..6488ead802334 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -106,9 +106,9 @@ <h5><a id="upgrade_370_notable" href="#upgrade_370_notable">Notable changes in 3
         </li>
     </ul>
 
-<h4><a id="upgrade_3_6_1" href="#upgrade_3_6_1">Upgrading to 3.6.1 from any version 0.8.x through 3.5.x</a></h4>
+<h4><a id="upgrade_3_6_2" href="#upgrade_3_6_2">Upgrading to 3.6.2 from any version 0.8.x through 3.5.x</a></h4>
 
-    <h5><a id="upgrade_361_zk" href="#upgrade_361_zk">Upgrading ZooKeeper-based clusters</a></h5>
+    <h5><a id="upgrade_362_zk" href="#upgrade_362_zk">Upgrading ZooKeeper-based clusters</a></h5>
     <p><b>If you are upgrading from a version prior to 2.1.x, please see the note in step 5 below about the change to the schema used to store consumer offsets.
         Once you have changed the inter.broker.protocol.version to the latest version, it will not be possible to downgrade to a version prior to 2.1.</b></p>
 
@@ -149,7 +149,7 @@ <h5><a id="upgrade_361_zk" href="#upgrade_361_zk">Upgrading ZooKeeper-based clus
         </li>
     </ol>
 
-    <h5><a id="upgrade_361_kraft" href="#upgrade_361_kraft">Upgrading KRaft-based clusters</a></h5>
+    <h5><a id="upgrade_362_kraft" href="#upgrade_362_kraft">Upgrading KRaft-based clusters</a></h5>
     <p><b>If you are upgrading from a version prior to 3.3.0, please see the note in step 3 below. Once you have changed the metadata.version to the latest version, it will not be possible to downgrade to a version prior to 3.3-IV0.</b></p>
 
     <p><b>For a rolling upgrade:</b></p>

From 34d365fd8a8e387e2650e3d9e492403d6bca9064 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Wed, 20 Mar 2024 14:53:23 +0800
Subject: [PATCH 176/258] KAFKA-16222: desanitize entity name when migrate
 client quotas (#15481)

The entity name is sanitized when it's in Zk mode.
We didn't desanitize it when we migrate client quotas. Add Sanitizer.desanitize to fix it.

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../scala/kafka/zk/ZkMigrationClient.scala    |  4 +++
 .../kafka/zk/ZkMigrationIntegrationTest.scala | 30 ++++++++++++-------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
index ee960bd35f8bf..e69719e7501a3 100644
--- a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
@@ -27,6 +27,7 @@ import org.apache.kafka.common.errors.ControllerMovedException
 import org.apache.kafka.common.metadata._
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.scram.ScramCredential
+import org.apache.kafka.common.utils.Sanitizer
 import org.apache.kafka.common.{TopicIdPartition, Uuid}
 import org.apache.kafka.metadata.DelegationTokenData
 import org.apache.kafka.metadata.PartitionRegistration
@@ -226,6 +227,9 @@ class ZkMigrationClient(
         entityDataList: util.List[ClientQuotaRecord.EntityData],
         quotas: util.Map[String, lang.Double]
       ): Unit = {
+        entityDataList.forEach(entityData => {
+          entityData.setEntityName(Sanitizer.desanitize(entityData.entityName()))
+        })
         val batch = new util.ArrayList[ApiMessageAndVersion]()
         quotas.forEach((key, value) => {
           batch.add(new ApiMessageAndVersion(new ClientQuotaRecord()
diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
index e52bc4e567a4b..d442df33adfdf 100644
--- a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -48,7 +48,7 @@ import org.apache.kafka.raft.RaftConfig
 import org.apache.kafka.security.PasswordEncoder
 import org.apache.kafka.server.ControllerRequestCompletionHandler
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotEquals, assertNotNull, assertTrue, fail}
 import org.junit.jupiter.api.{Assumptions, Timeout}
 import org.junit.jupiter.api.extension.ExtendWith
@@ -227,15 +227,19 @@ class ZkMigrationIntegrationTest {
     createTopicResult.all().get(60, TimeUnit.SECONDS)
 
     val quotas = new util.ArrayList[ClientQuotaAlteration]()
-    quotas.add(new ClientQuotaAlteration(
-      new ClientQuotaEntity(Map("user" -> "user1").asJava),
-      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
-    quotas.add(new ClientQuotaAlteration(
-      new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava),
+    val defaultUserEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> ConfigEntityName.DEFAULT).asJava)
+    quotas.add(new ClientQuotaAlteration(defaultUserEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 900.0)).asJava))
+    val defaultClientIdEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.CLIENT_ID -> ConfigEntityName.DEFAULT).asJava)
+    quotas.add(new ClientQuotaAlteration(defaultClientIdEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 900.0)).asJava))
+    val defaultIpEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> null.asInstanceOf[String]).asJava)
+    quotas.add(new ClientQuotaAlteration(defaultIpEntity, List(new ClientQuotaAlteration.Op("connection_creation_rate", 9.0)).asJava))
+    val userEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> "user/1@prod").asJava)
+    quotas.add(new ClientQuotaAlteration(userEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
+    val userClientEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> "user/1@prod", ClientQuotaEntity.CLIENT_ID -> "client/1@domain").asJava)
+    quotas.add(new ClientQuotaAlteration(userClientEntity,
       List(new ClientQuotaAlteration.Op("consumer_byte_rate", 800.0), new ClientQuotaAlteration.Op("producer_byte_rate", 100.0)).asJava))
-    quotas.add(new ClientQuotaAlteration(
-      new ClientQuotaEntity(Map("ip" -> "8.8.8.8").asJava),
-      List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
+    val ipEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> "8.8.8.8").asJava)
+    quotas.add(new ClientQuotaAlteration(ipEntity, List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
     admin.alterClientQuotas(quotas).all().get(60, TimeUnit.SECONDS)
 
     val zkClient = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient
@@ -271,7 +275,13 @@ class ZkMigrationIntegrationTest {
       assertEquals(10, image.topics().getTopic("test-topic-3").partitions().size())
 
       val clientQuotas = image.clientQuotas().entities()
-      assertEquals(3, clientQuotas.size())
+      assertEquals(6, clientQuotas.size())
+      assertEquals(true, clientQuotas.containsKey(defaultUserEntity))
+      assertEquals(true, clientQuotas.containsKey(defaultClientIdEntity))
+      assertEquals(true, clientQuotas.containsKey(new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> "").asJava))) // default ip
+      assertEquals(true, clientQuotas.containsKey(userEntity))
+      assertEquals(true, clientQuotas.containsKey(userClientEntity))
+      assertEquals(true, clientQuotas.containsKey(ipEntity))
     }
 
     migrationState = migrationClient.releaseControllerLeadership(migrationState)

From b6183a41342c765daec9c88f9d2723a221131960 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Wed, 20 Mar 2024 10:34:45 +0300
Subject: [PATCH 177/258] KAFKA-14589 ConsumerGroupCommand rewritten in java
 (#14471)

This PR contains changes to rewrite ConsumerGroupCommand in java and transfer it to tools module

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 bin/kafka-consumer-groups.sh                  |    2 +-
 bin/windows/kafka-consumer-groups.bat         |    2 +-
 build.gradle                                  |    3 +
 .../kafka/admin/ConsumerGroupCommand.scala    | 1220 ----------------
 .../main/scala/kafka/utils/ToolsUtils.scala   |    2 +-
 .../org/apache/kafka/tools/ToolsUtils.java    |   15 +
 .../consumer/group/ConsumerGroupCommand.java  | 1235 +++++++++++++++++
 .../group/ConsumerGroupCommandOptions.java    |  138 +-
 .../kafka/tools/consumer/group/CsvUtils.java  |  157 +++
 .../tools/consumer/group/GroupState.java      |   35 +
 .../consumer/group/MemberAssignmentState.java |   16 +-
 .../group/PartitionAssignmentState.java       |   28 +-
 .../apache/kafka/tools/ToolsTestUtils.java    |   62 +-
 .../group/AuthorizerIntegrationTest.java      |   27 +-
 .../group/ConsumerGroupCommandTest.java       |   16 +-
 .../group/ConsumerGroupServiceTest.java       |   60 +-
 .../group/DeleteConsumerGroupsTest.java       |   76 +-
 ...tsConsumerGroupCommandIntegrationTest.java |   17 +-
 .../group/DescribeConsumerGroupTest.java      |  301 ++--
 .../consumer/group/ListConsumerGroupTest.java |   28 +-
 .../group/ResetConsumerGroupOffsetTest.java   |   81 +-
 ...SaslClientsWithInvalidCredentialsTest.java |    6 +-
 22 files changed, 1871 insertions(+), 1656 deletions(-)
 delete mode 100755 core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/group/CsvUtils.java
 create mode 100644 tools/src/main/java/org/apache/kafka/tools/consumer/group/GroupState.java

diff --git a/bin/kafka-consumer-groups.sh b/bin/kafka-consumer-groups.sh
index feb063de75693..6dde7d708f7eb 100755
--- a/bin/kafka-consumer-groups.sh
+++ b/bin/kafka-consumer-groups.sh
@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-exec $(dirname $0)/kafka-run-class.sh kafka.admin.ConsumerGroupCommand "$@"
+exec $(dirname $0)/kafka-run-class.sh org.apache.kafka.tools.consumer.group.ConsumerGroupCommand "$@"
diff --git a/bin/windows/kafka-consumer-groups.bat b/bin/windows/kafka-consumer-groups.bat
index e027b9e6bfe59..bdec36be41d1f 100644
--- a/bin/windows/kafka-consumer-groups.bat
+++ b/bin/windows/kafka-consumer-groups.bat
@@ -14,4 +14,4 @@ rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 rem See the License for the specific language governing permissions and
 rem limitations under the License.
 
-"%~dp0kafka-run-class.bat" kafka.admin.ConsumerGroupCommand %*
+"%~dp0kafka-run-class.bat" org.apache.kafka.tools.consumer.group.ConsumerGroupCommand %*
diff --git a/build.gradle b/build.gradle
index fb5cbf5869d6d..88d919b27c638 100644
--- a/build.gradle
+++ b/build.gradle
@@ -1980,6 +1980,9 @@ project(':tools') {
     implementation project(':log4j-appender')
     implementation project(':tools:tools-api')
     implementation libs.argparse4j
+    implementation libs.jacksonDatabind
+    implementation libs.jacksonDataformatCsv
+    implementation libs.jacksonJDK8Datatypes
     implementation libs.slf4jApi
     implementation libs.log4j
     implementation libs.joptSimple
diff --git a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
deleted file mode 100755
index 160b9a70aaec9..0000000000000
--- a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
+++ /dev/null
@@ -1,1220 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.admin
-
-import com.fasterxml.jackson.databind.{ObjectReader, ObjectWriter}
-
-import java.time.{Duration, Instant}
-import java.util.{Collections, Properties}
-import com.fasterxml.jackson.dataformat.csv.CsvMapper
-import com.fasterxml.jackson.module.scala.DefaultScalaModule
-import kafka.utils._
-import kafka.utils.Implicits._
-import org.apache.kafka.clients.admin._
-import org.apache.kafka.clients.consumer.OffsetAndMetadata
-import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.common.{ConsumerGroupState, GroupType, KafkaException, Node, TopicPartition}
-import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
-
-import scala.jdk.CollectionConverters._
-import scala.collection.mutable.ListBuffer
-import scala.collection.{Map, Seq, immutable, mutable}
-import scala.util.{Failure, Success, Try}
-import joptsimple.{OptionException, OptionSpec, OptionSpecBuilder}
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection.immutable.TreeMap
-import scala.reflect.ClassTag
-import org.apache.kafka.common.requests.ListOffsetsResponse
-
-object ConsumerGroupCommand extends Logging {
-
-  def main(args: Array[String]): Unit = {
-
-    val opts = new ConsumerGroupCommandOptions(args)
-    try {
-      opts.checkArgs()
-      CommandLineUtils.maybePrintHelpOrVersion(opts, "This tool helps to list all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.")
-
-      // should have exactly one action
-      val actions = Seq(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt, opts.deleteOffsetsOpt).count(opts.options.has)
-      if (actions != 1)
-        CommandLineUtils.printUsageAndExit(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offsets, --delete-offsets")
-
-      run(opts)
-    } catch {
-      case e: OptionException =>
-        CommandLineUtils.printUsageAndExit(opts.parser, e.getMessage)
-    }
-  }
-
-  def run(opts: ConsumerGroupCommandOptions): Unit = {
-    val consumerGroupService = new ConsumerGroupService(opts)
-    try {
-      if (opts.options.has(opts.listOpt))
-        consumerGroupService.listGroups()
-      else if (opts.options.has(opts.describeOpt))
-        consumerGroupService.describeGroups()
-      else if (opts.options.has(opts.deleteOpt))
-        consumerGroupService.deleteGroups()
-      else if (opts.options.has(opts.resetOffsetsOpt)) {
-        val offsetsToReset = consumerGroupService.resetOffsets()
-        if (opts.options.has(opts.exportOpt)) {
-          val exported = consumerGroupService.exportOffsetsToCsv(offsetsToReset)
-          println(exported)
-        } else
-          printOffsetsToReset(offsetsToReset)
-      }
-      else if (opts.options.has(opts.deleteOffsetsOpt)) {
-        consumerGroupService.deleteOffsets()
-      }
-    } catch {
-      case e: IllegalArgumentException =>
-        CommandLineUtils.printUsageAndExit(opts.parser, e.getMessage)
-      case e: Throwable =>
-        printError(s"Executing consumer group command failed due to ${e.getMessage}", Some(e))
-    } finally {
-      consumerGroupService.close()
-    }
-  }
-
-  def consumerGroupStatesFromString(input: String): Set[ConsumerGroupState] = {
-    val parsedStates = input.split(',').map(s => ConsumerGroupState.parse(s.trim)).toSet
-    if (parsedStates.contains(ConsumerGroupState.UNKNOWN)) {
-      val validStates = ConsumerGroupState.values().filter(_ != ConsumerGroupState.UNKNOWN)
-      throw new IllegalArgumentException(s"Invalid state list '$input'. Valid states are: ${validStates.mkString(", ")}")
-    }
-    parsedStates
-  }
-
-  def consumerGroupTypesFromString(input: String): Set[GroupType] = {
-    val parsedTypes = input.toLowerCase.split(',').map(s => GroupType.parse(s.trim)).toSet
-    if (parsedTypes.contains(GroupType.UNKNOWN)) {
-      val validTypes = GroupType.values().filter(_ != GroupType.UNKNOWN)
-      throw new IllegalArgumentException(s"Invalid types list '$input'. Valid types are: ${validTypes.mkString(", ")}")
-    }
-    parsedTypes
-  }
-
-  val MISSING_COLUMN_VALUE = "-"
-
-  private def printError(msg: String, e: Option[Throwable] = None): Unit = {
-    println(s"\nError: $msg")
-    e.foreach(_.printStackTrace())
-  }
-
-  private def printOffsetsToReset(groupAssignmentsToReset: Map[String, Map[TopicPartition, OffsetAndMetadata]]): Unit = {
-    val format = "%-30s %-30s %-10s %-15s"
-    if (groupAssignmentsToReset.nonEmpty)
-      println("\n" + format.format("GROUP", "TOPIC", "PARTITION", "NEW-OFFSET"))
-    for {
-      (groupId, assignment) <- groupAssignmentsToReset
-      (consumerAssignment, offsetAndMetadata) <- assignment
-    } {
-      println(format.format(
-        groupId,
-        consumerAssignment.topic,
-        consumerAssignment.partition,
-        offsetAndMetadata.offset))
-    }
-  }
-
-  case class PartitionAssignmentState(group: String, coordinator: Option[Node], topic: Option[String],
-                                                partition: Option[Int], offset: Option[Long], lag: Option[Long],
-                                                consumerId: Option[String], host: Option[String],
-                                                clientId: Option[String], logEndOffset: Option[Long])
-
-  private[admin] case class MemberAssignmentState(group: String, consumerId: String, host: String, clientId: String, groupInstanceId: String,
-                                             numPartitions: Int, assignment: List[TopicPartition])
-
-  private[admin] case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
-
-  private[admin] sealed trait CsvRecord
-  private[admin] case class CsvRecordWithGroup(group: String, topic: String, partition: Int, offset: Long) extends CsvRecord
-  private[admin] case class CsvRecordNoGroup(topic: String, partition: Int, offset: Long) extends CsvRecord
-  private[admin] object CsvRecordWithGroup {
-    val fields: Array[String] = Array("group", "topic", "partition", "offset")
-  }
-  private[admin] object CsvRecordNoGroup {
-    val fields: Array[String] = Array("topic", "partition", "offset")
-  }
-  // Example: CsvUtils().readerFor[CsvRecordWithoutGroup]
-  private[admin] case class CsvUtils() {
-    val mapper = new CsvMapper
-    mapper.registerModule(DefaultScalaModule)
-    def readerFor[T <: CsvRecord : ClassTag]: ObjectReader = {
-      val schema = getSchema[T]
-      val clazz = implicitly[ClassTag[T]].runtimeClass
-      mapper.readerFor(clazz).`with`(schema)
-    }
-    def writerFor[T <: CsvRecord : ClassTag]: ObjectWriter = {
-      val schema = getSchema[T]
-      val clazz = implicitly[ClassTag[T]].runtimeClass
-      mapper.writerFor(clazz).`with`(schema)
-    }
-    private def getSchema[T <: CsvRecord : ClassTag] = {
-      val clazz = implicitly[ClassTag[T]].runtimeClass
-
-      val fields =
-        if (classOf[CsvRecordWithGroup] == clazz) CsvRecordWithGroup.fields
-        else if (classOf[CsvRecordNoGroup] == clazz) CsvRecordNoGroup.fields
-        else throw new IllegalStateException(s"Unhandled class $clazz")
-
-      val schema = mapper.schemaFor(clazz).sortedBy(fields: _*)
-      schema
-    }
-  }
-
-  class ConsumerGroupService(val opts: ConsumerGroupCommandOptions,
-                             private[admin] val configOverrides: Map[String, String] = Map.empty) {
-
-    private val adminClient = createAdminClient(configOverrides)
-
-    // We have to make sure it is evaluated once and available
-    private lazy val resetPlanFromFile: Option[Map[String, Map[TopicPartition, OffsetAndMetadata]]] = {
-      if (opts.options.has(opts.resetFromFileOpt)) {
-        val resetPlanPath = opts.options.valueOf(opts.resetFromFileOpt)
-        val resetPlanCsv = Utils.readFileAsString(resetPlanPath)
-        val resetPlan = parseResetPlan(resetPlanCsv)
-        Some(resetPlan)
-      } else None
-    }
-
-    def listGroups(): Unit = {
-      val includeType = opts.options.has(opts.typeOpt)
-      val includeState = opts.options.has(opts.stateOpt)
-
-      if (includeType || includeState) {
-        val types = typeValues()
-        val states = stateValues()
-        val listings = listConsumerGroupsWithFilters(types, states)
-
-        printGroupInfo(listings, includeType, includeState)
-
-      } else {
-        listConsumerGroups().foreach(println(_))
-      }
-    }
-
-    private def stateValues(): Set[ConsumerGroupState] = {
-      val stateValue = opts.options.valueOf(opts.stateOpt)
-      if (stateValue == null || stateValue.isEmpty)
-        Set[ConsumerGroupState]()
-      else
-        consumerGroupStatesFromString(stateValue)
-    }
-
-    private def typeValues(): Set[GroupType] = {
-      val typeValue = opts.options.valueOf(opts.typeOpt)
-      if (typeValue == null || typeValue.isEmpty)
-        Set[GroupType]()
-      else
-        consumerGroupTypesFromString(typeValue)
-    }
-
-    private def printGroupInfo(groups: List[ConsumerGroupListing], includeType: Boolean, includeState: Boolean): Unit = {
-      def groupId(groupListing: ConsumerGroupListing): String = groupListing.groupId
-      def groupType(groupListing: ConsumerGroupListing): String = groupListing.`type`().orElse(GroupType.UNKNOWN).toString
-      def groupState(groupListing: ConsumerGroupListing): String = groupListing.state.orElse(ConsumerGroupState.UNKNOWN).toString
-
-      val maxGroupLen = groups.foldLeft(15)((maxLen, groupListing) => Math.max(maxLen, groupId(groupListing).length)) + 10
-      var format = s"%-${maxGroupLen}s"
-      var header = List("GROUP")
-      var extractors: List[ConsumerGroupListing => String] = List(groupId)
-
-      if (includeType) {
-        header = header :+ "TYPE"
-        extractors = extractors :+ groupType _
-        format += " %-20s"
-      }
-
-      if (includeState) {
-        header = header :+ "STATE"
-        extractors = extractors :+ groupState _
-        format += " %-20s"
-      }
-
-      println(format.format(header: _*))
-
-      groups.foreach { groupListing =>
-        val info = extractors.map(extractor => extractor(groupListing))
-        println(format.format(info: _*))
-      }
-    }
-
-    def listConsumerGroups(): List[String] = {
-      val result = adminClient.listConsumerGroups(withTimeoutMs(new ListConsumerGroupsOptions))
-      val listings = result.all.get.asScala
-      listings.map(_.groupId).toList
-    }
-
-    def listConsumerGroupsWithFilters(types: Set[GroupType], states: Set[ConsumerGroupState]): List[ConsumerGroupListing] = {
-      val listConsumerGroupsOptions = withTimeoutMs(new ListConsumerGroupsOptions())
-      listConsumerGroupsOptions
-        .inStates(states.asJava)
-        .withTypes(types.asJava)
-      val result = adminClient.listConsumerGroups(listConsumerGroupsOptions)
-      result.all.get.asScala.toList
-    }
-
-    private def shouldPrintMemberState(group: String, state: Option[String], numRows: Option[Int]): Boolean = {
-      // numRows contains the number of data rows, if any, compiled from the API call in the caller method.
-      // if it's undefined or 0, there is no relevant group information to display.
-      numRows match {
-        case None =>
-          printError(s"The consumer group '$group' does not exist.")
-          false
-        case Some(num) => state match {
-          case Some("Dead") =>
-            printError(s"Consumer group '$group' does not exist.")
-          case Some("Empty") =>
-            Console.err.println(s"\nConsumer group '$group' has no active members.")
-          case Some("PreparingRebalance") | Some("CompletingRebalance") | Some("Assigning") | Some("Reconciling") =>
-            Console.err.println(s"\nWarning: Consumer group '$group' is rebalancing.")
-          case Some("Stable") =>
-          case other =>
-            // the control should never reach here
-            throw new KafkaException(s"Expected a valid consumer group state, but found '${other.getOrElse("NONE")}'.")
-        }
-        !state.contains("Dead") && num > 0
-      }
-    }
-
-    private def size(colOpt: Option[Seq[Object]]): Option[Int] = colOpt.map(_.size)
-
-    private def printOffsets(offsets: Map[String, (Option[String], Option[Seq[PartitionAssignmentState]])]): Unit = {
-      for ((groupId, (state, assignments)) <- offsets) {
-        if (shouldPrintMemberState(groupId, state, size(assignments))) {
-          // find proper columns width
-          var (maxGroupLen, maxTopicLen, maxConsumerIdLen, maxHostLen) = (15, 15, 15, 15)
-          assignments match {
-            case None => // do nothing
-            case Some(consumerAssignments) =>
-              consumerAssignments.foreach { consumerAssignment =>
-                maxGroupLen = Math.max(maxGroupLen, consumerAssignment.group.length)
-                maxTopicLen = Math.max(maxTopicLen, consumerAssignment.topic.getOrElse(MISSING_COLUMN_VALUE).length)
-                maxConsumerIdLen = Math.max(maxConsumerIdLen, consumerAssignment.consumerId.getOrElse(MISSING_COLUMN_VALUE).length)
-                maxHostLen = Math.max(maxHostLen, consumerAssignment.host.getOrElse(MISSING_COLUMN_VALUE).length)
-              }
-          }
-
-          val format = s"%${-maxGroupLen}s %${-maxTopicLen}s %-10s %-15s %-15s %-15s %${-maxConsumerIdLen}s %${-maxHostLen}s %s"
-          println("\n" + format
-            .format("GROUP", "TOPIC", "PARTITION", "CURRENT-OFFSET", "LOG-END-OFFSET", "LAG", "CONSUMER-ID", "HOST", "CLIENT-ID"))
-
-          assignments match {
-            case None => // do nothing
-            case Some(consumerAssignments) =>
-              consumerAssignments.foreach { consumerAssignment =>
-                println(format.format(
-                  consumerAssignment.group,
-                  consumerAssignment.topic.getOrElse(MISSING_COLUMN_VALUE), consumerAssignment.partition.getOrElse(MISSING_COLUMN_VALUE),
-                  consumerAssignment.offset.getOrElse(MISSING_COLUMN_VALUE), consumerAssignment.logEndOffset.getOrElse(MISSING_COLUMN_VALUE),
-                  consumerAssignment.lag.getOrElse(MISSING_COLUMN_VALUE), consumerAssignment.consumerId.getOrElse(MISSING_COLUMN_VALUE),
-                  consumerAssignment.host.getOrElse(MISSING_COLUMN_VALUE), consumerAssignment.clientId.getOrElse(MISSING_COLUMN_VALUE))
-                )
-              }
-          }
-        }
-      }
-    }
-
-    private def printMembers(members: Map[String, (Option[String], Option[Seq[MemberAssignmentState]])], verbose: Boolean): Unit = {
-      for ((groupId, (state, assignments)) <- members) {
-        if (shouldPrintMemberState(groupId, state, size(assignments))) {
-          // find proper columns width
-          var (maxGroupLen, maxConsumerIdLen, maxGroupInstanceIdLen, maxHostLen, maxClientIdLen, includeGroupInstanceId) = (15, 15, 17, 15, 15, false)
-          assignments match {
-            case None => // do nothing
-            case Some(memberAssignments) =>
-              memberAssignments.foreach { memberAssignment =>
-                maxGroupLen = Math.max(maxGroupLen, memberAssignment.group.length)
-                maxConsumerIdLen = Math.max(maxConsumerIdLen, memberAssignment.consumerId.length)
-                maxGroupInstanceIdLen =  Math.max(maxGroupInstanceIdLen, memberAssignment.groupInstanceId.length)
-                maxHostLen = Math.max(maxHostLen, memberAssignment.host.length)
-                maxClientIdLen = Math.max(maxClientIdLen, memberAssignment.clientId.length)
-                includeGroupInstanceId = includeGroupInstanceId || memberAssignment.groupInstanceId.nonEmpty
-              }
-          }
-
-          val wideFormat = s"%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxGroupInstanceIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s "
-          val shortFormat = s"%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s "
-
-          if (includeGroupInstanceId) {
-            print("\n" + wideFormat
-                .format("GROUP", "CONSUMER-ID", "GROUP-INSTANCE-ID", "HOST", "CLIENT-ID", "#PARTITIONS"))
-          } else {
-            print("\n" + shortFormat
-                .format("GROUP", "CONSUMER-ID", "HOST", "CLIENT-ID", "#PARTITIONS"))
-          }
-          if (verbose)
-            print(s"%s".format("ASSIGNMENT"))
-          println()
-
-          assignments match {
-            case None => // do nothing
-            case Some(memberAssignments) =>
-              memberAssignments.foreach { memberAssignment =>
-                if (includeGroupInstanceId) {
-                  print(wideFormat.format(
-                    memberAssignment.group, memberAssignment.consumerId, memberAssignment.groupInstanceId, memberAssignment.host,
-                    memberAssignment.clientId, memberAssignment.numPartitions))
-                } else {
-                  print(shortFormat.format(
-                    memberAssignment.group, memberAssignment.consumerId, memberAssignment.host, memberAssignment.clientId, memberAssignment.numPartitions))
-                }
-                if (verbose) {
-                  val partitions = memberAssignment.assignment match {
-                    case List() => MISSING_COLUMN_VALUE
-                    case assignment =>
-                      assignment.groupBy(_.topic).map {
-                        case (topic, partitionList) => topic + partitionList.map(_.partition).sorted.mkString("(", ",", ")")
-                      }.toList.sorted.mkString(", ")
-                  }
-                  print(s"%s".format(partitions))
-                }
-                println()
-              }
-          }
-        }
-      }
-    }
-
-    private def printStates(states: Map[String, GroupState]): Unit = {
-      for ((groupId, state) <- states) {
-        if (shouldPrintMemberState(groupId, Some(state.state), Some(1))) {
-          val coordinator = s"${state.coordinator.host}:${state.coordinator.port} (${state.coordinator.idString})"
-          val coordinatorColLen = Math.max(25, coordinator.length)
-          val format = s"\n%${-coordinatorColLen}s %-25s %-20s %-15s %s"
-          print(format.format("GROUP", "COORDINATOR (ID)", "ASSIGNMENT-STRATEGY", "STATE", "#MEMBERS"))
-          print(format.format(state.group, coordinator, state.assignmentStrategy, state.state, state.numMembers))
-          println()
-        }
-      }
-    }
-
-    def describeGroups(): Unit = {
-      val groupIds =
-        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
-        else opts.options.valuesOf(opts.groupOpt).asScala
-      val membersOptPresent = opts.options.has(opts.membersOpt)
-      val stateOptPresent = opts.options.has(opts.stateOpt)
-      val offsetsOptPresent = opts.options.has(opts.offsetsOpt)
-      val subActions = Seq(membersOptPresent, offsetsOptPresent, stateOptPresent).count(_ == true)
-
-      if (subActions == 0 || offsetsOptPresent) {
-        val offsets = collectGroupsOffsets(groupIds)
-        printOffsets(offsets)
-      } else if (membersOptPresent) {
-        val members = collectGroupsMembers(groupIds, opts.options.has(opts.verboseOpt))
-        printMembers(members, opts.options.has(opts.verboseOpt))
-      } else {
-        val states = collectGroupsState(groupIds)
-        printStates(states)
-      }
-    }
-
-    private def collectConsumerAssignment(group: String,
-                                          coordinator: Option[Node],
-                                          topicPartitions: Seq[TopicPartition],
-                                          getPartitionOffset: TopicPartition => Option[Long],
-                                          consumerIdOpt: Option[String],
-                                          hostOpt: Option[String],
-                                          clientIdOpt: Option[String]): Array[PartitionAssignmentState] = {
-      if (topicPartitions.isEmpty) {
-        Array[PartitionAssignmentState](
-          PartitionAssignmentState(group, coordinator, None, None, None, getLag(None, None), consumerIdOpt, hostOpt, clientIdOpt, None)
-        )
-      }
-      else
-        describePartitions(group, coordinator, topicPartitions.sortBy(_.partition), getPartitionOffset, consumerIdOpt, hostOpt, clientIdOpt)
-    }
-
-    private def getLag(offset: Option[Long], logEndOffset: Option[Long]): Option[Long] =
-      offset.filter(_ != -1).flatMap(offset => logEndOffset.map(_ - offset))
-
-    private def describePartitions(group: String,
-                                   coordinator: Option[Node],
-                                   topicPartitions: Seq[TopicPartition],
-                                   getPartitionOffset: TopicPartition => Option[Long],
-                                   consumerIdOpt: Option[String],
-                                   hostOpt: Option[String],
-                                   clientIdOpt: Option[String]): Array[PartitionAssignmentState] = {
-
-      def getDescribePartitionResult(topicPartition: TopicPartition, logEndOffsetOpt: Option[Long]): PartitionAssignmentState = {
-        val offset = getPartitionOffset(topicPartition)
-        PartitionAssignmentState(group, coordinator, Option(topicPartition.topic), Option(topicPartition.partition), offset,
-          getLag(offset, logEndOffsetOpt), consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt)
-      }
-
-      getLogEndOffsets(topicPartitions).map {
-        logEndOffsetResult =>
-          logEndOffsetResult._2 match {
-            case LogOffsetResult.LogOffset(logEndOffset) => getDescribePartitionResult(logEndOffsetResult._1, Some(logEndOffset))
-            case LogOffsetResult.Unknown => getDescribePartitionResult(logEndOffsetResult._1, None)
-            case LogOffsetResult.Ignore => null
-          }
-      }.toArray
-    }
-
-    def resetOffsets(): Map[String, Map[TopicPartition, OffsetAndMetadata]] = {
-      val groupIds =
-        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
-        else opts.options.valuesOf(opts.groupOpt).asScala
-
-      val consumerGroups = adminClient.describeConsumerGroups(
-        groupIds.asJava,
-        withTimeoutMs(new DescribeConsumerGroupsOptions)
-      ).describedGroups()
-
-      val result =
-        consumerGroups.asScala.foldLeft(immutable.Map[String, Map[TopicPartition, OffsetAndMetadata]]()) {
-          case (acc, (groupId, groupDescription)) =>
-            groupDescription.get.state().toString match {
-              case "Empty" | "Dead" =>
-                val partitionsToReset = getPartitionsToReset(groupId)
-                val preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset)
-
-                // Dry-run is the default behavior if --execute is not specified
-                val dryRun = opts.options.has(opts.dryRunOpt) || !opts.options.has(opts.executeOpt)
-                if (!dryRun) {
-                  adminClient.alterConsumerGroupOffsets(
-                    groupId,
-                    preparedOffsets.asJava,
-                    withTimeoutMs(new AlterConsumerGroupOffsetsOptions)
-                  ).all.get
-                }
-                acc.updated(groupId, preparedOffsets)
-              case currentState =>
-                printError(s"Assignments can only be reset if the group '$groupId' is inactive, but the current state is $currentState.")
-                acc.updated(groupId, Map.empty)
-            }
-        }
-      result
-    }
-
-    def deleteOffsets(groupId: String, topics: List[String]): (Errors, Map[TopicPartition, Throwable]) = {
-      val partitionLevelResult = mutable.Map[TopicPartition, Throwable]()
-
-      val (topicWithPartitions, topicWithoutPartitions) = topics.partition(_.contains(":"))
-      val knownPartitions = topicWithPartitions.flatMap(parseTopicsWithPartitions)
-
-      // Get the partitions of topics that the user did not explicitly specify the partitions
-      val describeTopicsResult = adminClient.describeTopics(
-        topicWithoutPartitions.asJava,
-        withTimeoutMs(new DescribeTopicsOptions))
-
-      val unknownPartitions = describeTopicsResult.topicNameValues().asScala.flatMap { case (topic, future) =>
-        Try(future.get()) match {
-          case Success(description) => description.partitions().asScala.map { partition =>
-            new TopicPartition(topic, partition.partition())
-          }
-          case Failure(e) =>
-            partitionLevelResult += new TopicPartition(topic, -1) -> e
-            List.empty
-        }
-      }
-
-      val partitions = knownPartitions ++ unknownPartitions
-
-      val deleteResult = adminClient.deleteConsumerGroupOffsets(
-        groupId,
-        partitions.toSet.asJava,
-        withTimeoutMs(new DeleteConsumerGroupOffsetsOptions)
-      )
-
-      var topLevelException = Errors.NONE
-      Try(deleteResult.all.get) match {
-        case Success(_) =>
-        case Failure(e) => topLevelException = Errors.forException(e.getCause)
-      }
-
-      partitions.foreach { partition =>
-        Try(deleteResult.partitionResult(partition).get()) match {
-          case Success(_) => partitionLevelResult += partition -> null
-          case Failure(e) => partitionLevelResult += partition -> e
-        }
-      }
-
-      (topLevelException, partitionLevelResult)
-    }
-
-    def deleteOffsets(): Unit = {
-      val groupId = opts.options.valueOf(opts.groupOpt)
-      val topics = opts.options.valuesOf(opts.topicOpt).asScala.toList
-
-      val (topLevelResult, partitionLevelResult) = deleteOffsets(groupId, topics)
-
-      topLevelResult match {
-        case Errors.NONE =>
-          println(s"Request succeed for deleting offsets with topic ${topics.mkString(", ")} group $groupId")
-        case Errors.INVALID_GROUP_ID =>
-          printError(s"'$groupId' is not valid.")
-        case Errors.GROUP_ID_NOT_FOUND =>
-          printError(s"'$groupId' does not exist.")
-        case Errors.GROUP_AUTHORIZATION_FAILED =>
-          printError(s"Access to '$groupId' is not authorized.")
-        case Errors.NON_EMPTY_GROUP =>
-          printError(s"Deleting offsets of a consumer group '$groupId' is forbidden if the group is not empty.")
-        case Errors.GROUP_SUBSCRIBED_TO_TOPIC |
-             Errors.TOPIC_AUTHORIZATION_FAILED |
-             Errors.UNKNOWN_TOPIC_OR_PARTITION =>
-          printError(s"Encounter some partition level error, see the follow-up details:")
-        case _ =>
-          printError(s"Encounter some unknown error: $topLevelResult")
-      }
-
-      val format = "%-30s %-15s %-15s"
-      println("\n" + format.format("TOPIC", "PARTITION", "STATUS"))
-      partitionLevelResult.toList.sortBy(t => t._1.topic + t._1.partition.toString).foreach { case (tp, error) =>
-        println(format.format(
-          tp.topic,
-          if (tp.partition >= 0) tp.partition else "Not Provided",
-          if (error != null) s"Error: ${error.getMessage}" else "Successful"
-        ))
-      }
-    }
-
-    private[admin] def describeConsumerGroups(groupIds: Seq[String]): mutable.Map[String, ConsumerGroupDescription] = {
-      adminClient.describeConsumerGroups(
-        groupIds.asJava,
-        withTimeoutMs(new DescribeConsumerGroupsOptions)
-      ).describedGroups().asScala.map {
-        case (groupId, groupDescriptionFuture) => (groupId, groupDescriptionFuture.get())
-      }
-    }
-
-    /**
-      * Returns the state of the specified consumer group and partition assignment states
-      */
-    def collectGroupOffsets(groupId: String): (Option[String], Option[Seq[PartitionAssignmentState]]) = {
-      collectGroupsOffsets(List(groupId)).getOrElse(groupId, (None, None))
-    }
-
-    /**
-      * Returns states of the specified consumer groups and partition assignment states
-      */
-    private def collectGroupsOffsets(groupIds: Seq[String]): TreeMap[String, (Option[String], Option[Seq[PartitionAssignmentState]])] = {
-      val consumerGroups = describeConsumerGroups(groupIds)
-
-      val groupOffsets = TreeMap[String, (Option[String], Option[Seq[PartitionAssignmentState]])]() ++ (for ((groupId, consumerGroup) <- consumerGroups) yield {
-        val state = consumerGroup.state
-        val committedOffsets = getCommittedOffsets(groupId)
-        // The admin client returns `null` as a value to indicate that there is not committed offset for a partition.
-        def getPartitionOffset(tp: TopicPartition): Option[Long] = committedOffsets.get(tp).filter(_ != null).map(_.offset)
-        var assignedTopicPartitions = ListBuffer[TopicPartition]()
-        val rowsWithConsumer = consumerGroup.members.asScala.filterNot(_.assignment.topicPartitions.isEmpty).toSeq
-          .sortBy(_.assignment.topicPartitions.size)(Ordering[Int].reverse).flatMap { consumerSummary =>
-          val topicPartitions = consumerSummary.assignment.topicPartitions.asScala
-          assignedTopicPartitions = assignedTopicPartitions ++ topicPartitions
-          collectConsumerAssignment(groupId, Option(consumerGroup.coordinator), topicPartitions.toList,
-            getPartitionOffset, Some(s"${consumerSummary.consumerId}"), Some(s"${consumerSummary.host}"),
-            Some(s"${consumerSummary.clientId}"))
-        }
-        val unassignedPartitions = committedOffsets.filterNot { case (tp, _) => assignedTopicPartitions.contains(tp) }
-        val rowsWithoutConsumer = if (unassignedPartitions.nonEmpty) {
-          collectConsumerAssignment(
-            groupId,
-            Option(consumerGroup.coordinator),
-            unassignedPartitions.keySet.toSeq,
-            getPartitionOffset,
-            Some(MISSING_COLUMN_VALUE),
-            Some(MISSING_COLUMN_VALUE),
-            Some(MISSING_COLUMN_VALUE)).toSeq
-        } else
-          Seq.empty
-
-        groupId -> (Some(state.toString), Some(rowsWithConsumer ++ rowsWithoutConsumer))
-      }).toMap
-
-      groupOffsets
-    }
-
-    private[admin] def collectGroupMembers(groupId: String, verbose: Boolean): (Option[String], Option[Seq[MemberAssignmentState]]) = {
-      collectGroupsMembers(Seq(groupId), verbose)(groupId)
-    }
-
-    private[admin] def collectGroupsMembers(groupIds: Seq[String], verbose: Boolean): TreeMap[String, (Option[String], Option[Seq[MemberAssignmentState]])] = {
-      val consumerGroups = describeConsumerGroups(groupIds)
-      TreeMap[String, (Option[String], Option[Seq[MemberAssignmentState]])]() ++ (for ((groupId, consumerGroup) <- consumerGroups) yield {
-        val state = consumerGroup.state.toString
-        val memberAssignmentStates = consumerGroup.members().asScala.map(consumer =>
-          MemberAssignmentState(
-            groupId,
-            consumer.consumerId,
-            consumer.host,
-            consumer.clientId,
-            consumer.groupInstanceId.orElse(""),
-            consumer.assignment.topicPartitions.size(),
-            if (verbose) consumer.assignment.topicPartitions.asScala.toList else List()
-          )).toList
-        groupId -> (Some(state), Option(memberAssignmentStates))
-      }).toMap
-    }
-
-    private[admin] def collectGroupState(groupId: String): GroupState = {
-      collectGroupsState(Seq(groupId))(groupId)
-    }
-
-    private[admin] def collectGroupsState(groupIds: Seq[String]): TreeMap[String, GroupState] = {
-      val consumerGroups = describeConsumerGroups(groupIds)
-      TreeMap[String, GroupState]() ++ (for ((groupId, groupDescription) <- consumerGroups) yield {
-        groupId -> GroupState(
-          groupId,
-          groupDescription.coordinator,
-          groupDescription.partitionAssignor(),
-          groupDescription.state.toString,
-          groupDescription.members().size
-        )
-      }).toMap
-    }
-
-    private def getLogEndOffsets(topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
-      val endOffsets = topicPartitions.map { topicPartition =>
-        topicPartition -> OffsetSpec.latest
-      }.toMap
-      val offsets = adminClient.listOffsets(
-        endOffsets.asJava,
-        withTimeoutMs(new ListOffsetsOptions)
-      ).all.get
-      topicPartitions.map { topicPartition =>
-        Option(offsets.get(topicPartition)) match {
-          case Some(listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
-          case _ => topicPartition -> LogOffsetResult.Unknown
-        }
-      }.toMap
-    }
-
-    private def getLogStartOffsets(topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
-      val startOffsets = topicPartitions.map { topicPartition =>
-        topicPartition -> OffsetSpec.earliest
-      }.toMap
-      val offsets = adminClient.listOffsets(
-        startOffsets.asJava,
-        withTimeoutMs(new ListOffsetsOptions)
-      ).all.get
-      topicPartitions.map { topicPartition =>
-        Option(offsets.get(topicPartition)) match {
-          case Some(listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
-          case _ => topicPartition -> LogOffsetResult.Unknown
-        }
-      }.toMap
-    }
-
-    private def getLogTimestampOffsets(topicPartitions: Seq[TopicPartition], timestamp: java.lang.Long): Map[TopicPartition, LogOffsetResult] = {
-      val timestampOffsets = topicPartitions.map { topicPartition =>
-        topicPartition -> OffsetSpec.forTimestamp(timestamp)
-      }.toMap
-      val offsets = adminClient.listOffsets(
-        timestampOffsets.asJava,
-        withTimeoutMs(new ListOffsetsOptions)
-      ).all.get
-      val (successfulOffsetsForTimes, unsuccessfulOffsetsForTimes) =
-        offsets.asScala.partition(_._2.offset != ListOffsetsResponse.UNKNOWN_OFFSET)
-
-      val successfulLogTimestampOffsets = successfulOffsetsForTimes.map {
-        case (topicPartition, listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
-      }.toMap
-
-      unsuccessfulOffsetsForTimes.foreach { entry =>
-        println(s"\nWarn: Partition " + entry._1.partition() + " from topic " + entry._1.topic() +
-          " is empty. Falling back to latest known offset.")
-      }
-
-      successfulLogTimestampOffsets ++ getLogEndOffsets(unsuccessfulOffsetsForTimes.keySet.toSeq)
-    }
-
-    def close(): Unit = {
-      adminClient.close()
-    }
-
-    // Visibility for testing
-    protected def createAdminClient(configOverrides: Map[String, String]): Admin = {
-      val props = if (opts.options.has(opts.commandConfigOpt)) Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt)) else new Properties()
-      props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
-      configOverrides.forKeyValue { (k, v) => props.put(k, v)}
-      Admin.create(props)
-    }
-
-    private def withTimeoutMs [T <: AbstractOptions[T]] (options : T) =  {
-      val t = opts.options.valueOf(opts.timeoutMsOpt).intValue()
-      options.timeoutMs(t)
-    }
-
-    private def parseTopicsWithPartitions(topicArg: String): Seq[TopicPartition] = {
-      def partitionNum(partition: String): Int = {
-        try {
-          partition.toInt
-        } catch {
-          case _: NumberFormatException =>
-            throw new IllegalArgumentException(s"Invalid partition '$partition' specified in topic arg '$topicArg''")
-        }
-      }
-      topicArg.split(":") match {
-        case Array(topic, partitions) =>
-          partitions.split(",").map(partition => new TopicPartition(topic, partitionNum(partition)))
-        case _ =>
-          throw new IllegalArgumentException(s"Invalid topic arg '$topicArg', expected topic name and partitions")
-      }
-    }
-
-    private def parseTopicPartitionsToReset(topicArgs: Seq[String]): Seq[TopicPartition] = {
-      val (topicsWithPartitions, topics) = topicArgs.partition(_.contains(":"))
-      val specifiedPartitions = topicsWithPartitions.flatMap(parseTopicsWithPartitions)
-
-      val unspecifiedPartitions = if (topics.nonEmpty) {
-        val descriptionMap = adminClient.describeTopics(
-          topics.asJava,
-          withTimeoutMs(new DescribeTopicsOptions)
-        ).allTopicNames().get.asScala
-        descriptionMap.flatMap { case (topic, description) =>
-          description.partitions().asScala.map { tpInfo =>
-            new TopicPartition(topic, tpInfo.partition)
-          }
-        }
-      } else
-        Seq.empty
-      specifiedPartitions ++ unspecifiedPartitions
-    }
-
-    private def getPartitionsToReset(groupId: String): Seq[TopicPartition] = {
-      if (opts.options.has(opts.allTopicsOpt)) {
-        getCommittedOffsets(groupId).keys.toSeq
-      } else if (opts.options.has(opts.topicOpt)) {
-        val topics = opts.options.valuesOf(opts.topicOpt).asScala
-        parseTopicPartitionsToReset(topics)
-      } else {
-        if (opts.options.has(opts.resetFromFileOpt))
-          Nil
-        else
-          ToolsUtils.printUsageAndExit(opts.parser, "One of the reset scopes should be defined: --all-topics, --topic.")
-      }
-    }
-
-    private def getCommittedOffsets(groupId: String): Map[TopicPartition, OffsetAndMetadata] = {
-      adminClient.listConsumerGroupOffsets(
-        Collections.singletonMap(groupId, new ListConsumerGroupOffsetsSpec),
-        withTimeoutMs(new ListConsumerGroupOffsetsOptions())
-      ).partitionsToOffsetAndMetadata(groupId).get().asScala
-    }
-
-    type GroupMetadata = immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]
-    private def parseResetPlan(resetPlanCsv: String): GroupMetadata = {
-      def updateGroupMetadata(group: String, topic: String, partition: Int, offset: Long, acc: GroupMetadata) = {
-        val topicPartition = new TopicPartition(topic, partition)
-        val offsetAndMetadata = new OffsetAndMetadata(offset)
-        val dataMap = acc.getOrElse(group, immutable.Map()).updated(topicPartition, offsetAndMetadata)
-        acc.updated(group, dataMap)
-      }
-      val csvReader = CsvUtils().readerFor[CsvRecordNoGroup]
-      val lines = resetPlanCsv.split("\n")
-      val isSingleGroupQuery = opts.options.valuesOf(opts.groupOpt).size() == 1
-      val isOldCsvFormat = lines.headOption.flatMap(line =>
-        Try(csvReader.readValue[CsvRecordNoGroup](line)).toOption).nonEmpty
-      // Single group CSV format: "topic,partition,offset"
-      val dataMap = if (isSingleGroupQuery && isOldCsvFormat) {
-        val group = opts.options.valueOf(opts.groupOpt)
-        lines.foldLeft(immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]()) { (acc, line) =>
-          val CsvRecordNoGroup(topic, partition, offset) = csvReader.readValue[CsvRecordNoGroup](line)
-          updateGroupMetadata(group, topic, partition, offset, acc)
-        }
-        // Multiple group CSV format: "group,topic,partition,offset"
-      } else {
-        val csvReader = CsvUtils().readerFor[CsvRecordWithGroup]
-        lines.foldLeft(immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]()) { (acc, line) =>
-          val CsvRecordWithGroup(group, topic, partition, offset) = csvReader.readValue[CsvRecordWithGroup](line)
-          updateGroupMetadata(group, topic, partition, offset, acc)
-        }
-      }
-      dataMap
-    }
-
-    private def prepareOffsetsToReset(groupId: String,
-                                      partitionsToReset: Seq[TopicPartition]): Map[TopicPartition, OffsetAndMetadata] = {
-      if (opts.options.has(opts.resetToOffsetOpt)) {
-        val offset = opts.options.valueOf(opts.resetToOffsetOpt)
-        checkOffsetsRange(partitionsToReset.map((_, offset)).toMap).map {
-          case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
-        }
-      } else if (opts.options.has(opts.resetToEarliestOpt)) {
-        val logStartOffsets = getLogStartOffsets(partitionsToReset)
-        partitionsToReset.map { topicPartition =>
-          logStartOffsets.get(topicPartition) match {
-            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
-            case _ => ToolsUtils.printUsageAndExit(opts.parser, s"Error getting starting offset of topic partition: $topicPartition")
-          }
-        }.toMap
-      } else if (opts.options.has(opts.resetToLatestOpt)) {
-        val logEndOffsets = getLogEndOffsets(partitionsToReset)
-        partitionsToReset.map { topicPartition =>
-          logEndOffsets.get(topicPartition) match {
-            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
-            case _ => ToolsUtils.printUsageAndExit(opts.parser, s"Error getting ending offset of topic partition: $topicPartition")
-          }
-        }.toMap
-      } else if (opts.options.has(opts.resetShiftByOpt)) {
-        val currentCommittedOffsets = getCommittedOffsets(groupId)
-        val requestedOffsets = partitionsToReset.map { topicPartition =>
-          val shiftBy = opts.options.valueOf(opts.resetShiftByOpt)
-          val currentOffset = currentCommittedOffsets.getOrElse(topicPartition,
-            throw new IllegalArgumentException(s"Cannot shift offset for partition $topicPartition since there is no current committed offset")).offset
-          (topicPartition, currentOffset + shiftBy)
-        }.toMap
-        checkOffsetsRange(requestedOffsets).map {
-          case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
-        }
-      } else if (opts.options.has(opts.resetToDatetimeOpt)) {
-        val timestamp = Utils.getDateTime(opts.options.valueOf(opts.resetToDatetimeOpt))
-        val logTimestampOffsets = getLogTimestampOffsets(partitionsToReset, timestamp)
-        partitionsToReset.map { topicPartition =>
-          val logTimestampOffset = logTimestampOffsets.get(topicPartition)
-          logTimestampOffset match {
-            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
-            case _ => ToolsUtils.printUsageAndExit(opts.parser, s"Error getting offset by timestamp of topic partition: $topicPartition")
-          }
-        }.toMap
-      } else if (opts.options.has(opts.resetByDurationOpt)) {
-        val duration = opts.options.valueOf(opts.resetByDurationOpt)
-        val durationParsed = Duration.parse(duration)
-        val now = Instant.now()
-        durationParsed.negated().addTo(now)
-        val timestamp = now.minus(durationParsed).toEpochMilli
-        val logTimestampOffsets = getLogTimestampOffsets(partitionsToReset, timestamp)
-        partitionsToReset.map { topicPartition =>
-          val logTimestampOffset = logTimestampOffsets.get(topicPartition)
-          logTimestampOffset match {
-            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
-            case _ => ToolsUtils.printUsageAndExit(opts.parser, s"Error getting offset by timestamp of topic partition: $topicPartition")
-          }
-        }.toMap
-      } else if (resetPlanFromFile.isDefined) {
-        resetPlanFromFile.map(resetPlan => resetPlan.get(groupId).map { resetPlanForGroup =>
-          val requestedOffsets = resetPlanForGroup.keySet.map { topicPartition =>
-            topicPartition -> resetPlanForGroup(topicPartition).offset
-          }.toMap
-          checkOffsetsRange(requestedOffsets).map {
-            case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
-          }
-        } match {
-          case Some(resetPlanForGroup) => resetPlanForGroup
-          case None =>
-            printError(s"No reset plan for group $groupId found")
-            Map[TopicPartition, OffsetAndMetadata]()
-        }).getOrElse(Map.empty)
-      } else if (opts.options.has(opts.resetToCurrentOpt)) {
-        val currentCommittedOffsets = getCommittedOffsets(groupId)
-        val (partitionsToResetWithCommittedOffset, partitionsToResetWithoutCommittedOffset) =
-          partitionsToReset.partition(currentCommittedOffsets.keySet.contains(_))
-
-        val preparedOffsetsForPartitionsWithCommittedOffset = partitionsToResetWithCommittedOffset.map { topicPartition =>
-          (topicPartition, new OffsetAndMetadata(currentCommittedOffsets.get(topicPartition) match {
-            case Some(offset) => offset.offset
-            case None => throw new IllegalStateException(s"Expected a valid current offset for topic partition: $topicPartition")
-          }))
-        }.toMap
-
-        val preparedOffsetsForPartitionsWithoutCommittedOffset = getLogEndOffsets(partitionsToResetWithoutCommittedOffset).map {
-          case (topicPartition, LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
-          case (topicPartition, _) => ToolsUtils.printUsageAndExit(opts.parser, s"Error getting ending offset of topic partition: $topicPartition")
-        }
-
-        preparedOffsetsForPartitionsWithCommittedOffset ++ preparedOffsetsForPartitionsWithoutCommittedOffset
-      } else {
-        ToolsUtils.printUsageAndExit(opts.parser, "Option '%s' requires one of the following scenarios: %s".format(opts.resetOffsetsOpt, opts.allResetOffsetScenarioOpts))
-      }
-    }
-
-    private def checkOffsetsRange(requestedOffsets: Map[TopicPartition, Long]) = {
-      val logStartOffsets = getLogStartOffsets(requestedOffsets.keySet.toSeq)
-      val logEndOffsets = getLogEndOffsets(requestedOffsets.keySet.toSeq)
-      requestedOffsets.map { case (topicPartition, offset) => (topicPartition,
-        logEndOffsets.get(topicPartition) match {
-          case Some(LogOffsetResult.LogOffset(endOffset)) if offset > endOffset =>
-            warn(s"New offset ($offset) is higher than latest offset for topic partition $topicPartition. Value will be set to $endOffset")
-            endOffset
-
-          case Some(_) => logStartOffsets.get(topicPartition) match {
-            case Some(LogOffsetResult.LogOffset(startOffset)) if offset < startOffset =>
-              warn(s"New offset ($offset) is lower than earliest offset for topic partition $topicPartition. Value will be set to $startOffset")
-              startOffset
-
-            case _ => offset
-          }
-
-          case None => // the control should not reach here
-            throw new IllegalStateException(s"Unexpected non-existing offset value for topic partition $topicPartition")
-        })
-      }
-    }
-
-    def exportOffsetsToCsv(assignments: Map[String, Map[TopicPartition, OffsetAndMetadata]]): String = {
-      val isSingleGroupQuery = opts.options.valuesOf(opts.groupOpt).size() == 1
-      val csvWriter =
-        if (isSingleGroupQuery) CsvUtils().writerFor[CsvRecordNoGroup]
-        else CsvUtils().writerFor[CsvRecordWithGroup]
-      val rows = assignments.flatMap { case (groupId, partitionInfo) =>
-        partitionInfo.map { case (k: TopicPartition, v: OffsetAndMetadata) =>
-          val csvRecord =
-            if (isSingleGroupQuery) CsvRecordNoGroup(k.topic, k.partition, v.offset)
-            else CsvRecordWithGroup(groupId, k.topic, k.partition, v.offset)
-          csvWriter.writeValueAsString(csvRecord)
-        }
-      }
-      rows.mkString("")
-    }
-
-    def deleteGroups(): Map[String, Throwable] = {
-      val groupIds =
-        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
-        else opts.options.valuesOf(opts.groupOpt).asScala
-
-      val groupsToDelete = adminClient.deleteConsumerGroups(
-        groupIds.asJava,
-        withTimeoutMs(new DeleteConsumerGroupsOptions)
-      ).deletedGroups().asScala
-
-      val result = groupsToDelete.map { case (g, f) =>
-        Try(f.get) match {
-          case Success(_) => g -> null
-          case Failure(e) => g -> e
-        }
-      }
-
-      val (success, failed) = result.partition {
-        case (_, error) => error == null
-      }
-
-      if (failed.isEmpty) {
-        println(s"Deletion of requested consumer groups (${success.keySet.mkString("'", "', '", "'")}) was successful.")
-      }
-      else {
-        printError("Deletion of some consumer groups failed:")
-        failed.foreach {
-          case (group, error) => println(s"* Group '$group' could not be deleted due to: ${error.toString}")
-        }
-        if (success.nonEmpty)
-          println(s"\nThese consumer groups were deleted successfully: ${success.keySet.mkString("'", "', '", "'")}")
-      }
-
-      result.toMap
-    }
-  }
-
-  sealed trait LogOffsetResult
-
-  private object LogOffsetResult {
-    case class LogOffset(value: Long) extends LogOffsetResult
-    case object Unknown extends LogOffsetResult
-    case object Ignore extends LogOffsetResult
-  }
-
-  class ConsumerGroupCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
-    val BootstrapServerDoc = "REQUIRED: The server(s) to connect to."
-    private val GroupDoc = "The consumer group we wish to act on."
-    private val TopicDoc = "The topic whose consumer group information should be deleted or topic whose should be included in the reset offset process. " +
-      "In `reset-offsets` case, partitions can be specified using this format: `topic1:0,1,2`, where 0,1,2 are the partition to be included in the process. " +
-      "Reset-offsets also supports multiple topic inputs."
-    private val AllTopicsDoc = "Consider all topics assigned to a group in the `reset-offsets` process."
-    private val ListDoc = "List all consumer groups."
-    private val DescribeDoc = "Describe consumer group and list offset lag (number of messages not yet processed) related to given group."
-    private val AllGroupsDoc = "Apply to all consumer groups."
-    val nl: String = System.getProperty("line.separator")
-    private val DeleteDoc = "Pass in groups to delete topic partition offsets and ownership information " +
-      "over the entire consumer group. For instance --group g1 --group g2"
-    private val TimeoutMsDoc = "The timeout that can be set for some use cases. For example, it can be used when describing the group " +
-      "to specify the maximum amount of time in milliseconds to wait before the group stabilizes (when the group is just created, " +
-      "or is going through some changes)."
-    val CommandConfigDoc: String = "Property file containing configs to be passed to Admin Client and Consumer."
-    private val ResetOffsetsDoc = "Reset offsets of consumer group. Supports one consumer group at the time, and instances should be inactive" + nl +
-      "Has 2 execution options: --dry-run (the default) to plan which offsets to reset, and --execute to update the offsets. " +
-      "Additionally, the --export option is used to export the results to a CSV format." + nl +
-      "You must choose one of the following reset specifications: --to-datetime, --by-duration, --to-earliest, " +
-      "--to-latest, --shift-by, --from-file, --to-current, --to-offset." + nl +
-      "To define the scope use --all-topics or --topic. One scope must be specified unless you use '--from-file'."
-    private val DryRunDoc = "Only show results without executing changes on Consumer Groups. Supported operations: reset-offsets."
-    private val ExecuteDoc = "Execute operation. Supported operations: reset-offsets."
-    private val ExportDoc = "Export operation execution to a CSV file. Supported operations: reset-offsets."
-    private val ResetToOffsetDoc = "Reset offsets to a specific offset."
-    private val ResetFromFileDoc = "Reset offsets to values defined in CSV file."
-    private val ResetToDatetimeDoc = "Reset offsets to offset from datetime. Format: 'YYYY-MM-DDTHH:mm:SS.sss'"
-    private val ResetByDurationDoc = "Reset offsets to offset by duration from current timestamp. Format: 'PnDTnHnMnS'"
-    private val ResetToEarliestDoc = "Reset offsets to earliest offset."
-    private val ResetToLatestDoc = "Reset offsets to latest offset."
-    private val ResetToCurrentDoc = "Reset offsets to current offset."
-    private val ResetShiftByDoc = "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative."
-    private val MembersDoc = "Describe members of the group. This option may be used with '--describe' and '--bootstrap-server' options only." + nl +
-      "Example: --bootstrap-server localhost:9092 --describe --group group1 --members"
-    private val VerboseDoc = "Provide additional information, if any, when describing the group. This option may be used " +
-      "with '--offsets'/'--members'/'--state' and '--bootstrap-server' options only." + nl + "Example: --bootstrap-server localhost:9092 --describe --group group1 --members --verbose"
-    private val OffsetsDoc = "Describe the group and list all topic partitions in the group along with their offset lag. " +
-      "This is the default sub-action of and may be used with '--describe' and '--bootstrap-server' options only." + nl +
-      "Example: --bootstrap-server localhost:9092 --describe --group group1 --offsets"
-    private val StateDoc = "When specified with '--describe', includes the state of the group." + nl +
-      "Example: --bootstrap-server localhost:9092 --describe --group group1 --state" + nl +
-      "When specified with '--list', it displays the state of all groups. It can also be used to list groups with specific states." + nl +
-      "Example: --bootstrap-server localhost:9092 --list --state stable,empty" + nl +
-      "This option may be used with '--describe', '--list' and '--bootstrap-server' options only."
-    private val TypeDoc = "When specified with '--list', it displays the types of all the groups. It can also be used to list groups with specific types." + nl +
-      "Example: --bootstrap-server localhost:9092 --list --type classic,consumer" + nl +
-      "This option may be used with the '--list' option only."
-    private val DeleteOffsetsDoc = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics."
-
-    val bootstrapServerOpt: OptionSpec[String] = parser.accepts("bootstrap-server", BootstrapServerDoc)
-                                   .withRequiredArg
-                                   .describedAs("server to connect to")
-                                   .ofType(classOf[String])
-    val groupOpt: OptionSpec[String] = parser.accepts("group", GroupDoc)
-                         .withRequiredArg
-                         .describedAs("consumer group")
-                         .ofType(classOf[String])
-    val topicOpt: OptionSpec[String] = parser.accepts("topic", TopicDoc)
-                         .withRequiredArg
-                         .describedAs("topic")
-                         .ofType(classOf[String])
-    val allTopicsOpt: OptionSpecBuilder = parser.accepts("all-topics", AllTopicsDoc)
-    val listOpt: OptionSpecBuilder = parser.accepts("list", ListDoc)
-    val describeOpt: OptionSpecBuilder = parser.accepts("describe", DescribeDoc)
-    val allGroupsOpt: OptionSpecBuilder = parser.accepts("all-groups", AllGroupsDoc)
-    val deleteOpt: OptionSpecBuilder = parser.accepts("delete", DeleteDoc)
-    val timeoutMsOpt: OptionSpec[Long] = parser.accepts("timeout", TimeoutMsDoc)
-                             .withRequiredArg
-                             .describedAs("timeout (ms)")
-                             .ofType(classOf[Long])
-                             .defaultsTo(5000)
-    val commandConfigOpt: OptionSpec[String] = parser.accepts("command-config", CommandConfigDoc)
-                                  .withRequiredArg
-                                  .describedAs("command config property file")
-                                  .ofType(classOf[String])
-    val resetOffsetsOpt: OptionSpecBuilder = parser.accepts("reset-offsets", ResetOffsetsDoc)
-    val deleteOffsetsOpt: OptionSpecBuilder = parser.accepts("delete-offsets", DeleteOffsetsDoc)
-    val dryRunOpt: OptionSpecBuilder = parser.accepts("dry-run", DryRunDoc)
-    val executeOpt: OptionSpecBuilder = parser.accepts("execute", ExecuteDoc)
-    val exportOpt: OptionSpecBuilder = parser.accepts("export", ExportDoc)
-    val resetToOffsetOpt: OptionSpec[Long] = parser.accepts("to-offset", ResetToOffsetDoc)
-                           .withRequiredArg()
-                           .describedAs("offset")
-                           .ofType(classOf[Long])
-    val resetFromFileOpt: OptionSpec[String] = parser.accepts("from-file", ResetFromFileDoc)
-                                 .withRequiredArg()
-                                 .describedAs("path to CSV file")
-                                 .ofType(classOf[String])
-    val resetToDatetimeOpt: OptionSpec[String] = parser.accepts("to-datetime", ResetToDatetimeDoc)
-                                   .withRequiredArg()
-                                   .describedAs("datetime")
-                                   .ofType(classOf[String])
-    val resetByDurationOpt: OptionSpec[String] = parser.accepts("by-duration", ResetByDurationDoc)
-                                   .withRequiredArg()
-                                   .describedAs("duration")
-                                   .ofType(classOf[String])
-    val resetToEarliestOpt: OptionSpecBuilder = parser.accepts("to-earliest", ResetToEarliestDoc)
-    val resetToLatestOpt: OptionSpecBuilder = parser.accepts("to-latest", ResetToLatestDoc)
-    val resetToCurrentOpt: OptionSpecBuilder = parser.accepts("to-current", ResetToCurrentDoc)
-    val resetShiftByOpt: OptionSpec[Long] = parser.accepts("shift-by", ResetShiftByDoc)
-                             .withRequiredArg()
-                             .describedAs("number-of-offsets")
-                             .ofType(classOf[Long])
-    val membersOpt: OptionSpecBuilder = parser.accepts("members", MembersDoc)
-                           .availableIf(describeOpt)
-    val verboseOpt: OptionSpecBuilder = parser.accepts("verbose", VerboseDoc)
-                           .availableIf(describeOpt)
-    val offsetsOpt: OptionSpecBuilder = parser.accepts("offsets", OffsetsDoc)
-                           .availableIf(describeOpt)
-    val stateOpt: OptionSpec[String] = parser.accepts("state", StateDoc)
-                         .availableIf(describeOpt, listOpt)
-                         .withOptionalArg()
-                         .ofType(classOf[String])
-    val typeOpt: OptionSpec[String] = parser.accepts("type", TypeDoc)
-                        .availableIf(listOpt)
-                        .withOptionalArg()
-                        .ofType(classOf[String])
-
-    options = parser.parse(args : _*)
-
-    private val allGroupSelectionScopeOpts = immutable.Set[OptionSpec[_]](groupOpt, allGroupsOpt)
-    private val allConsumerGroupLevelOpts = immutable.Set[OptionSpec[_]](listOpt, describeOpt, deleteOpt, resetOffsetsOpt)
-    val allResetOffsetScenarioOpts: Set[OptionSpec[_]] = immutable.Set[OptionSpec[_]](resetToOffsetOpt, resetShiftByOpt,
-      resetToDatetimeOpt, resetByDurationOpt, resetToEarliestOpt, resetToLatestOpt, resetToCurrentOpt, resetFromFileOpt)
-    private val allDeleteOffsetsOpts = immutable.Set[OptionSpec[_]](groupOpt, topicOpt)
-
-    def checkArgs(): Unit = {
-
-      CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
-
-      if (options.has(describeOpt)) {
-        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $describeOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
-        val mutuallyExclusiveOpts: Set[OptionSpec[_]] = Set(membersOpt, offsetsOpt, stateOpt)
-        if (mutuallyExclusiveOpts.toList.map(o => if (options.has(o)) 1 else 0).sum > 1) {
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $describeOpt takes at most one of these options: ${mutuallyExclusiveOpts.mkString(", ")}")
-        }
-        if (options.has(stateOpt) && options.valueOf(stateOpt) != null)
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $describeOpt does not take a value for $stateOpt")
-      } else {
-        if (options.has(timeoutMsOpt))
-          debug(s"Option $timeoutMsOpt is applicable only when $describeOpt is used.")
-      }
-
-      if (options.has(deleteOpt)) {
-        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $deleteOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
-        if (options.has(topicOpt))
-          CommandLineUtils.printUsageAndExit(parser, s"The consumer does not support topic-specific offset " +
-            "deletion from a consumer group.")
-      }
-
-      if (options.has(deleteOffsetsOpt)) {
-        if (!options.has(groupOpt) || !options.has(topicOpt))
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $deleteOffsetsOpt takes the following options: ${allDeleteOffsetsOpts.mkString(", ")}")
-      }
-
-      if (options.has(resetOffsetsOpt)) {
-        if (options.has(dryRunOpt) && options.has(executeOpt))
-          CommandLineUtils.printUsageAndExit(parser, s"Option $resetOffsetsOpt only accepts one of $executeOpt and $dryRunOpt")
-
-        if (!options.has(dryRunOpt) && !options.has(executeOpt)) {
-          Console.err.println("WARN: No action will be performed as the --execute option is missing." +
-            "In a future major release, the default behavior of this command will be to prompt the user before " +
-            "executing the reset rather than doing a dry run. You should add the --dry-run option explicitly " +
-            "if you are scripting this command and want to keep the current default behavior without prompting.")
-        }
-
-        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
-          CommandLineUtils.printUsageAndExit(parser,
-            s"Option $resetOffsetsOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToOffsetOpt,   (allResetOffsetScenarioOpts - resetToOffsetOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToDatetimeOpt, (allResetOffsetScenarioOpts - resetToDatetimeOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetByDurationOpt, (allResetOffsetScenarioOpts - resetByDurationOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToEarliestOpt, (allResetOffsetScenarioOpts - resetToEarliestOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToLatestOpt,   (allResetOffsetScenarioOpts - resetToLatestOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToCurrentOpt,  (allResetOffsetScenarioOpts - resetToCurrentOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetShiftByOpt,    (allResetOffsetScenarioOpts - resetShiftByOpt).asJava)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetFromFileOpt,   (allResetOffsetScenarioOpts - resetFromFileOpt).asJava)
-      }
-
-      CommandLineUtils.checkInvalidArgs(parser, options, groupOpt, (allGroupSelectionScopeOpts - groupOpt).asJava)
-      CommandLineUtils.checkInvalidArgs(parser, options, groupOpt, (allConsumerGroupLevelOpts - describeOpt - deleteOpt - resetOffsetsOpt).asJava)
-      CommandLineUtils.checkInvalidArgs(parser, options, topicOpt, (allConsumerGroupLevelOpts - deleteOpt - resetOffsetsOpt).asJava )
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/utils/ToolsUtils.scala b/core/src/main/scala/kafka/utils/ToolsUtils.scala
index edd79c00030b9..ffb6214d0b69f 100644
--- a/core/src/main/scala/kafka/utils/ToolsUtils.scala
+++ b/core/src/main/scala/kafka/utils/ToolsUtils.scala
@@ -69,7 +69,7 @@ object ToolsUtils {
   /**
    * This is a simple wrapper around `CommandLineUtils.printUsageAndExit`.
    * It is needed for tools migration (KAFKA-14525), as there is no Java equivalent for return type `Nothing`.
-   * Can be removed once [[kafka.admin.ConsumerGroupCommand]] and [[kafka.tools.ConsoleProducer]] are migrated.
+   * Can be removed once [[kafka.tools.ConsoleProducer]] are migrated.
    *
    * @param parser Command line options parser.
    * @param message Error message.
diff --git a/tools/src/main/java/org/apache/kafka/tools/ToolsUtils.java b/tools/src/main/java/org/apache/kafka/tools/ToolsUtils.java
index 9c6a7a0d1ccb6..394f5078c4690 100644
--- a/tools/src/main/java/org/apache/kafka/tools/ToolsUtils.java
+++ b/tools/src/main/java/org/apache/kafka/tools/ToolsUtils.java
@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.tools;
 
+import joptsimple.OptionParser;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.util.CommandLineUtils;
 
 import java.io.PrintStream;
 import java.util.Arrays;
@@ -155,4 +157,17 @@ public static <T> Set<T> minus(Set<T> set, T...toRemove) {
         return res;
     }
 
+    /**
+     * This is a simple wrapper around `CommandLineUtils.printUsageAndExit`.
+     * It is needed for tools migration (KAFKA-14525), as there is no Java equivalent for return type `Nothing`.
+     * Can be removed once [[kafka.tools.ConsoleConsumer]]
+     * and [[kafka.tools.ConsoleProducer]] are migrated.
+     *
+     * @param parser Command line options parser.
+     * @param message Error message.
+     */
+    public static void printUsageAndExit(OptionParser parser, String message) {
+        CommandLineUtils.printUsageAndExit(parser, message);
+        throw new AssertionError("printUsageAndExit should not return, but it did.");
+    }
 }
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
new file mode 100644
index 0000000000000..2b055de616191
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
@@ -0,0 +1,1235 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectReader;
+import com.fasterxml.jackson.databind.ObjectWriter;
+import joptsimple.OptionException;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.AbstractOptions;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AlterConsumerGroupOffsetsOptions;
+import org.apache.kafka.clients.admin.ConsumerGroupDescription;
+import org.apache.kafka.clients.admin.ConsumerGroupListing;
+import org.apache.kafka.clients.admin.DeleteConsumerGroupOffsetsOptions;
+import org.apache.kafka.clients.admin.DeleteConsumerGroupOffsetsResult;
+import org.apache.kafka.clients.admin.DeleteConsumerGroupsOptions;
+import org.apache.kafka.clients.admin.DescribeConsumerGroupsOptions;
+import org.apache.kafka.clients.admin.DescribeTopicsOptions;
+import org.apache.kafka.clients.admin.DescribeTopicsResult;
+import org.apache.kafka.clients.admin.ListConsumerGroupOffsetsOptions;
+import org.apache.kafka.clients.admin.ListConsumerGroupOffsetsSpec;
+import org.apache.kafka.clients.admin.ListConsumerGroupsOptions;
+import org.apache.kafka.clients.admin.ListConsumerGroupsResult;
+import org.apache.kafka.clients.admin.ListOffsetsOptions;
+import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.MemberDescription;
+import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.GroupType;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.ListOffsetsResponse;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.util.CommandLineUtils;
+import org.apache.kafka.tools.ToolsUtils;
+import org.apache.kafka.tools.Tuple2;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.text.ParseException;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
+import java.util.Properties;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.concurrent.ExecutionException;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.ToIntFunction;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class ConsumerGroupCommand {
+    private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerGroupCommand.class);
+
+    static final String MISSING_COLUMN_VALUE = "-";
+
+    public static void main(String[] args) {
+        ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(args);
+        try {
+            CommandLineUtils.maybePrintHelpOrVersion(opts, "This tool helps to list all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.");
+
+            // should have exactly one action
+            long actions = Stream.of(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt, opts.deleteOffsetsOpt).filter(opts.options::has).count();
+            if (actions != 1)
+                CommandLineUtils.printUsageAndExit(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offsets, --delete-offsets");
+
+            run(opts);
+        } catch (OptionException e) {
+            CommandLineUtils.printUsageAndExit(opts.parser, e.getMessage());
+        }
+    }
+
+    static void run(ConsumerGroupCommandOptions opts) {
+        try (ConsumerGroupService consumerGroupService = new ConsumerGroupService(opts, Collections.emptyMap())) {
+            if (opts.options.has(opts.listOpt))
+                consumerGroupService.listGroups();
+            else if (opts.options.has(opts.describeOpt))
+                consumerGroupService.describeGroups();
+            else if (opts.options.has(opts.deleteOpt))
+                consumerGroupService.deleteGroups();
+            else if (opts.options.has(opts.resetOffsetsOpt)) {
+                Map<String, Map<TopicPartition, OffsetAndMetadata>> offsetsToReset = consumerGroupService.resetOffsets();
+                if (opts.options.has(opts.exportOpt)) {
+                    String exported = consumerGroupService.exportOffsetsToCsv(offsetsToReset);
+                    System.out.println(exported);
+                } else
+                    printOffsetsToReset(offsetsToReset);
+            } else if (opts.options.has(opts.deleteOffsetsOpt)) {
+                consumerGroupService.deleteOffsets();
+            }
+        } catch (IllegalArgumentException e) {
+            CommandLineUtils.printUsageAndExit(opts.parser, e.getMessage());
+        } catch (Throwable e) {
+            printError("Executing consumer group command failed due to " + e.getMessage(), Optional.of(e));
+        }
+    }
+
+    static Set<ConsumerGroupState> consumerGroupStatesFromString(String input) {
+        Set<ConsumerGroupState> parsedStates = Arrays.stream(input.split(",")).map(s -> ConsumerGroupState.parse(s.trim())).collect(Collectors.toSet());
+        if (parsedStates.contains(ConsumerGroupState.UNKNOWN)) {
+            Collection<ConsumerGroupState> validStates = Arrays.stream(ConsumerGroupState.values()).filter(s -> s != ConsumerGroupState.UNKNOWN).collect(Collectors.toList());
+            throw new IllegalArgumentException("Invalid state list '" + input + "'. Valid states are: " + Utils.join(validStates, ", "));
+        }
+        return parsedStates;
+    }
+
+    @SuppressWarnings("Regexp")
+    static Set<GroupType> consumerGroupTypesFromString(String input) {
+        Set<GroupType> parsedTypes = Stream.of(input.toLowerCase().split(",")).map(s -> GroupType.parse(s.trim())).collect(Collectors.toSet());
+        if (parsedTypes.contains(GroupType.UNKNOWN)) {
+            List<String> validTypes = Arrays.stream(GroupType.values()).filter(t -> t != GroupType.UNKNOWN).map(Object::toString).collect(Collectors.toList());
+            throw new IllegalArgumentException("Invalid types list '" + input + "'. Valid types are: " + String.join(", ", validTypes));
+        }
+        return parsedTypes;
+    }
+
+    static void printError(String msg, Optional<Throwable> e) {
+        System.out.println("\nError: " + msg);
+        e.ifPresent(Throwable::printStackTrace);
+    }
+
+    static void printOffsetsToReset(Map<String, Map<TopicPartition, OffsetAndMetadata>> groupAssignmentsToReset) {
+        String format = "%-30s %-30s %-10s %-15s";
+        if (!groupAssignmentsToReset.isEmpty())
+            System.out.printf("\n" + format, "GROUP", "TOPIC", "PARTITION", "NEW-OFFSET");
+
+        groupAssignmentsToReset.forEach((groupId, assignment) ->
+            assignment.forEach((consumerAssignment, offsetAndMetadata) ->
+                System.out.printf(format,
+                    groupId,
+                    consumerAssignment.topic(),
+                    consumerAssignment.partition(),
+                    offsetAndMetadata.offset())));
+    }
+
+    @SuppressWarnings("ClassFanOutComplexity")
+    static class ConsumerGroupService implements AutoCloseable {
+        final ConsumerGroupCommandOptions opts;
+        final Map<String, String> configOverrides;
+        private final Admin adminClient;
+
+        ConsumerGroupService(ConsumerGroupCommandOptions opts, Map<String, String> configOverrides) {
+            this.opts = opts;
+            this.configOverrides = configOverrides;
+            try {
+                this.adminClient = createAdminClient(configOverrides);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        Optional<Map<String, Map<TopicPartition, OffsetAndMetadata>>> resetPlanFromFile() {
+            if (opts.options.has(opts.resetFromFileOpt)) {
+                try {
+                    String resetPlanPath = opts.options.valueOf(opts.resetFromFileOpt);
+                    String resetPlanCsv = Utils.readFileAsString(resetPlanPath);
+                    Map<String, Map<TopicPartition, OffsetAndMetadata>> resetPlan = parseResetPlan(resetPlanCsv);
+                    return Optional.of(resetPlan);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            } else return Optional.empty();
+        }
+
+        void listGroups() throws ExecutionException, InterruptedException {
+            boolean includeType = opts.options.has(opts.typeOpt);
+            boolean includeState = opts.options.has(opts.stateOpt);
+
+            if (includeType || includeState) {
+                Set<GroupType> types = typeValues();
+                Set<ConsumerGroupState> states = stateValues();
+                List<ConsumerGroupListing> listings = listConsumerGroupsWithFilters(types, states);
+
+                printGroupInfo(listings, includeType, includeState);
+            } else {
+                listConsumerGroups().forEach(System.out::println);
+            }
+        }
+
+        private Set<ConsumerGroupState> stateValues() {
+            String stateValue = opts.options.valueOf(opts.stateOpt);
+            return (stateValue == null || stateValue.isEmpty())
+                ? Collections.emptySet()
+                : consumerGroupStatesFromString(stateValue);
+        }
+
+        private Set<GroupType> typeValues() {
+            String typeValue = opts.options.valueOf(opts.typeOpt);
+            return (typeValue == null || typeValue.isEmpty())
+                ? Collections.emptySet()
+                : consumerGroupTypesFromString(typeValue);
+        }
+
+        private void printGroupInfo(List<ConsumerGroupListing> groups, boolean includeType, boolean includeState) {
+            Function<ConsumerGroupListing, String> groupId = ConsumerGroupListing::groupId;
+            Function<ConsumerGroupListing, String> groupType = groupListing -> groupListing.type().orElse(GroupType.UNKNOWN).toString();
+            Function<ConsumerGroupListing, String> groupState = groupListing -> groupListing.state().orElse(ConsumerGroupState.UNKNOWN).toString();
+
+            OptionalInt maybeMax = groups.stream().mapToInt(groupListing -> Math.max(15, groupId.apply(groupListing).length())).max();
+            int maxGroupLen = maybeMax.orElse(15) + 10;
+            String format = "%-" + maxGroupLen + "s";
+            List<String> header = new ArrayList<>();
+            header.add("GROUP");
+            List<Function<ConsumerGroupListing, String>> extractors = new ArrayList<>();
+            extractors.add(groupId);
+
+            if (includeType) {
+                header.add("TYPE");
+                extractors.add(groupType);
+                format += " %-20s";
+            }
+
+            if (includeState) {
+                header.add("STATE");
+                extractors.add(groupState);
+                format += " %-20s";
+            }
+
+            System.out.printf(format + "%n", header.toArray(new Object[0]));
+
+            for (ConsumerGroupListing groupListing : groups) {
+                Object[] info = extractors.stream().map(extractor -> extractor.apply(groupListing)).toArray(Object[]::new);
+                System.out.printf(format + "%n", info);
+            }
+        }
+
+        List<String> listConsumerGroups() {
+            try {
+                ListConsumerGroupsResult result = adminClient.listConsumerGroups(withTimeoutMs(new ListConsumerGroupsOptions()));
+                Collection<ConsumerGroupListing> listings = result.all().get();
+                return listings.stream().map(ConsumerGroupListing::groupId).collect(Collectors.toList());
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        List<ConsumerGroupListing> listConsumerGroupsWithFilters(Set<GroupType> types, Set<ConsumerGroupState> states) throws ExecutionException, InterruptedException {
+            ListConsumerGroupsOptions listConsumerGroupsOptions = withTimeoutMs(new ListConsumerGroupsOptions());
+            listConsumerGroupsOptions
+                .inStates(states)
+                .withTypes(types);
+            ListConsumerGroupsResult result = adminClient.listConsumerGroups(listConsumerGroupsOptions);
+            return new ArrayList<>(result.all().get());
+        }
+
+        private boolean shouldPrintMemberState(String group, Optional<String> state, Optional<Integer> numRows) {
+            // numRows contains the number of data rows, if any, compiled from the API call in the caller method.
+            // if it's undefined or 0, there is no relevant group information to display.
+            if (!numRows.isPresent()) {
+                printError("The consumer group '" + group + "' does not exist.", Optional.empty());
+                return false;
+            }
+
+            int num = numRows.get();
+
+            String state0 = state.orElse("NONE");
+            switch (state0) {
+                case "Dead":
+                    printError("Consumer group '" + group + "' does not exist.", Optional.empty());
+                    break;
+                case "Empty":
+                    System.err.println("\nConsumer group '" + group + "' has no active members.");
+                    break;
+                case "PreparingRebalance":
+                case "CompletingRebalance":
+                case "Assigning":
+                case "Reconciling":
+                    System.err.println("\nWarning: Consumer group '" + group + "' is rebalancing.");
+                    break;
+                case "Stable":
+                    break;
+                default:
+                    // the control should never reach here
+                    throw new KafkaException("Expected a valid consumer group state, but found '" + state0 + "'.");
+            }
+
+            return !state0.contains("Dead") && num > 0;
+        }
+
+        private Optional<Integer> size(Optional<? extends Collection<?>> colOpt) {
+            return colOpt.map(Collection::size);
+        }
+
+        private void printOffsets(Map<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets) {
+            offsets.forEach((groupId, tuple) -> {
+                Optional<String> state = tuple.v1;
+                Optional<Collection<PartitionAssignmentState>> assignments = tuple.v2;
+
+                if (shouldPrintMemberState(groupId, state, size(assignments))) {
+                    String format = printOffsetFormat(assignments);
+
+                    System.out.printf(format, "GROUP", "TOPIC", "PARTITION", "CURRENT-OFFSET", "LOG-END-OFFSET", "LAG", "CONSUMER-ID", "HOST", "CLIENT-ID");
+
+                    if (assignments.isPresent()) {
+                        Collection<PartitionAssignmentState> consumerAssignments = assignments.get();
+                        for (PartitionAssignmentState consumerAssignment : consumerAssignments) {
+                            System.out.printf(format,
+                                consumerAssignment.group,
+                                consumerAssignment.topic.orElse(MISSING_COLUMN_VALUE), consumerAssignment.partition.map(Object::toString).orElse(MISSING_COLUMN_VALUE),
+                                consumerAssignment.offset.map(Object::toString).orElse(MISSING_COLUMN_VALUE), consumerAssignment.logEndOffset.map(Object::toString).orElse(MISSING_COLUMN_VALUE),
+                                consumerAssignment.lag.map(Object::toString).orElse(MISSING_COLUMN_VALUE), consumerAssignment.consumerId.orElse(MISSING_COLUMN_VALUE),
+                                consumerAssignment.host.orElse(MISSING_COLUMN_VALUE), consumerAssignment.clientId.orElse(MISSING_COLUMN_VALUE)
+                            );
+                        }
+                    }
+                }
+            });
+        }
+
+        private static String printOffsetFormat(Optional<Collection<PartitionAssignmentState>> assignments) {
+            // find proper columns width
+            int maxGroupLen = 15, maxTopicLen = 15, maxConsumerIdLen = 15, maxHostLen = 15;
+            if (assignments.isPresent()) {
+                Collection<PartitionAssignmentState> consumerAssignments = assignments.get();
+                for (PartitionAssignmentState consumerAssignment : consumerAssignments) {
+                    maxGroupLen = Math.max(maxGroupLen, consumerAssignment.group.length());
+                    maxTopicLen = Math.max(maxTopicLen, consumerAssignment.topic.orElse(MISSING_COLUMN_VALUE).length());
+                    maxConsumerIdLen = Math.max(maxConsumerIdLen, consumerAssignment.consumerId.orElse(MISSING_COLUMN_VALUE).length());
+                    maxHostLen = Math.max(maxHostLen, consumerAssignment.host.orElse(MISSING_COLUMN_VALUE).length());
+
+                }
+            }
+
+            String format = "\n%" + (-maxGroupLen) + "s %" + (-maxTopicLen) + "s %-10s %-15s %-15s %-15s %" + (-maxConsumerIdLen) + "s %" + (-maxHostLen) + "s %s";
+            return format;
+        }
+
+        private void printMembers(Map<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members, boolean verbose) {
+            members.forEach((groupId, tuple) -> {
+                Optional<String> state = tuple.v1;
+                Optional<Collection<MemberAssignmentState>> assignments = tuple.v2;
+                int maxGroupLen = 15, maxConsumerIdLen = 15, maxGroupInstanceIdLen = 17, maxHostLen = 15, maxClientIdLen = 15;
+                boolean includeGroupInstanceId = false;
+
+                if (shouldPrintMemberState(groupId, state, size(assignments))) {
+                    // find proper columns width
+                    if (assignments.isPresent()) {
+                        for (MemberAssignmentState memberAssignment : assignments.get()) {
+                            maxGroupLen = Math.max(maxGroupLen, memberAssignment.group.length());
+                            maxConsumerIdLen = Math.max(maxConsumerIdLen, memberAssignment.consumerId.length());
+                            maxGroupInstanceIdLen =  Math.max(maxGroupInstanceIdLen, memberAssignment.groupInstanceId.length());
+                            maxHostLen = Math.max(maxHostLen, memberAssignment.host.length());
+                            maxClientIdLen = Math.max(maxClientIdLen, memberAssignment.clientId.length());
+                            includeGroupInstanceId = includeGroupInstanceId || !memberAssignment.groupInstanceId.isEmpty();
+                        }
+                    }
+                }
+
+                String format0 = "%" + -maxGroupLen + "s %" + -maxConsumerIdLen + "s %" + -maxGroupInstanceIdLen + "s %" + -maxHostLen + "s %" + -maxClientIdLen + "s %-15s ";
+                String format1 = "%" + -maxGroupLen + "s %" + -maxConsumerIdLen + "s %" + -maxHostLen + "s %" + -maxClientIdLen + "s %-15s ";
+
+                if (includeGroupInstanceId) {
+                    System.out.printf("\n" + format0, "GROUP", "CONSUMER-ID", "GROUP-INSTANCE-ID", "HOST", "CLIENT-ID", "#PARTITIONS");
+                } else {
+                    System.out.printf("\n" + format1, "GROUP", "CONSUMER-ID", "HOST", "CLIENT-ID", "#PARTITIONS");
+                }
+                if (verbose)
+                    System.out.printf("%s", "ASSIGNMENT");
+                System.out.println();
+
+                if (assignments.isPresent()) {
+                    for (MemberAssignmentState memberAssignment : assignments.get()) {
+                        if (includeGroupInstanceId) {
+                            System.out.printf(format0, memberAssignment.group, memberAssignment.consumerId,
+                                memberAssignment.groupInstanceId, memberAssignment.host, memberAssignment.clientId,
+                                memberAssignment.numPartitions);
+                        } else {
+                            System.out.printf(format1, memberAssignment.group, memberAssignment.consumerId,
+                                memberAssignment.host, memberAssignment.clientId, memberAssignment.numPartitions);
+                        }
+                        if (verbose) {
+                            String partitions;
+
+                            if (memberAssignment.assignment.isEmpty())
+                                partitions = MISSING_COLUMN_VALUE;
+                            else {
+                                Map<String, List<TopicPartition>> grouped = new HashMap<>();
+                                memberAssignment.assignment.forEach(
+                                    tp -> grouped.computeIfAbsent(tp.topic(), key -> new ArrayList<>()).add(tp));
+                                partitions = grouped.values().stream().map(topicPartitions ->
+                                    topicPartitions.stream().map(TopicPartition::partition).map(Object::toString).sorted().collect(Collectors.joining(",", "(", ")"))
+                                ).sorted().collect(Collectors.joining(", "));
+                            }
+                            System.out.printf("%s", partitions);
+                        }
+                        System.out.println();
+                    }
+                }
+            });
+        }
+
+        private void printStates(Map<String, GroupState> states) {
+            states.forEach((groupId, state) -> {
+                if (shouldPrintMemberState(groupId, Optional.of(state.state), Optional.of(1))) {
+                    String coordinator = state.coordinator.host() + ":" + state.coordinator.port() + "  (" + state.coordinator.idString() + ")";
+                    int coordinatorColLen = Math.max(25, coordinator.length());
+
+                    String format = "\n%" + -coordinatorColLen + "s %-25s %-20s %-15s %s";
+
+                    System.out.printf(format, "GROUP", "COORDINATOR (ID)", "ASSIGNMENT-STRATEGY", "STATE", "#MEMBERS");
+                    System.out.printf(format, state.group, coordinator, state.assignmentStrategy, state.state, state.numMembers);
+                    System.out.println();
+                }
+            });
+        }
+
+        void describeGroups() throws Exception {
+            Collection<String> groupIds = opts.options.has(opts.allGroupsOpt)
+                ? listConsumerGroups()
+                : opts.options.valuesOf(opts.groupOpt);
+            boolean membersOptPresent = opts.options.has(opts.membersOpt);
+            boolean stateOptPresent = opts.options.has(opts.stateOpt);
+            boolean offsetsOptPresent = opts.options.has(opts.offsetsOpt);
+            long subActions = Stream.of(membersOptPresent, offsetsOptPresent, stateOptPresent).filter(x -> x).count();
+
+            if (subActions == 0 || offsetsOptPresent) {
+                TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets
+                    = collectGroupsOffsets(groupIds);
+                printOffsets(offsets);
+            } else if (membersOptPresent) {
+                TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members
+                    = collectGroupsMembers(groupIds, opts.options.has(opts.verboseOpt));
+                printMembers(members, opts.options.has(opts.verboseOpt));
+            } else {
+                TreeMap<String, GroupState> states = collectGroupsState(groupIds);
+                printStates(states);
+            }
+        }
+
+        private Collection<PartitionAssignmentState> collectConsumerAssignment(
+            String group,
+            Optional<Node> coordinator,
+            Collection<TopicPartition> topicPartitions,
+            Function<TopicPartition, Optional<Long>> getPartitionOffset,
+            Optional<String> consumerIdOpt,
+            Optional<String> hostOpt,
+            Optional<String> clientIdOpt
+        ) {
+            if (topicPartitions.isEmpty()) {
+                return Collections.singleton(
+                    new PartitionAssignmentState(group, coordinator, Optional.empty(), Optional.empty(), Optional.empty(),
+                        getLag(Optional.empty(), Optional.empty()), consumerIdOpt, hostOpt, clientIdOpt, Optional.empty())
+                );
+            } else {
+                List<TopicPartition> topicPartitionsSorted = topicPartitions.stream().sorted(Comparator.comparingInt(TopicPartition::partition)).collect(Collectors.toList());
+                return describePartitions(group, coordinator, topicPartitionsSorted, getPartitionOffset, consumerIdOpt, hostOpt, clientIdOpt);
+            }
+        }
+
+        private Optional<Long> getLag(Optional<Long> offset, Optional<Long> logEndOffset) {
+            return offset.filter(o -> o != -1).flatMap(offset0 -> logEndOffset.map(end -> end - offset0));
+        }
+
+        private Collection<PartitionAssignmentState> describePartitions(String group,
+                                                              Optional<Node> coordinator,
+                                                              List<TopicPartition> topicPartitions,
+                                                              Function<TopicPartition, Optional<Long>> getPartitionOffset,
+                                                              Optional<String> consumerIdOpt,
+                                                              Optional<String> hostOpt,
+                                                              Optional<String> clientIdOpt) {
+            BiFunction<TopicPartition, Optional<Long>, PartitionAssignmentState> getDescribePartitionResult = (topicPartition, logEndOffsetOpt) -> {
+                Optional<Long> offset = getPartitionOffset.apply(topicPartition);
+                return new PartitionAssignmentState(group, coordinator, Optional.of(topicPartition.topic()),
+                    Optional.of(topicPartition.partition()), offset, getLag(offset, logEndOffsetOpt),
+                    consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt);
+            };
+
+            return getLogEndOffsets(topicPartitions).entrySet().stream().map(logEndOffsetResult -> {
+                if (logEndOffsetResult.getValue() instanceof LogOffset)
+                    return getDescribePartitionResult.apply(
+                        logEndOffsetResult.getKey(),
+                        Optional.of(((LogOffset) logEndOffsetResult.getValue()).value)
+                    );
+                else if (logEndOffsetResult.getValue() instanceof Unknown)
+                    return getDescribePartitionResult.apply(logEndOffsetResult.getKey(), Optional.empty());
+                else if (logEndOffsetResult.getValue() instanceof Ignore)
+                    return null;
+
+                throw new IllegalStateException("Unknown LogOffset subclass: " + logEndOffsetResult.getValue());
+            }).collect(Collectors.toList());
+        }
+
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> resetOffsets() {
+            List<String> groupIds = opts.options.has(opts.allGroupsOpt)
+                ? listConsumerGroups()
+                : opts.options.valuesOf(opts.groupOpt);
+
+            Map<String, KafkaFuture<ConsumerGroupDescription>> consumerGroups = adminClient.describeConsumerGroups(
+                groupIds,
+                withTimeoutMs(new DescribeConsumerGroupsOptions())
+            ).describedGroups();
+
+            Map<String, Map<TopicPartition, OffsetAndMetadata>> result = new HashMap<>();
+
+            consumerGroups.forEach((groupId, groupDescription) -> {
+                try {
+                    String state = groupDescription.get().state().toString();
+                    switch (state) {
+                        case "Empty":
+                        case "Dead":
+                            Collection<TopicPartition> partitionsToReset = getPartitionsToReset(groupId);
+                            Map<TopicPartition, OffsetAndMetadata> preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset);
+
+                            // Dry-run is the default behavior if --execute is not specified
+                            boolean dryRun = opts.options.has(opts.dryRunOpt) || !opts.options.has(opts.executeOpt);
+                            if (!dryRun) {
+                                adminClient.alterConsumerGroupOffsets(
+                                    groupId,
+                                    preparedOffsets,
+                                    withTimeoutMs(new AlterConsumerGroupOffsetsOptions())
+                                ).all().get();
+                            }
+
+                            result.put(groupId, preparedOffsets);
+
+                            break;
+                        default:
+                            printError("Assignments can only be reset if the group '" + groupId + "' is inactive, but the current state is " + state + ".", Optional.empty());
+                            result.put(groupId, Collections.emptyMap());
+                    }
+                } catch (InterruptedException | ExecutionException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            return result;
+        }
+
+        Tuple2<Errors, Map<TopicPartition, Throwable>> deleteOffsets(String groupId, List<String> topics) {
+            Map<TopicPartition, Throwable> partitionLevelResult = new HashMap<>();
+            Set<String> topicWithPartitions = new HashSet<>();
+            Set<String> topicWithoutPartitions = new HashSet<>();
+
+            for (String topic : topics) {
+                if (topic.contains(":"))
+                    topicWithPartitions.add(topic);
+                else
+                    topicWithoutPartitions.add(topic);
+            }
+
+            List<TopicPartition> knownPartitions = topicWithPartitions.stream().flatMap(this::parseTopicsWithPartitions).collect(Collectors.toList());
+
+            // Get the partitions of topics that the user did not explicitly specify the partitions
+            DescribeTopicsResult describeTopicsResult = adminClient.describeTopics(
+                topicWithoutPartitions,
+                withTimeoutMs(new DescribeTopicsOptions()));
+
+            Iterator<TopicPartition> unknownPartitions = describeTopicsResult.topicNameValues().entrySet().stream().flatMap(e -> {
+                String topic = e.getKey();
+                try {
+                    return e.getValue().get().partitions().stream().map(partition ->
+                        new TopicPartition(topic, partition.partition()));
+                } catch (ExecutionException | InterruptedException err) {
+                    partitionLevelResult.put(new TopicPartition(topic, -1), err);
+                    return Stream.empty();
+                }
+            }).iterator();
+
+            Set<TopicPartition> partitions = new HashSet<>(knownPartitions);
+
+            unknownPartitions.forEachRemaining(partitions::add);
+
+            DeleteConsumerGroupOffsetsResult deleteResult = adminClient.deleteConsumerGroupOffsets(
+                groupId,
+                partitions,
+                withTimeoutMs(new DeleteConsumerGroupOffsetsOptions())
+            );
+
+            Errors topLevelException = Errors.NONE;
+
+            try {
+                deleteResult.all().get();
+            } catch (ExecutionException | InterruptedException e) {
+                topLevelException = Errors.forException(e.getCause());
+            }
+
+            partitions.forEach(partition -> {
+                try {
+                    deleteResult.partitionResult(partition).get();
+                    partitionLevelResult.put(partition, null);
+                } catch (ExecutionException | InterruptedException e) {
+                    partitionLevelResult.put(partition, e);
+                }
+            });
+
+            return new Tuple2<>(topLevelException, partitionLevelResult);
+        }
+
+        void deleteOffsets() {
+            String groupId = opts.options.valueOf(opts.groupOpt);
+            List<String> topics = opts.options.valuesOf(opts.topicOpt);
+
+            Tuple2<Errors, Map<TopicPartition, Throwable>> res = deleteOffsets(groupId, topics);
+
+            Errors topLevelResult = res.v1;
+            Map<TopicPartition, Throwable> partitionLevelResult = res.v2;
+
+            switch (topLevelResult) {
+                case NONE:
+                    System.out.println("Request succeed for deleting offsets with topic " + Utils.mkString(topics.stream(), "", "", ", ") + " group " + groupId);
+                    break;
+                case INVALID_GROUP_ID:
+                    printError("'" + groupId + "' is not valid.", Optional.empty());
+                    break;
+                case GROUP_ID_NOT_FOUND:
+                    printError("'" + groupId + "' does not exist.", Optional.empty());
+                    break;
+                case GROUP_AUTHORIZATION_FAILED:
+                    printError("Access to '" + groupId + "' is not authorized.", Optional.empty());
+                    break;
+                case NON_EMPTY_GROUP:
+                    printError("Deleting offsets of a consumer group '" + groupId + "' is forbidden if the group is not empty.", Optional.empty());
+                    break;
+                case GROUP_SUBSCRIBED_TO_TOPIC:
+                case TOPIC_AUTHORIZATION_FAILED:
+                case UNKNOWN_TOPIC_OR_PARTITION:
+                    printError("Encounter some partition level error, see the follow-up details:", Optional.empty());
+                    break;
+                default:
+                    printError("Encounter some unknown error: " + topLevelResult, Optional.empty());
+            }
+
+            String format = "%-30s %-15s %-15s";
+
+            System.out.printf("\n" + format, "TOPIC", "PARTITION", "STATUS");
+            partitionLevelResult.entrySet().stream()
+                .sorted(Comparator.comparing(e -> e.getKey().topic() + e.getKey().partition()))
+                .forEach(e -> {
+                    TopicPartition tp = e.getKey();
+                    Throwable error = e.getValue();
+                    System.out.printf(format,
+                        tp.topic(),
+                        tp.partition() >= 0 ? tp.partition() : "Not Provided",
+                        error != null ? "Error: :" + error.getMessage() : "Successful"
+                    );
+                });
+        }
+
+        Map<String, ConsumerGroupDescription> describeConsumerGroups(Collection<String> groupIds) throws Exception {
+            Map<String, ConsumerGroupDescription> res = new HashMap<>();
+            Map<String, KafkaFuture<ConsumerGroupDescription>> stringKafkaFutureMap = adminClient.describeConsumerGroups(
+                groupIds,
+                withTimeoutMs(new DescribeConsumerGroupsOptions())
+            ).describedGroups();
+
+            for (Map.Entry<String, KafkaFuture<ConsumerGroupDescription>> e : stringKafkaFutureMap.entrySet()) {
+                res.put(e.getKey(), e.getValue().get());
+            }
+            return res;
+        }
+
+        /**
+         * Returns the state of the specified consumer group and partition assignment states
+         */
+        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> collectGroupOffsets(String groupId) throws Exception {
+            return collectGroupsOffsets(Collections.singletonList(groupId)).getOrDefault(groupId, new Tuple2<>(Optional.empty(), Optional.empty()));
+        }
+
+        /**
+         * Returns states of the specified consumer groups and partition assignment states
+         */
+        TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> collectGroupsOffsets(Collection<String> groupIds) throws Exception {
+            Map<String, ConsumerGroupDescription> consumerGroups = describeConsumerGroups(groupIds);
+            TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> groupOffsets = new TreeMap<>();
+
+            consumerGroups.forEach((groupId, consumerGroup) -> {
+                ConsumerGroupState state = consumerGroup.state();
+                Map<TopicPartition, OffsetAndMetadata> committedOffsets = getCommittedOffsets(groupId);
+                // The admin client returns `null` as a value to indicate that there is not committed offset for a partition.
+                Function<TopicPartition, Optional<Long>> getPartitionOffset = tp -> Optional.ofNullable(committedOffsets.get(tp)).map(OffsetAndMetadata::offset);
+                List<TopicPartition> assignedTopicPartitions = new ArrayList<>();
+                Comparator<MemberDescription> comparator =
+                    Comparator.<MemberDescription>comparingInt(m -> m.assignment().topicPartitions().size()).reversed();
+                List<PartitionAssignmentState> rowsWithConsumer = new ArrayList<>();
+                consumerGroup.members().stream().filter(m -> !m.assignment().topicPartitions().isEmpty())
+                    .sorted(comparator)
+                    .forEach(consumerSummary -> {
+                        Set<TopicPartition> topicPartitions = consumerSummary.assignment().topicPartitions();
+                        assignedTopicPartitions.addAll(topicPartitions);
+                        rowsWithConsumer.addAll(collectConsumerAssignment(
+                            groupId,
+                            Optional.of(consumerGroup.coordinator()),
+                            topicPartitions,
+                            getPartitionOffset,
+                            Optional.of(consumerSummary.consumerId()),
+                            Optional.of(consumerSummary.host()),
+                            Optional.of(consumerSummary.clientId()))
+                        );
+                    });
+                Map<TopicPartition, OffsetAndMetadata> unassignedPartitions = new HashMap<>();
+                committedOffsets.entrySet().stream().filter(e -> !assignedTopicPartitions.contains(e.getKey()))
+                    .forEach(e -> unassignedPartitions.put(e.getKey(), e.getValue()));
+                Collection<PartitionAssignmentState> rowsWithoutConsumer = !unassignedPartitions.isEmpty()
+                    ? collectConsumerAssignment(
+                        groupId,
+                        Optional.of(consumerGroup.coordinator()),
+                        unassignedPartitions.keySet(),
+                        getPartitionOffset,
+                        Optional.of(MISSING_COLUMN_VALUE),
+                        Optional.of(MISSING_COLUMN_VALUE),
+                        Optional.of(MISSING_COLUMN_VALUE))
+                    : Collections.emptyList();
+
+                rowsWithConsumer.addAll(rowsWithoutConsumer);
+
+                groupOffsets.put(groupId, new Tuple2<>(Optional.of(state.toString()), Optional.of(rowsWithConsumer)));
+            });
+
+            return groupOffsets;
+        }
+
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> collectGroupMembers(String groupId, boolean verbose) throws Exception {
+            return collectGroupsMembers(Collections.singleton(groupId), verbose).get(groupId);
+        }
+
+        TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> collectGroupsMembers(Collection<String> groupIds, boolean verbose) throws Exception {
+            Map<String, ConsumerGroupDescription> consumerGroups = describeConsumerGroups(groupIds);
+            TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> res = new TreeMap<>();
+
+            consumerGroups.forEach((groupId, consumerGroup) -> {
+                String state = consumerGroup.state().toString();
+                List<MemberAssignmentState> memberAssignmentStates = consumerGroup.members().stream().map(consumer ->
+                    new MemberAssignmentState(
+                        groupId,
+                        consumer.consumerId(),
+                        consumer.host(),
+                        consumer.clientId(),
+                        consumer.groupInstanceId().orElse(""),
+                        consumer.assignment().topicPartitions().size(),
+                        new ArrayList<>(verbose ? consumer.assignment().topicPartitions() : Collections.emptySet())
+                )).collect(Collectors.toList());
+                res.put(groupId, new Tuple2<>(Optional.of(state), Optional.of(memberAssignmentStates)));
+            });
+            return res;
+        }
+
+        GroupState collectGroupState(String groupId) throws Exception {
+            return collectGroupsState(Collections.singleton(groupId)).get(groupId);
+        }
+
+        TreeMap<String, GroupState> collectGroupsState(Collection<String> groupIds) throws Exception {
+            Map<String, ConsumerGroupDescription> consumerGroups = describeConsumerGroups(groupIds);
+            TreeMap<String, GroupState> res = new TreeMap<>();
+            consumerGroups.forEach((groupId, groupDescription) ->
+                res.put(groupId, new GroupState(
+                    groupId,
+                    groupDescription.coordinator(),
+                    groupDescription.partitionAssignor(),
+                    groupDescription.state().toString(),
+                    groupDescription.members().size()
+            )));
+            return res;
+        }
+
+        private Map<TopicPartition, LogOffsetResult> getLogEndOffsets(Collection<TopicPartition> topicPartitions) {
+            return getLogOffsets(topicPartitions, OffsetSpec.latest());
+        }
+
+        private Map<TopicPartition, LogOffsetResult> getLogStartOffsets(Collection<TopicPartition> topicPartitions) {
+            return getLogOffsets(topicPartitions, OffsetSpec.earliest());
+        }
+
+        private Map<TopicPartition, LogOffsetResult> getLogOffsets(Collection<TopicPartition> topicPartitions, OffsetSpec offsetSpec) {
+            try {
+                Map<TopicPartition, OffsetSpec> startOffsets = topicPartitions.stream()
+                    .collect(Collectors.toMap(Function.identity(), tp -> offsetSpec));
+
+                Map<TopicPartition, ListOffsetsResultInfo> offsets = adminClient.listOffsets(
+                    startOffsets,
+                    withTimeoutMs(new ListOffsetsOptions())
+                ).all().get();
+
+                return topicPartitions.stream().collect(Collectors.toMap(
+                    Function.identity(),
+                    tp -> offsets.containsKey(tp)
+                        ? new LogOffset(offsets.get(tp).offset())
+                        : new Unknown()
+                ));
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        private Map<TopicPartition, LogOffsetResult> getLogTimestampOffsets(Collection<TopicPartition> topicPartitions, long timestamp) {
+            try {
+                Map<TopicPartition, OffsetSpec> timestampOffsets = topicPartitions.stream()
+                    .collect(Collectors.toMap(Function.identity(), tp -> OffsetSpec.forTimestamp(timestamp)));
+
+                Map<TopicPartition, ListOffsetsResultInfo> offsets = adminClient.listOffsets(
+                    timestampOffsets,
+                    withTimeoutMs(new ListOffsetsOptions())
+                ).all().get();
+
+                Map<TopicPartition, ListOffsetsResultInfo> successfulOffsetsForTimes = new HashMap<>();
+                Map<TopicPartition, ListOffsetsResultInfo> unsuccessfulOffsetsForTimes = new HashMap<>();
+
+                offsets.forEach((tp, offsetsResultInfo) -> {
+                    if (offsetsResultInfo.offset() != ListOffsetsResponse.UNKNOWN_OFFSET)
+                        successfulOffsetsForTimes.put(tp, offsetsResultInfo);
+                    else
+                        unsuccessfulOffsetsForTimes.put(tp, offsetsResultInfo);
+                });
+
+                Map<TopicPartition, LogOffsetResult> successfulLogTimestampOffsets = successfulOffsetsForTimes.entrySet().stream()
+                    .collect(Collectors.toMap(Map.Entry::getKey, e -> new LogOffset(e.getValue().offset())));
+
+                unsuccessfulOffsetsForTimes.forEach((tp, offsetResultInfo) ->
+                    System.out.println("\nWarn: Partition " + tp.partition() + " from topic " + tp.topic() +
+                    " is empty. Falling back to latest known offset."));
+
+                successfulLogTimestampOffsets.putAll(getLogEndOffsets(unsuccessfulOffsetsForTimes.keySet()));
+
+                return successfulLogTimestampOffsets;
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @Override
+        public void close() {
+            adminClient.close();
+        }
+
+        // Visibility for testing
+        protected Admin createAdminClient(Map<String, String> configOverrides) throws IOException {
+            Properties props = opts.options.has(opts.commandConfigOpt) ? Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt)) : new Properties();
+            props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt));
+            props.putAll(configOverrides);
+            return Admin.create(props);
+        }
+
+        private <T extends AbstractOptions<T>> T withTimeoutMs(T options) {
+            int t = opts.options.valueOf(opts.timeoutMsOpt).intValue();
+            return options.timeoutMs(t);
+        }
+
+        private Stream<TopicPartition> parseTopicsWithPartitions(String topicArg) {
+            ToIntFunction<String> partitionNum = partition -> {
+                try {
+                    return Integer.parseInt(partition);
+                } catch (NumberFormatException e) {
+                    throw new IllegalArgumentException("Invalid partition '" + partition + "' specified in topic arg '" + topicArg + "''");
+                }
+            };
+
+            String[] arr = topicArg.split(":");
+
+            if (arr.length != 2)
+                throw new IllegalArgumentException("Invalid topic arg '" + topicArg + "', expected topic name and partitions");
+
+            String topic = arr[0];
+            String partitions = arr[1];
+
+            return Arrays.stream(partitions.split(",")).
+                map(partition -> new TopicPartition(topic, partitionNum.applyAsInt(partition)));
+        }
+
+        private List<TopicPartition> parseTopicPartitionsToReset(List<String> topicArgs) throws ExecutionException, InterruptedException {
+            List<String> topicsWithPartitions = new ArrayList<>();
+            List<String> topics = new ArrayList<>();
+
+            topicArgs.forEach(topicArg -> {
+                if (topicArg.contains(":"))
+                    topicsWithPartitions.add(topicArg);
+                else
+                    topics.add(topicArg);
+            });
+
+            List<TopicPartition> specifiedPartitions = topicsWithPartitions.stream().flatMap(this::parseTopicsWithPartitions).collect(Collectors.toList());
+
+            List<TopicPartition> unspecifiedPartitions = new ArrayList<>();
+
+            if (!topics.isEmpty()) {
+                Map<String, TopicDescription> descriptionMap = adminClient.describeTopics(
+                    topics,
+                    withTimeoutMs(new DescribeTopicsOptions())
+                ).allTopicNames().get();
+
+                descriptionMap.forEach((topic, description) ->
+                    description.partitions().forEach(tpInfo -> unspecifiedPartitions.add(new TopicPartition(topic, tpInfo.partition())))
+                );
+            }
+
+            specifiedPartitions.addAll(unspecifiedPartitions);
+
+            return specifiedPartitions;
+        }
+
+        private Collection<TopicPartition> getPartitionsToReset(String groupId) throws ExecutionException, InterruptedException {
+            if (opts.options.has(opts.allTopicsOpt)) {
+                return getCommittedOffsets(groupId).keySet();
+            } else if (opts.options.has(opts.topicOpt)) {
+                List<String> topics = opts.options.valuesOf(opts.topicOpt);
+                return parseTopicPartitionsToReset(topics);
+            } else {
+                if (!opts.options.has(opts.resetFromFileOpt))
+                    CommandLineUtils.printUsageAndExit(opts.parser, "One of the reset scopes should be defined: --all-topics, --topic.");
+
+                return Collections.emptyList();
+            }
+        }
+
+        private Map<TopicPartition, OffsetAndMetadata> getCommittedOffsets(String groupId) {
+            try {
+                return adminClient.listConsumerGroupOffsets(
+                    Collections.singletonMap(groupId, new ListConsumerGroupOffsetsSpec()),
+                    withTimeoutMs(new ListConsumerGroupOffsetsOptions())
+                ).partitionsToOffsetAndMetadata(groupId).get();
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        private Map<String, Map<TopicPartition, OffsetAndMetadata>> parseResetPlan(String resetPlanCsv) {
+            ObjectReader csvReader = CsvUtils.readerFor(CsvUtils.CsvRecordNoGroup.class);
+            String[] lines = resetPlanCsv.split("\n");
+            boolean isSingleGroupQuery = opts.options.valuesOf(opts.groupOpt).size() == 1;
+            boolean isOldCsvFormat = false;
+            try {
+                if (lines.length > 0) {
+                    csvReader.readValue(lines[0], CsvUtils.CsvRecordNoGroup.class);
+                    isOldCsvFormat = true;
+                }
+            } catch (IOException e) {
+                e.printStackTrace();
+                // Ignore.
+            }
+
+            Map<String, Map<TopicPartition, OffsetAndMetadata>> dataMap = new HashMap<>();
+
+            try {
+                // Single group CSV format: "topic,partition,offset"
+                if (isSingleGroupQuery && isOldCsvFormat) {
+                    String group = opts.options.valueOf(opts.groupOpt);
+                    for (String line : lines) {
+                        CsvUtils.CsvRecordNoGroup rec = csvReader.readValue(line, CsvUtils.CsvRecordNoGroup.class);
+                        dataMap.computeIfAbsent(group, k -> new HashMap<>())
+                            .put(new TopicPartition(rec.getTopic(), rec.getPartition()), new OffsetAndMetadata(rec.getOffset()));
+                    }
+                } else {
+                    csvReader = CsvUtils.readerFor(CsvUtils.CsvRecordWithGroup.class);
+                    for (String line : lines) {
+                        CsvUtils.CsvRecordWithGroup rec = csvReader.readValue(line, CsvUtils.CsvRecordWithGroup.class);
+                        dataMap.computeIfAbsent(rec.getGroup(), k -> new HashMap<>())
+                            .put(new TopicPartition(rec.getTopic(), rec.getPartition()), new OffsetAndMetadata(rec.getOffset()));
+                    }
+                }
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+
+            return dataMap;
+        }
+
+        @SuppressWarnings("CyclomaticComplexity")
+        private Map<TopicPartition, OffsetAndMetadata> prepareOffsetsToReset(String groupId, Collection<TopicPartition> partitionsToReset) {
+            if (opts.options.has(opts.resetToOffsetOpt)) {
+                long offset = opts.options.valueOf(opts.resetToOffsetOpt);
+                return checkOffsetsRange(partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), tp -> offset)))
+                    .entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+            } else if (opts.options.has(opts.resetToEarliestOpt)) {
+                Map<TopicPartition, LogOffsetResult> logStartOffsets = getLogStartOffsets(partitionsToReset);
+                return partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                    LogOffsetResult logOffsetResult = logStartOffsets.get(topicPartition);
+
+                    if (!(logOffsetResult instanceof LogOffset)) {
+                        ToolsUtils.printUsageAndExit(opts.parser, "Error getting starting offset of topic partition: " + topicPartition);
+                        return null;
+                    }
+
+                    return new OffsetAndMetadata(((LogOffset) logOffsetResult).value);
+                }));
+            } else if (opts.options.has(opts.resetToLatestOpt)) {
+                Map<TopicPartition, LogOffsetResult> logEndOffsets = getLogEndOffsets(partitionsToReset);
+                return partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                    LogOffsetResult logOffsetResult = logEndOffsets.get(topicPartition);
+
+                    if (!(logOffsetResult instanceof LogOffset)) {
+                        ToolsUtils.printUsageAndExit(opts.parser, "Error getting ending offset of topic partition: " + topicPartition);
+                        return null;
+                    }
+
+                    return new OffsetAndMetadata(((LogOffset) logOffsetResult).value);
+                }));
+            } else if (opts.options.has(opts.resetShiftByOpt)) {
+                Map<TopicPartition, OffsetAndMetadata> currentCommittedOffsets = getCommittedOffsets(groupId);
+                Map<TopicPartition, Long> requestedOffsets = partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                    long shiftBy = opts.options.valueOf(opts.resetShiftByOpt);
+                    OffsetAndMetadata currentOffset = currentCommittedOffsets.get(topicPartition);
+
+                    if (currentOffset == null) {
+                        throw new IllegalArgumentException("Cannot shift offset for partition " + topicPartition + " since there is no current committed offset");
+                    }
+
+                    return currentOffset.offset() + shiftBy;
+                }));
+                return checkOffsetsRange(requestedOffsets).entrySet().stream()
+                    .collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+            } else if (opts.options.has(opts.resetToDatetimeOpt)) {
+                try {
+                    long timestamp = Utils.getDateTime(opts.options.valueOf(opts.resetToDatetimeOpt));
+                    Map<TopicPartition, LogOffsetResult> logTimestampOffsets = getLogTimestampOffsets(partitionsToReset, timestamp);
+                    return partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                        LogOffsetResult logTimestampOffset = logTimestampOffsets.get(topicPartition);
+
+                        if (!(logTimestampOffset instanceof LogOffset)) {
+                            ToolsUtils.printUsageAndExit(opts.parser, "Error getting offset by timestamp of topic partition: " + topicPartition);
+                            return null;
+                        }
+
+                        return new OffsetAndMetadata(((LogOffset) logTimestampOffset).value);
+                    }));
+                } catch (ParseException e) {
+                    throw new RuntimeException(e);
+                }
+            } else if (opts.options.has(opts.resetByDurationOpt)) {
+                String duration = opts.options.valueOf(opts.resetByDurationOpt);
+                Duration durationParsed = Duration.parse(duration);
+                Instant now = Instant.now();
+                durationParsed.negated().addTo(now);
+                long timestamp = now.minus(durationParsed).toEpochMilli();
+                Map<TopicPartition, LogOffsetResult> logTimestampOffsets = getLogTimestampOffsets(partitionsToReset, timestamp);
+                return partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                    LogOffsetResult logTimestampOffset = logTimestampOffsets.get(topicPartition);
+
+                    if (!(logTimestampOffset instanceof LogOffset)) {
+                        ToolsUtils.printUsageAndExit(opts.parser, "Error getting offset by timestamp of topic partition: " + topicPartition);
+                        return null;
+                    }
+
+                    return new OffsetAndMetadata(((LogOffset) logTimestampOffset).value);
+                }));
+            } else if (resetPlanFromFile().isPresent()) {
+                return resetPlanFromFile().map(resetPlan -> {
+                    Map<TopicPartition, OffsetAndMetadata> resetPlanForGroup = resetPlan.get(groupId);
+
+                    if (resetPlanForGroup == null) {
+                        printError("No reset plan for group " + groupId + " found", Optional.empty());
+                        return Collections.<TopicPartition, OffsetAndMetadata>emptyMap();
+                    }
+
+                    Map<TopicPartition, Long> requestedOffsets = resetPlanForGroup.keySet().stream().collect(Collectors.toMap(
+                        Function.identity(),
+                        topicPartition -> resetPlanForGroup.get(topicPartition).offset()));
+
+                    return checkOffsetsRange(requestedOffsets).entrySet().stream()
+                        .collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+                }).orElseGet(Collections::emptyMap);
+            } else if (opts.options.has(opts.resetToCurrentOpt)) {
+                Map<TopicPartition, OffsetAndMetadata> currentCommittedOffsets = getCommittedOffsets(groupId);
+                Collection<TopicPartition> partitionsToResetWithCommittedOffset = new ArrayList<>();
+                Collection<TopicPartition> partitionsToResetWithoutCommittedOffset = new ArrayList<>();
+
+                for (TopicPartition topicPartition : partitionsToReset) {
+                    if (currentCommittedOffsets.containsKey(topicPartition))
+                        partitionsToResetWithCommittedOffset.add(topicPartition);
+                    else
+                        partitionsToResetWithoutCommittedOffset.add(topicPartition);
+                }
+
+                Map<TopicPartition, OffsetAndMetadata> preparedOffsetsForPartitionsWithCommittedOffset = partitionsToResetWithCommittedOffset.stream()
+                    .collect(Collectors.toMap(Function.identity(), topicPartition -> {
+                        OffsetAndMetadata committedOffset = currentCommittedOffsets.get(topicPartition);
+
+                        if (committedOffset == null) {
+                            throw new IllegalStateException("Expected a valid current offset for topic partition: " + topicPartition);
+                        }
+
+                        return new OffsetAndMetadata(committedOffset.offset());
+                    }));
+
+                Map<TopicPartition, OffsetAndMetadata> preparedOffsetsForPartitionsWithoutCommittedOffset = getLogEndOffsets(partitionsToResetWithoutCommittedOffset)
+                    .entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> {
+                        if (!(e.getValue() instanceof LogOffset)) {
+                            ToolsUtils.printUsageAndExit(opts.parser, "Error getting ending offset of topic partition: " + e.getKey());
+                            return null;
+                        }
+
+                        return new OffsetAndMetadata(((LogOffset) e.getValue()).value);
+                    }));
+
+                preparedOffsetsForPartitionsWithCommittedOffset.putAll(preparedOffsetsForPartitionsWithoutCommittedOffset);
+
+                return preparedOffsetsForPartitionsWithCommittedOffset;
+            }
+
+            ToolsUtils.printUsageAndExit(opts.parser, String.format("Option '%s' requires one of the following scenarios: %s", opts.resetOffsetsOpt, opts.allResetOffsetScenarioOpts));
+            return null;
+        }
+
+        private Map<TopicPartition, Long> checkOffsetsRange(Map<TopicPartition, Long> requestedOffsets) {
+            Map<TopicPartition, LogOffsetResult> logStartOffsets = getLogStartOffsets(requestedOffsets.keySet());
+            Map<TopicPartition, LogOffsetResult> logEndOffsets = getLogEndOffsets(requestedOffsets.keySet());
+
+            Map<TopicPartition, Long> res = new HashMap<>();
+
+            requestedOffsets.forEach((topicPartition, offset) -> {
+                LogOffsetResult logEndOffset = logEndOffsets.get(topicPartition);
+
+                if (logEndOffset != null) {
+                    if (logEndOffset instanceof LogOffset && offset > ((LogOffset) logEndOffset).value) {
+                        long endOffset = ((LogOffset) logEndOffset).value;
+                        LOGGER.warn("New offset (" + offset + ") is higher than latest offset for topic partition " + topicPartition + ". Value will be set to " + endOffset);
+                        res.put(topicPartition, endOffset);
+                    } else {
+                        LogOffsetResult logStartOffset = logStartOffsets.get(topicPartition);
+
+                        if (logStartOffset instanceof LogOffset && offset < ((LogOffset) logStartOffset).value) {
+                            long startOffset = ((LogOffset) logStartOffset).value;
+                            LOGGER.warn("New offset (" + offset + ") is lower than earliest offset for topic partition " + topicPartition + ". Value will be set to " + startOffset);
+                            res.put(topicPartition, startOffset);
+                        } else
+                            res.put(topicPartition, offset);
+                    }
+                } else {
+                    // the control should not reach here
+                    throw new IllegalStateException("Unexpected non-existing offset value for topic partition " + topicPartition);
+                }
+            });
+
+            return res;
+        }
+
+        String exportOffsetsToCsv(Map<String, Map<TopicPartition, OffsetAndMetadata>> assignments) {
+            boolean isSingleGroupQuery = opts.options.valuesOf(opts.groupOpt).size() == 1;
+            ObjectWriter csvWriter = isSingleGroupQuery
+                ? CsvUtils.writerFor(CsvUtils.CsvRecordNoGroup.class)
+                : CsvUtils.writerFor(CsvUtils.CsvRecordWithGroup.class);
+
+            return Utils.mkString(assignments.entrySet().stream().flatMap(e -> {
+                String groupId = e.getKey();
+                Map<TopicPartition, OffsetAndMetadata> partitionInfo = e.getValue();
+
+                return partitionInfo.entrySet().stream().map(e1 -> {
+                    TopicPartition k = e1.getKey();
+                    OffsetAndMetadata v = e1.getValue();
+                    Object csvRecord = isSingleGroupQuery
+                        ? new CsvUtils.CsvRecordNoGroup(k.topic(), k.partition(), v.offset())
+                        : new CsvUtils.CsvRecordWithGroup(groupId, k.topic(), k.partition(), v.offset());
+
+                    try {
+                        return csvWriter.writeValueAsString(csvRecord);
+                    } catch (JsonProcessingException err) {
+                        throw new RuntimeException(err);
+                    }
+                });
+            }), "", "", "");
+        }
+
+        Map<String, Throwable> deleteGroups() {
+            List<String> groupIds = opts.options.has(opts.allGroupsOpt)
+                ? listConsumerGroups()
+                : opts.options.valuesOf(opts.groupOpt);
+
+            Map<String, KafkaFuture<Void>> groupsToDelete = adminClient.deleteConsumerGroups(
+                groupIds,
+                withTimeoutMs(new DeleteConsumerGroupsOptions())
+            ).deletedGroups();
+
+            Map<String, Throwable> success = new HashMap<>();
+            Map<String, Throwable> failed = new HashMap<>();
+
+            groupsToDelete.forEach((g, f) -> {
+                try {
+                    f.get();
+                    success.put(g, null);
+                } catch (ExecutionException | InterruptedException e) {
+                    failed.put(g, e);
+                }
+            });
+
+            if (failed.isEmpty())
+                System.out.println("Deletion of requested consumer groups (" + Utils.mkString(success.keySet().stream(), "'", "'", "', '") + ") was successful.");
+            else {
+                printError("Deletion of some consumer groups failed:", Optional.empty());
+                failed.forEach((group, error) -> System.out.println("* Group '" + group + "' could not be deleted due to: " + error));
+
+                if (!success.isEmpty())
+                    System.out.println("\nThese consumer groups were deleted successfully: " + Utils.mkString(success.keySet().stream(), "'", "', '", "'"));
+            }
+
+            failed.putAll(success);
+
+            return failed;
+        }
+    }
+
+    interface LogOffsetResult { }
+
+    private static class LogOffset implements LogOffsetResult {
+        final long value;
+
+        LogOffset(long value) {
+            this.value = value;
+        }
+    }
+
+    private static class Unknown implements LogOffsetResult { }
+
+    private static class Ignore implements LogOffsetResult { }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
index 045d296444dfb..3eb811bcaf829 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
@@ -31,89 +31,99 @@
 import static org.apache.kafka.tools.ToolsUtils.minus;
 
 public class ConsumerGroupCommandOptions extends CommandDefaultOptions {
-    public static final Logger LOGGER = LoggerFactory.getLogger(ConsumerGroupCommandOptions.class);
+    private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerGroupCommandOptions.class);
 
-    public static final String BOOTSTRAP_SERVER_DOC = "REQUIRED: The server(s) to connect to.";
-    public static final String GROUP_DOC = "The consumer group we wish to act on.";
-    public static final String TOPIC_DOC = "The topic whose consumer group information should be deleted or topic whose should be included in the reset offset process. " +
+    private static final String BOOTSTRAP_SERVER_DOC = "REQUIRED: The server(s) to connect to.";
+    private static final String GROUP_DOC = "The consumer group we wish to act on.";
+    private static final String TOPIC_DOC = "The topic whose consumer group information should be deleted or topic whose should be included in the reset offset process. " +
         "In `reset-offsets` case, partitions can be specified using this format: `topic1:0,1,2`, where 0,1,2 are the partition to be included in the process. " +
         "Reset-offsets also supports multiple topic inputs.";
-    public static final String ALL_TOPICS_DOC = "Consider all topics assigned to a group in the `reset-offsets` process.";
-    public static final String LIST_DOC = "List all consumer groups.";
-    public static final String DESCRIBE_DOC = "Describe consumer group and list offset lag (number of messages not yet processed) related to given group.";
-    public static final String ALL_GROUPS_DOC = "Apply to all consumer groups.";
-    public static final String NL = System.lineSeparator();
-    public static final String DELETE_DOC = "Pass in groups to delete topic partition offsets and ownership information " +
+    private static final String ALL_TOPICS_DOC = "Consider all topics assigned to a group in the `reset-offsets` process.";
+    private static final String LIST_DOC = "List all consumer groups.";
+    private static final String DESCRIBE_DOC = "Describe consumer group and list offset lag (number of messages not yet processed) related to given group.";
+    private static final String ALL_GROUPS_DOC = "Apply to all consumer groups.";
+    private static final String NL = System.lineSeparator();
+    private static final String DELETE_DOC = "Pass in groups to delete topic partition offsets and ownership information " +
         "over the entire consumer group. For instance --group g1 --group g2";
-    public static final String TIMEOUT_MS_DOC = "The timeout that can be set for some use cases. For example, it can be used when describing the group " +
+    private static final String TIMEOUT_MS_DOC = "The timeout that can be set for some use cases. For example, it can be used when describing the group " +
         "to specify the maximum amount of time in milliseconds to wait before the group stabilizes (when the group is just created, " +
         "or is going through some changes).";
-    public static final String COMMAND_CONFIG_DOC = "Property file containing configs to be passed to Admin Client and Consumer.";
-    public static final String RESET_OFFSETS_DOC = "Reset offsets of consumer group. Supports one consumer group at the time, and instances should be inactive" + NL +
+    private static final String COMMAND_CONFIG_DOC = "Property file containing configs to be passed to Admin Client and Consumer.";
+    private static final String RESET_OFFSETS_DOC = "Reset offsets of consumer group. Supports one consumer group at the time, and instances should be inactive" + NL +
         "Has 2 execution options: --dry-run (the default) to plan which offsets to reset, and --execute to update the offsets. " +
         "Additionally, the --export option is used to export the results to a CSV format." + NL +
         "You must choose one of the following reset specifications: --to-datetime, --by-duration, --to-earliest, " +
         "--to-latest, --shift-by, --from-file, --to-current, --to-offset." + NL +
         "To define the scope use --all-topics or --topic. One scope must be specified unless you use '--from-file'.";
-    public static final String DRY_RUN_DOC = "Only show results without executing changes on Consumer Groups. Supported operations: reset-offsets.";
-    public static final String EXECUTE_DOC = "Execute operation. Supported operations: reset-offsets.";
-    public static final String EXPORT_DOC = "Export operation execution to a CSV file. Supported operations: reset-offsets.";
-    public static final String RESET_TO_OFFSET_DOC = "Reset offsets to a specific offset.";
-    public static final String RESET_FROM_FILE_DOC = "Reset offsets to values defined in CSV file.";
-    public static final String RESET_TO_DATETIME_DOC = "Reset offsets to offset from datetime. Format: 'YYYY-MM-DDTHH:mm:SS.sss'";
-    public static final String RESET_BY_DURATION_DOC = "Reset offsets to offset by duration from current timestamp. Format: 'PnDTnHnMnS'";
-    public static final String RESET_TO_EARLIEST_DOC = "Reset offsets to earliest offset.";
-    public static final String RESET_TO_LATEST_DOC = "Reset offsets to latest offset.";
-    public static final String RESET_TO_CURRENT_DOC = "Reset offsets to current offset.";
-    public static final String RESET_SHIFT_BY_DOC = "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative.";
-    public static final String MEMBERS_DOC = "Describe members of the group. This option may be used with '--describe' and '--bootstrap-server' options only." + NL +
+    private static final String DRY_RUN_DOC = "Only show results without executing changes on Consumer Groups. Supported operations: reset-offsets.";
+    private static final String EXECUTE_DOC = "Execute operation. Supported operations: reset-offsets.";
+    private static final String EXPORT_DOC = "Export operation execution to a CSV file. Supported operations: reset-offsets.";
+    private static final String RESET_TO_OFFSET_DOC = "Reset offsets to a specific offset.";
+    private static final String RESET_FROM_FILE_DOC = "Reset offsets to values defined in CSV file.";
+    private static final String RESET_TO_DATETIME_DOC = "Reset offsets to offset from datetime. Format: 'YYYY-MM-DDTHH:mm:SS.sss'";
+    private static final String RESET_BY_DURATION_DOC = "Reset offsets to offset by duration from current timestamp. Format: 'PnDTnHnMnS'";
+    private static final String RESET_TO_EARLIEST_DOC = "Reset offsets to earliest offset.";
+    private static final String RESET_TO_LATEST_DOC = "Reset offsets to latest offset.";
+    private static final String RESET_TO_CURRENT_DOC = "Reset offsets to current offset.";
+    private static final String RESET_SHIFT_BY_DOC = "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative.";
+    private static final String MEMBERS_DOC = "Describe members of the group. This option may be used with '--describe' and '--bootstrap-server' options only." + NL +
         "Example: --bootstrap-server localhost:9092 --describe --group group1 --members";
-    public static final String VERBOSE_DOC = "Provide additional information, if any, when describing the group. This option may be used " +
+    private static final String VERBOSE_DOC = "Provide additional information, if any, when describing the group. This option may be used " +
         "with '--offsets'/'--members'/'--state' and '--bootstrap-server' options only." + NL + "Example: --bootstrap-server localhost:9092 --describe --group group1 --members --verbose";
-    public static final String OFFSETS_DOC = "Describe the group and list all topic partitions in the group along with their offset lag. " +
+    private static final String OFFSETS_DOC = "Describe the group and list all topic partitions in the group along with their offset lag. " +
         "This is the default sub-action of and may be used with '--describe' and '--bootstrap-server' options only." + NL +
         "Example: --bootstrap-server localhost:9092 --describe --group group1 --offsets";
-    public static final String STATE_DOC = "When specified with '--describe', includes the state of the group." + NL +
+    private static final String STATE_DOC = "When specified with '--describe', includes the state of the group." + NL +
         "Example: --bootstrap-server localhost:9092 --describe --group group1 --state" + NL +
         "When specified with '--list', it displays the state of all groups. It can also be used to list groups with specific states." + NL +
         "Example: --bootstrap-server localhost:9092 --list --state stable,empty" + NL +
         "This option may be used with '--describe', '--list' and '--bootstrap-server' options only.";
-    public static final String DELETE_OFFSETS_DOC = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics.";
+    private static final String TYPE_DOC = "When specified with '--list', it displays the types of all the groups. It can also be used to list groups with specific types." + NL +
+        "Example: --bootstrap-server localhost:9092 --list --type classic,consumer" + NL +
+        "This option may be used with the '--list' option only.";
+    private static final String DELETE_OFFSETS_DOC = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics.";
 
-    public final OptionSpec<String> bootstrapServerOpt;
-    public final OptionSpec<String> groupOpt;
-    public final OptionSpec<String> topicOpt;
-    public final OptionSpec<Void> allTopicsOpt;
-    public final OptionSpec<Void> listOpt;
-    public final OptionSpec<Void> describeOpt;
-    public final OptionSpec<Void> allGroupsOpt;
-    public final OptionSpec<Void> deleteOpt;
-    public final OptionSpec<Long> timeoutMsOpt;
-    public final OptionSpec<String> commandConfigOpt;
-    public final OptionSpec<Void> resetOffsetsOpt;
-    public final OptionSpec<Void> deleteOffsetsOpt;
-    public final OptionSpec<Void> dryRunOpt;
-    public final OptionSpec<Void> executeOpt;
-    public final OptionSpec<Void> exportOpt;
-    public final OptionSpec<Long> resetToOffsetOpt;
-    public final OptionSpec<String> resetFromFileOpt;
-    public final OptionSpec<String> resetToDatetimeOpt;
-    public final OptionSpec<String> resetByDurationOpt;
-    public final OptionSpec<Void> resetToEarliestOpt;
-    public final OptionSpec<Void> resetToLatestOpt;
-    public final OptionSpec<Void> resetToCurrentOpt;
-    public final OptionSpec<Long> resetShiftByOpt;
-    public final OptionSpec<Void> membersOpt;
-    public final OptionSpec<Void> verboseOpt;
-    public final OptionSpec<Void> offsetsOpt;
-    public final OptionSpec<String> stateOpt;
+    final OptionSpec<String> bootstrapServerOpt;
+    final OptionSpec<String> groupOpt;
+    final OptionSpec<String> topicOpt;
+    final OptionSpec<Void> allTopicsOpt;
+    final OptionSpec<Void> listOpt;
+    final OptionSpec<Void> describeOpt;
+    final OptionSpec<Void> allGroupsOpt;
+    final OptionSpec<Void> deleteOpt;
+    final OptionSpec<Long> timeoutMsOpt;
+    final OptionSpec<String> commandConfigOpt;
+    final OptionSpec<Void> resetOffsetsOpt;
+    final OptionSpec<Void> deleteOffsetsOpt;
+    final OptionSpec<Void> dryRunOpt;
+    final OptionSpec<Void> executeOpt;
+    final OptionSpec<Void> exportOpt;
+    final OptionSpec<Long> resetToOffsetOpt;
+    final OptionSpec<String> resetFromFileOpt;
+    final OptionSpec<String> resetToDatetimeOpt;
+    final OptionSpec<String> resetByDurationOpt;
+    final OptionSpec<Void> resetToEarliestOpt;
+    final OptionSpec<Void> resetToLatestOpt;
+    final OptionSpec<Void> resetToCurrentOpt;
+    final OptionSpec<Long> resetShiftByOpt;
+    final OptionSpec<Void> membersOpt;
+    final OptionSpec<Void> verboseOpt;
+    final OptionSpec<Void> offsetsOpt;
+    final OptionSpec<String> stateOpt;
+    final OptionSpec<String> typeOpt;
 
-    public final Set<OptionSpec<?>> allGroupSelectionScopeOpts;
-    public final Set<OptionSpec<?>> allConsumerGroupLevelOpts;
-    public final Set<OptionSpec<?>> allResetOffsetScenarioOpts;
-    public final Set<OptionSpec<?>> allDeleteOffsetsOpts;
+    final Set<OptionSpec<?>> allGroupSelectionScopeOpts;
+    final Set<OptionSpec<?>> allConsumerGroupLevelOpts;
+    final Set<OptionSpec<?>> allResetOffsetScenarioOpts;
+    final Set<OptionSpec<?>> allDeleteOffsetsOpts;
 
-    public ConsumerGroupCommandOptions(String[] args) {
+    public static ConsumerGroupCommandOptions fromArgs(String[] args) {
+        ConsumerGroupCommandOptions opts = new ConsumerGroupCommandOptions(args);
+        opts.checkArgs();
+        return opts;
+    }
+
+    private ConsumerGroupCommandOptions(String[] args) {
         super(args);
 
         bootstrapServerOpt = parser.accepts("bootstrap-server", BOOTSTRAP_SERVER_DOC)
@@ -180,6 +190,10 @@ public ConsumerGroupCommandOptions(String[] args) {
             .availableIf(describeOpt, listOpt)
             .withOptionalArg()
             .ofType(String.class);
+        typeOpt = parser.accepts("type", TYPE_DOC)
+            .availableIf(listOpt)
+            .withOptionalArg()
+            .ofType(String.class);
 
         allGroupSelectionScopeOpts = new HashSet<>(Arrays.asList(groupOpt, allGroupsOpt));
         allConsumerGroupLevelOpts = new HashSet<>(Arrays.asList(listOpt, describeOpt, deleteOpt, resetOffsetsOpt));
@@ -191,7 +205,7 @@ public ConsumerGroupCommandOptions(String[] args) {
     }
 
     @SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"})
-    public void checkArgs() {
+    void checkArgs() {
         CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt);
 
         if (options.has(describeOpt)) {
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/CsvUtils.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/CsvUtils.java
new file mode 100644
index 0000000000000..60a8bb22be950
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/CsvUtils.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.ObjectReader;
+import com.fasterxml.jackson.databind.ObjectWriter;
+import com.fasterxml.jackson.dataformat.csv.CsvMapper;
+import com.fasterxml.jackson.dataformat.csv.CsvSchema;
+
+public class CsvUtils {
+    private final static CsvMapper MAPPER = new CsvMapper();
+
+    static ObjectReader readerFor(Class<?> clazz) {
+        return MAPPER.readerFor(clazz).with(getSchema(clazz));
+    }
+
+    static ObjectWriter writerFor(Class<?> clazz) {
+        return MAPPER.writerFor(clazz).with(getSchema(clazz));
+    }
+
+    private static CsvSchema getSchema(Class<?> clazz) {
+        String[] fields;
+        if (CsvRecordWithGroup.class == clazz)
+            fields = CsvRecordWithGroup.FIELDS;
+        else if (CsvRecordNoGroup.class == clazz)
+            fields = CsvRecordNoGroup.FIELDS;
+        else
+            throw new IllegalStateException("Unhandled class " + clazz);
+
+        return MAPPER.schemaFor(clazz).sortedBy(fields);
+    }
+
+    public static class CsvRecordWithGroup {
+        public static final String[] FIELDS = new String[] {"group", "topic", "partition", "offset"};
+
+        @JsonProperty
+        private String group;
+
+        @JsonProperty
+        private String topic;
+
+        @JsonProperty
+        private int partition;
+
+        @JsonProperty
+        private long offset;
+
+        /**
+         * Required for jackson.
+         */
+        public CsvRecordWithGroup() {
+        }
+
+        public CsvRecordWithGroup(String group, String topic, int partition, long offset) {
+            this.group = group;
+            this.topic = topic;
+            this.partition = partition;
+            this.offset = offset;
+        }
+
+        public void setGroup(String group) {
+            this.group = group;
+        }
+
+        public String getGroup() {
+            return group;
+        }
+
+        public String getTopic() {
+            return topic;
+        }
+
+        public void setTopic(String topic) {
+            this.topic = topic;
+        }
+
+        public int getPartition() {
+            return partition;
+        }
+
+        public void setPartition(int partition) {
+            this.partition = partition;
+        }
+
+        public long getOffset() {
+            return offset;
+        }
+
+        public void setOffset(long offset) {
+            this.offset = offset;
+        }
+    }
+
+    public static class CsvRecordNoGroup {
+        public static final String[] FIELDS = new String[]{"topic", "partition", "offset"};
+
+        @JsonProperty
+        private String topic;
+
+        @JsonProperty
+        private int partition;
+
+        @JsonProperty
+        private long offset;
+
+        /**
+         * Required for jackson.
+         */
+        public CsvRecordNoGroup() {
+        }
+
+        public CsvRecordNoGroup(String topic, int partition, long offset) {
+            this.topic = topic;
+            this.partition = partition;
+            this.offset = offset;
+        }
+
+        public String getTopic() {
+            return topic;
+        }
+
+        public void setTopic(String topic) {
+            this.topic = topic;
+        }
+
+        public int getPartition() {
+            return partition;
+        }
+
+        public void setPartition(int partition) {
+            this.partition = partition;
+        }
+
+        public long getOffset() {
+            return offset;
+        }
+
+        public void setOffset(long offset) {
+            this.offset = offset;
+        }
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/GroupState.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/GroupState.java
new file mode 100644
index 0000000000000..04a3b2eae195c
--- /dev/null
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/GroupState.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.tools.consumer.group;
+
+import org.apache.kafka.common.Node;
+
+class GroupState {
+    final String group;
+    final Node coordinator;
+    final String assignmentStrategy;
+    final String state;
+    final int numMembers;
+
+    GroupState(String group, Node coordinator, String assignmentStrategy, String state, int numMembers) {
+        this.group = group;
+        this.coordinator = coordinator;
+        this.assignmentStrategy = assignmentStrategy;
+        this.state = state;
+        this.numMembers = numMembers;
+    }
+}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/MemberAssignmentState.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/MemberAssignmentState.java
index 040cb1c741ec0..9ac46428338e4 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/MemberAssignmentState.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/MemberAssignmentState.java
@@ -21,15 +21,15 @@
 import java.util.List;
 
 class MemberAssignmentState {
-    public final String group;
-    public final String consumerId;
-    public final String host;
-    public final String clientId;
-    public final String groupInstanceId;
-    public final int numPartitions;
-    public final List<TopicPartition> assignment;
+    final String group;
+    final String consumerId;
+    final String host;
+    final String clientId;
+    final String groupInstanceId;
+    final int numPartitions;
+    final List<TopicPartition> assignment;
 
-    public MemberAssignmentState(String group, String consumerId, String host, String clientId, String groupInstanceId,
+    MemberAssignmentState(String group, String consumerId, String host, String clientId, String groupInstanceId,
                                  int numPartitions, List<TopicPartition> assignment) {
         this.group = group;
         this.consumerId = consumerId;
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/PartitionAssignmentState.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/PartitionAssignmentState.java
index 396032f0a0c29..9e45f2b05e772 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/PartitionAssignmentState.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/PartitionAssignmentState.java
@@ -19,25 +19,23 @@
 import org.apache.kafka.common.Node;
 
 import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
 
 class PartitionAssignmentState {
-    public final String group;
-    public final Optional<Node> coordinator;
-    public final Optional<String> topic;
-    public final OptionalInt partition;
-    public final OptionalLong offset;
-    public final OptionalLong lag;
-    public final Optional<String> consumerId;
-    public final Optional<String> host;
-    public final Optional<String> clientId;
-    public final OptionalLong logEndOffset;
+    final String group;
+    final Optional<Node> coordinator;
+    final Optional<String> topic;
+    final Optional<Integer> partition;
+    final Optional<Long> offset;
+    final Optional<Long> lag;
+    final Optional<String> consumerId;
+    final Optional<String> host;
+    final Optional<String> clientId;
+    final Optional<Long> logEndOffset;
 
-    public PartitionAssignmentState(String group, Optional<Node> coordinator, Optional<String> topic,
-                                    OptionalInt partition, OptionalLong offset, OptionalLong lag,
+    PartitionAssignmentState(String group, Optional<Node> coordinator, Optional<String> topic,
+                                    Optional<Integer> partition, Optional<Long> offset, Optional<Long> lag,
                                     Optional<String> consumerId, Optional<String> host, Optional<String> clientId,
-                                    OptionalLong logEndOffset) {
+                                    Optional<Long> logEndOffset) {
         this.group = group;
         this.coordinator = coordinator;
         this.topic = topic;
diff --git a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
index 83fa31bf5e795..0bb2c9c9d1b71 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
@@ -16,8 +16,8 @@
  */
 package org.apache.kafka.tools;
 
-import kafka.utils.TestInfoUtils;
 import kafka.server.DynamicConfig;
+import kafka.utils.TestInfoUtils;
 import kafka.utils.TestUtils;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AlterConfigOp;
@@ -210,6 +210,66 @@ public static File tempPropertiesFile(Map<String, String> properties) throws IOE
         return org.apache.kafka.test.TestUtils.tempFile(sb.toString());
     }
 
+    /**
+     * Capture the console output during the execution of the provided function.
+     */
+    public static String grabConsoleOutput(Runnable f) {
+        ByteArrayOutputStream buf = new ByteArrayOutputStream();
+        PrintStream out = new PrintStream(buf);
+        PrintStream out0 = System.out;
+
+        System.setOut(out);
+        try {
+            f.run();
+        } finally {
+            System.setOut(out0);
+        }
+        out.flush();
+        return buf.toString();
+    }
+
+    /**
+     * Capture the console error during the execution of the provided function.
+     */
+    public static String grabConsoleError(Runnable f) {
+        ByteArrayOutputStream buf = new ByteArrayOutputStream();
+        PrintStream err = new PrintStream(buf);
+        PrintStream err0 = System.err;
+
+        System.setErr(err);
+        try {
+            f.run();
+        } finally {
+            System.setErr(err0);
+        }
+        err.flush();
+        return buf.toString();
+    }
+
+    /**
+     * Capture both the console output and console error during the execution of the provided function.
+     */
+    public static Tuple2<String, String> grabConsoleOutputAndError(Runnable f) {
+        ByteArrayOutputStream outBuf = new ByteArrayOutputStream();
+        ByteArrayOutputStream errBuf = new ByteArrayOutputStream();
+        PrintStream out = new PrintStream(outBuf);
+        PrintStream err = new PrintStream(errBuf);
+        PrintStream out0 = System.out;
+        PrintStream err0 = System.err;
+
+        System.setOut(out);
+        System.setErr(err);
+        try {
+            f.run();
+        } finally {
+            System.setOut(out0);
+            System.setErr(err0);
+        }
+        out.flush();
+        err.flush();
+        return new Tuple2<>(outBuf.toString(), errBuf.toString());
+    }
+
     public static class MockExitProcedure implements Exit.Procedure {
         private boolean hasExited = false;
         private int statusCode;
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
index 3094572b04921..a0ad01bad955d 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
@@ -16,45 +16,30 @@
  */
 package org.apache.kafka.tools.consumer.group;
 
-import kafka.admin.ConsumerGroupCommand;
 import kafka.api.AbstractAuthorizerIntegrationTest;
 import kafka.security.authorizer.AclEntry;
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
-import scala.collection.immutable.Map$;
+import scala.collection.JavaConverters;
 
 import java.util.Collections;
-import java.util.Properties;
 
 import static org.apache.kafka.common.acl.AclOperation.DESCRIBE;
 import static org.apache.kafka.common.acl.AclPermissionType.ALLOW;
 import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_NAME;
-import static org.apache.kafka.tools.consumer.group.ConsumerGroupCommandTest.set;
 
 public class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
+    @SuppressWarnings({"deprecation"})
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
-    public void testDescribeGroupCliWithGroupDescribe(String quorum) {
-        addAndVerifyAcls(set(Collections.singleton(new AccessControlEntry(ClientPrincipal().toString(), AclEntry.WildcardHost(), DESCRIBE, ALLOW))), groupResource());
+    public void testDescribeGroupCliWithGroupDescribe(String quorum) throws Exception {
+        addAndVerifyAcls(JavaConverters.asScalaSet(Collections.singleton(new AccessControlEntry(ClientPrincipal().toString(), AclEntry.WildcardHost(), DESCRIBE, ALLOW))).toSet(), groupResource());
 
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group()};
-        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(cgcArgs);
-        ConsumerGroupCommand.ConsumerGroupService consumerGroupService = new ConsumerGroupCommand.ConsumerGroupService(opts, Map$.MODULE$.empty());
+        ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(cgcArgs);
+        ConsumerGroupCommand.ConsumerGroupService consumerGroupService = new ConsumerGroupCommand.ConsumerGroupService(opts, Collections.emptyMap());
         consumerGroupService.describeGroups();
         consumerGroupService.close();
     }
-
-    private void createTopicWithBrokerPrincipal(String topic) {
-        // Note the principal builder implementation maps all connections on the
-        // inter-broker listener to the broker principal.
-        createTopic(
-            topic,
-            1,
-            1,
-            new Properties(),
-            interBrokerListenerName(),
-            new Properties()
-        );
-    }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
index e3ee39f6b1d04..e1daa37f1976b 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.tools.consumer.group;
 
 import kafka.api.BaseConsumerTest;
-import kafka.admin.ConsumerGroupCommand;
 import kafka.server.KafkaConfig;
 import kafka.utils.TestUtils;
 import org.apache.kafka.clients.admin.AdminClientConfig;
@@ -41,7 +40,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -130,10 +128,10 @@ Consumer<String, String> createNoAutoCommitConsumer(String group) {
     }
 
     ConsumerGroupCommand.ConsumerGroupService getConsumerGroupService(String[] args) {
-        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(args);
+        ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(args);
         ConsumerGroupCommand.ConsumerGroupService service = new ConsumerGroupCommand.ConsumerGroupService(
             opts,
-            asScala(Collections.singletonMap(AdminClientConfig.RETRIES_CONFIG, Integer.toString(Integer.MAX_VALUE)))
+            Collections.singletonMap(AdminClientConfig.RETRIES_CONFIG, Integer.toString(Integer.MAX_VALUE))
         );
 
         consumerGroupService.add(0, service);
@@ -352,14 +350,4 @@ public static Stream<Arguments> getTestQuorumAndGroupProtocolParametersConsumerG
     static <T> Seq<T> seq(Collection<T> seq) {
         return JavaConverters.asScalaIteratorConverter(seq.iterator()).asScala().toSeq();
     }
-
-    @SuppressWarnings("deprecation")
-    static <K, V> scala.collection.Map<K, V> asScala(Map<K, V> jmap) {
-        return JavaConverters.mapAsScalaMap(jmap);
-    }
-
-    @SuppressWarnings({"deprecation"})
-    static <T> scala.collection.immutable.Set<T> set(final Collection<T> set) {
-        return JavaConverters.asScalaSet(new HashSet<>(set)).toSet();
-    }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
index 8c75823648f83..b7583c1c44f55 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.tools.consumer.group;
 
-import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AdminClientTestUtils;
 import org.apache.kafka.clients.admin.ConsumerGroupDescription;
@@ -39,15 +38,10 @@
 import org.apache.kafka.common.TopicPartitionInfo;
 import org.apache.kafka.common.internals.KafkaFutureImpl;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentMatcher;
 import org.mockito.ArgumentMatchers;
-import scala.Option;
-import scala.Some;
-import scala.Tuple2;
-import scala.collection.JavaConverters;
-import scala.collection.Seq;
-import scala.collection.immutable.Map$;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -86,7 +80,7 @@ public class ConsumerGroupServiceTest {
     private final Admin admin = mock(Admin.class);
 
     @Test
-    public void testAdminRequestsForDescribeOffsets() {
+    public void testAdminRequestsForDescribeOffsets() throws Exception {
         String[] args = new String[]{"--bootstrap-server", "localhost:9092", "--group", GROUP, "--describe", "--offsets"};
         ConsumerGroupCommand.ConsumerGroupService groupService = consumerGroupService(args);
 
@@ -97,10 +91,10 @@ public void testAdminRequestsForDescribeOffsets() {
         when(admin.listOffsets(offsetsArgMatcher(), any()))
                 .thenReturn(listOffsetsResult());
 
-        Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
-        assertEquals(Some.apply("Stable"), statesAndAssignments._1);
-        assertTrue(statesAndAssignments._2.isDefined());
-        assertEquals(TOPIC_PARTITIONS.size(), statesAndAssignments._2.get().size());
+        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
+        assertEquals(Optional.of("Stable"), statesAndAssignments.v1);
+        assertTrue(statesAndAssignments.v2.isPresent());
+        assertEquals(TOPIC_PARTITIONS.size(), statesAndAssignments.v2.get().size());
 
         verify(admin, times(1)).describeConsumerGroups(ArgumentMatchers.eq(Collections.singletonList(GROUP)), any());
         verify(admin, times(1)).listConsumerGroupOffsets(ArgumentMatchers.eq(listConsumerGroupOffsetsSpec()), any());
@@ -108,7 +102,7 @@ public void testAdminRequestsForDescribeOffsets() {
     }
 
     @Test
-    public void testAdminRequestsForDescribeNegativeOffsets() {
+    public void testAdminRequestsForDescribeNegativeOffsets() throws Exception {
         String[] args = new String[]{"--bootstrap-server", "localhost:9092", "--group", GROUP, "--describe", "--offsets"};
         ConsumerGroupCommand.ConsumerGroupService groupService = consumerGroupService(args);
 
@@ -170,20 +164,15 @@ public void testAdminRequestsForDescribeNegativeOffsets() {
         )).thenReturn(new ListOffsetsResult(endOffsets.entrySet().stream().filter(e -> unassignedTopicPartitions.contains(e.getKey()))
                 .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))));
 
-        Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
-        Option<String> state = statesAndAssignments._1;
-        Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>> assignments = statesAndAssignments._2;
-
-        Map<TopicPartition, Optional<Long>> returnedOffsets = new HashMap<>();
-        assignments.foreach(results -> {
-            results.foreach(assignment -> {
-                returnedOffsets.put(
-                        new TopicPartition(assignment.topic().get(), (Integer) assignment.partition().get()),
-                        assignment.offset().isDefined() ? Optional.of((Long) assignment.offset().get()) : Optional.empty());
-                return null;
-            });
-            return null;
-        });
+        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
+        Optional<String> state = statesAndAssignments.v1;
+        Optional<Collection<PartitionAssignmentState>> assignments = statesAndAssignments.v2;
+
+        Map<TopicPartition, Optional<Long>> returnedOffsets = assignments.map(results ->
+            results.stream().collect(Collectors.toMap(
+                assignment -> new TopicPartition(assignment.topic.get(), assignment.partition.get()),
+                assignment -> assignment.offset))
+        ).orElse(Collections.emptyMap());
 
         Map<TopicPartition, Optional<Long>> expectedOffsets = new HashMap<>();
 
@@ -194,7 +183,7 @@ public void testAdminRequestsForDescribeNegativeOffsets() {
         expectedOffsets.put(testTopicPartition4, Optional.of(100L));
         expectedOffsets.put(testTopicPartition5, Optional.empty());
 
-        assertEquals(Some.apply("Stable"), state);
+        assertEquals(Optional.of("Stable"), state);
         assertEquals(expectedOffsets, returnedOffsets);
 
         verify(admin, times(1)).describeConsumerGroups(ArgumentMatchers.eq(Collections.singletonList(GROUP)), any());
@@ -220,9 +209,9 @@ public void testAdminRequestsForResetOffsets() {
         when(admin.listOffsets(offsetsArgMatcher(), any()))
                 .thenReturn(listOffsetsResult());
 
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> resetResult = groupService.resetOffsets();
-        assertEquals(set(Collections.singletonList(GROUP)), resetResult.keySet());
-        assertEquals(set(TOPIC_PARTITIONS), resetResult.get(GROUP).get().keys().toSet());
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> resetResult = groupService.resetOffsets();
+        assertEquals(Collections.singleton(GROUP), resetResult.keySet());
+        assertEquals(new HashSet<>(TOPIC_PARTITIONS), resetResult.get(GROUP).keySet());
 
         verify(admin, times(1)).describeConsumerGroups(ArgumentMatchers.eq(Collections.singletonList(GROUP)), any());
         verify(admin, times(1)).describeTopics(ArgumentMatchers.eq(topicsWithoutPartitionsSpecified), any());
@@ -230,9 +219,9 @@ public void testAdminRequestsForResetOffsets() {
     }
 
     private ConsumerGroupCommand.ConsumerGroupService consumerGroupService(String[] args) {
-        return new ConsumerGroupCommand.ConsumerGroupService(new kafka.admin.ConsumerGroupCommand.ConsumerGroupCommandOptions(args), Map$.MODULE$.empty()) {
+        return new ConsumerGroupCommand.ConsumerGroupService(ConsumerGroupCommandOptions.fromArgs(args), Collections.emptyMap()) {
             @Override
-            public Admin createAdminClient(scala.collection.Map<String, String> configOverrides) {
+            protected Admin createAdminClient(Map<String, String> configOverrides) {
                 return admin;
             }
         };
@@ -291,9 +280,4 @@ private DescribeTopicsResult describeTopicsResult(Collection<String> topics) {
     private Map<String, ListConsumerGroupOffsetsSpec> listConsumerGroupOffsetsSpec() {
         return Collections.singletonMap(GROUP, new ListConsumerGroupOffsetsSpec());
     }
-
-    @SuppressWarnings({"deprecation"})
-    private static <T> scala.collection.immutable.Set<T> set(final Collection<T> set) {
-        return JavaConverters.asScalaSet(new HashSet<>(set)).toSet();
-    }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
index 0a1bcb2daef28..95fe212c53133 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
@@ -17,17 +17,18 @@
 package org.apache.kafka.tools.consumer.group;
 
 import joptsimple.OptionException;
-import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.consumer.GroupProtocol;
 import org.apache.kafka.clients.consumer.RangeAssignor;
 import org.apache.kafka.common.errors.GroupIdNotFoundException;
 import org.apache.kafka.common.errors.GroupNotEmptyException;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.ToolsTestUtils;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -60,10 +61,7 @@ public void testDeleteCmdNonExistingGroup(String quorum) {
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", missingGroup};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            service.deleteGroups();
-            return null;
-        });
+        String output = ToolsTestUtils.grabConsoleOutput(service::deleteGroups);
         assertTrue(output.contains("Group '" + missingGroup + "' could not be deleted due to:") && output.contains(Errors.GROUP_ID_NOT_FOUND.message()),
             "The expected error (" + Errors.GROUP_ID_NOT_FOUND + ") was not detected while deleting consumer group");
     }
@@ -78,8 +76,8 @@ public void testDeleteNonExistingGroup(String quorum) {
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", missingGroup};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        scala.collection.Map<String, Throwable> result = service.deleteGroups();
-        assertTrue(result.size() == 1 && result.contains(missingGroup) && result.get(missingGroup).get().getCause() instanceof GroupIdNotFoundException,
+        Map<String, Throwable> result = service.deleteGroups();
+        assertTrue(result.size() == 1 && result.containsKey(missingGroup) && result.get(missingGroup).getCause() instanceof GroupIdNotFoundException,
             "The expected error (" + Errors.GROUP_ID_NOT_FOUND + ") was not detected while deleting consumer group");
     }
 
@@ -94,14 +92,11 @@ public void testDeleteCmdNonEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.collectGroupMembers(GROUP, false)._2.get().size() == 1,
+            () -> service.collectGroupMembers(GROUP, false).v2.get().size() == 1,
             "The group did not initialize as expected."
         );
 
-        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            service.deleteGroups();
-            return null;
-        });
+        String output = ToolsTestUtils.grabConsoleOutput(service::deleteGroups);
         assertTrue(output.contains("Group '" + GROUP + "' could not be deleted due to:") && output.contains(Errors.NON_EMPTY_GROUP.message()),
             "The expected error (" + Errors.NON_EMPTY_GROUP + ") was not detected while deleting consumer group. Output was: (" + output + ")");
     }
@@ -117,14 +112,14 @@ public void testDeleteNonEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.collectGroupMembers(GROUP, false)._2.get().size() == 1,
+            () -> service.collectGroupMembers(GROUP, false).v2.get().size() == 1,
             "The group did not initialize as expected."
         );
 
-        scala.collection.Map<String, Throwable> result = service.deleteGroups();
-        assertNotNull(result.get(GROUP).get(),
+        Map<String, Throwable> result = service.deleteGroups();
+        assertNotNull(result.get(GROUP),
             "Group was deleted successfully, but it shouldn't have been. Result was:(" + result + ")");
-        assertTrue(result.size() == 1 && result.contains(GROUP) && result.get(GROUP).get().getCause() instanceof GroupNotEmptyException,
+        assertTrue(result.size() == 1 && result.containsKey(GROUP) && result.get(GROUP).getCause() instanceof GroupNotEmptyException,
             "The expected error (" + Errors.NON_EMPTY_GROUP + ") was not detected while deleting consumer group. Result was:(" + result + ")");
     }
 
@@ -139,21 +134,18 @@ public void testDeleteCmdEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state, "Stable"),
             "The group did not initialize as expected."
         );
 
         executor.shutdown();
 
         TestUtils.waitForCondition(
-            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            () -> Objects.equals(service.collectGroupState(GROUP).state, "Empty"),
             "The group did not become empty as expected."
         );
 
-        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            service.deleteGroups();
-            return null;
-        });
+        String output = ToolsTestUtils.grabConsoleOutput(service::deleteGroups);
         assertTrue(output.contains("Deletion of requested consumer groups ('" + GROUP + "') was successful."),
             "The consumer group could not be deleted as expected");
     }
@@ -173,10 +165,10 @@ public void testDeleteCmdAllGroups(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() ->
-            Objects.equals(service.listConsumerGroups().toSet(), set(groups.keySet())) &&
+            new HashSet<>(service.listConsumerGroups()).equals(groups.keySet()) &&
                 groups.keySet().stream().allMatch(groupId -> {
                     try {
-                        return Objects.equals(service.collectGroupState(groupId).state(), "Stable");
+                        return Objects.equals(service.collectGroupState(groupId).state, "Stable");
                     } catch (Exception e) {
                         throw new RuntimeException(e);
                     }
@@ -189,17 +181,14 @@ public void testDeleteCmdAllGroups(String quorum) throws Exception {
         TestUtils.waitForCondition(() ->
             groups.keySet().stream().allMatch(groupId -> {
                 try {
-                    return Objects.equals(service.collectGroupState(groupId).state(), "Empty");
+                    return Objects.equals(service.collectGroupState(groupId).state, "Empty");
                 } catch (Exception e) {
                     throw new RuntimeException(e);
                 }
             }),
             "The group did not become empty as expected.");
 
-        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            service.deleteGroups();
-            return null;
-        }).trim();
+        String output = ToolsTestUtils.grabConsoleOutput(service::deleteGroups).trim();
         Set<String> expectedGroupsForDeletion = groups.keySet();
         Set<String> deletedGroupsGrepped = Arrays.stream(output.substring(output.indexOf('(') + 1, output.indexOf(')')).split(","))
             .map(str -> str.replaceAll("'", "").trim()).collect(Collectors.toSet());
@@ -220,17 +209,17 @@ public void testDeleteEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state, "Stable"),
             "The group did not initialize as expected.");
 
         executor.shutdown();
 
         TestUtils.waitForCondition(
-            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            () -> Objects.equals(service.collectGroupState(GROUP).state, "Empty"),
             "The group did not become empty as expected.");
 
-        scala.collection.Map<String, Throwable> result = service.deleteGroups();
-        assertTrue(result.size() == 1 && result.contains(GROUP) && result.get(GROUP).get() == null,
+        Map<String, Throwable> result = service.deleteGroups();
+        assertTrue(result.size() == 1 && result.containsKey(GROUP) && result.get(GROUP) == null,
             "The consumer group could not be deleted as expected");
     }
 
@@ -246,23 +235,20 @@ public void testDeleteCmdWithMixOfSuccessAndError(String quorum) throws Exceptio
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state, "Stable"),
             "The group did not initialize as expected.");
 
         executor.shutdown();
 
         TestUtils.waitForCondition(
-            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            () -> Objects.equals(service.collectGroupState(GROUP).state, "Empty"),
             "The group did not become empty as expected.");
 
         cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP, "--group", missingGroup};
 
         ConsumerGroupCommand.ConsumerGroupService service2 = getConsumerGroupService(cgcArgs);
 
-        String output = kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            service2.deleteGroups();
-            return null;
-        });
+        String output = ToolsTestUtils.grabConsoleOutput(service2::deleteGroups);
         assertTrue(output.contains("Group '" + missingGroup + "' could not be deleted due to:")
             && output.contains(Errors.GROUP_ID_NOT_FOUND.message())
             && output.contains("These consumer groups were deleted successfully: '" + GROUP + "'"),
@@ -281,23 +267,23 @@ public void testDeleteWithMixOfSuccessAndError(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state(), "Stable"),
+            () -> service.listConsumerGroups().contains(GROUP) && Objects.equals(service.collectGroupState(GROUP).state, "Stable"),
             "The group did not initialize as expected.");
 
         executor.shutdown();
 
         TestUtils.waitForCondition(
-            () -> Objects.equals(service.collectGroupState(GROUP).state(), "Empty"),
+            () -> Objects.equals(service.collectGroupState(GROUP).state, "Empty"),
             "The group did not become empty as expected.");
 
         cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--delete", "--group", GROUP, "--group", missingGroup};
 
         ConsumerGroupCommand.ConsumerGroupService service2 = getConsumerGroupService(cgcArgs);
-        scala.collection.Map<String, Throwable> result = service2.deleteGroups();
+        Map<String, Throwable> result = service2.deleteGroups();
         assertTrue(result.size() == 2 &&
-                result.contains(GROUP) && result.get(GROUP).get() == null &&
-                result.contains(missingGroup) &&
-                result.get(missingGroup).get().getMessage().contains(Errors.GROUP_ID_NOT_FOUND.message()),
+                result.containsKey(GROUP) && result.get(GROUP) == null &&
+                result.containsKey(missingGroup) &&
+                result.get(missingGroup).getMessage().contains(Errors.GROUP_ID_NOT_FOUND.message()),
             "The consumer group deletion did not work as expected");
     }
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
index 2ac093923ce38..05fd00071e15e 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.tools.consumer.group;
 
-import kafka.admin.ConsumerGroupCommand;
 import kafka.utils.TestUtils;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -30,10 +29,12 @@
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.config.Defaults;
+import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.util.Collections;
+import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.ExecutionException;
 
@@ -59,8 +60,8 @@ public void testDeleteOffsetsNonExistingGroup() {
         String topic = "foo:1";
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(group, topic));
 
-        scala.Tuple2<Errors, scala.collection.Map<TopicPartition, Throwable>> res = service.deleteOffsets(group, seq(Collections.singleton(topic)).toList());
-        assertEquals(Errors.GROUP_ID_NOT_FOUND, res._1);
+        Tuple2<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(group, Collections.singletonList(topic));
+        assertEquals(Errors.GROUP_ID_NOT_FOUND, res.v1);
     }
 
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -144,18 +145,18 @@ private void testWithConsumerGroup(java.util.function.Consumer<Runnable> withCon
         withConsumerGroup.accept(() -> {
             String topic = inputPartition >= 0 ? inputTopic + ":" + inputPartition : inputTopic;
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(GROUP, topic));
-            scala.Tuple2<Errors, scala.collection.Map<TopicPartition, Throwable>> res = service.deleteOffsets(GROUP, seq(Collections.singletonList(topic)).toList());
-            Errors topLevelError = res._1;
-            scala.collection.Map<TopicPartition, Throwable> partitions = res._2;
+            Tuple2<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(GROUP, Collections.singletonList(topic));
+            Errors topLevelError = res.v1;
+            Map<TopicPartition, Throwable> partitions = res.v2;
             TopicPartition tp = new TopicPartition(inputTopic, expectedPartition);
             // Partition level error should propagate to top level, unless this is due to a missed partition attempt.
             if (inputPartition >= 0) {
                 assertEquals(expectedError, topLevelError);
             }
             if (expectedError == Errors.NONE)
-                assertNull(partitions.get(tp).get());
+                assertNull(partitions.get(tp));
             else
-                assertEquals(expectedError.exception(), partitions.get(tp).get().getCause());
+                assertEquals(expectedError.exception(), partitions.get(tp).getCause());
         });
     }
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
index f0277d18cd639..85f9f5b836a16 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.tools.consumer.group;
 
-import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.RangeAssignor;
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
@@ -24,17 +23,15 @@
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.tools.ToolsTestUtils;
+import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
-import scala.Function0;
-import scala.Function1;
-import scala.Option;
-import scala.collection.Seq;
-import scala.runtime.BoxedUnit;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -43,6 +40,8 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static org.apache.kafka.test.TestUtils.RANDOM;
 import static org.apache.kafka.tools.ToolsTestUtils.TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES;
@@ -81,7 +80,7 @@ public void testDescribeNonExistingGroup(String quorum, String groupProtocol) {
             cgcArgs.addAll(describeType);
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
-            String output = kafka.utils.TestUtils.grabConsoleOutput(describeGroups(service));
+            String output = ToolsTestUtils.grabConsoleOutput(describeGroups(service));
             assertTrue(output.contains("Consumer group '" + missingGroup + "' does not exist."),
                 "Expected error was not detected for describe option '" + String.join(" ", describeType) + "'");
         }
@@ -129,7 +128,7 @@ public void testDescribeWithStateValue(String quorum) {
 
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
     @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
-    public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupProtocol) {
+    public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupProtocol) throws Exception {
         String group = "missing.group";
         createOffsetsTopic(listenerName(), new Properties());
 
@@ -139,14 +138,14 @@ public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupPro
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(group);
-        assertTrue(res._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res._2.map(Seq::isEmpty).getOrElse(() -> false),
+        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(group);
+        assertTrue(res.v1.map(s -> s.contains("Dead")).orElse(false) && res.v2.map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
     }
 
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
     @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
-    public void testDescribeMembersOfNonExistingGroup(String quorum, String groupProtocol) {
+    public void testDescribeMembersOfNonExistingGroup(String quorum, String groupProtocol) throws Exception {
         String group = "missing.group";
         createOffsetsTopic(listenerName(), new Properties());
 
@@ -156,18 +155,18 @@ public void testDescribeMembersOfNonExistingGroup(String quorum, String groupPro
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(group, false);
-        assertTrue(res._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res._2.map(Seq::isEmpty).getOrElse(() -> false),
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(group, false);
+        assertTrue(res.v1.map(s -> s.contains("Dead")).orElse(false) && res.v2.map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res2 = service.collectGroupMembers(group, true);
-        assertTrue(res2._1.map(s -> s.contains("Dead")).getOrElse(() -> false) && res2._2.map(Seq::isEmpty).getOrElse(() -> false),
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res2 = service.collectGroupMembers(group, true);
+        assertTrue(res2.v1.map(s -> s.contains("Dead")).orElse(false) && res2.v2.map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "' (verbose option).");
     }
 
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
     @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
-    public void testDescribeStateOfNonExistingGroup(String quorum, String groupProtocol) {
+    public void testDescribeStateOfNonExistingGroup(String quorum, String groupProtocol) throws Exception {
         String group = "missing.group";
         createOffsetsTopic(listenerName(), new Properties());
 
@@ -177,9 +176,9 @@ public void testDescribeStateOfNonExistingGroup(String quorum, String groupProto
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        ConsumerGroupCommand.GroupState state = service.collectGroupState(group);
-        assertTrue(Objects.equals(state.state(), "Dead") && state.numMembers() == 0 &&
-                state.coordinator() != null && !brokers().filter(s -> s.config().brokerId() == state.coordinator().id()).isEmpty(),
+        GroupState state = service.collectGroupState(group);
+        assertTrue(Objects.equals(state.state, "Dead") && state.numMembers == 0 &&
+                state.coordinator != null && !brokers().filter(s -> s.config().brokerId() == state.coordinator.id()).isEmpty(),
             "Expected the state to be 'Dead', with no members in the group '" + group + "'."
         );
     }
@@ -198,8 +197,8 @@ public void testDescribeExistingGroup(String quorum, String groupProtocol) throw
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
-                return res._1.trim().split("\n").length == 2 && res._2.isEmpty();
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res.v1.trim().split("\n").length == 2 && res.v2.isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -227,9 +226,9 @@ public void testDescribeExistingGroups(String quorum, String groupProtocol) thro
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
-                long numLines = Arrays.stream(res._1.trim().split("\n")).filter(line -> !line.isEmpty()).count();
-                return (numLines == expectedNumLines) && res._2.isEmpty();
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res.v1.trim().split("\n")).filter(line -> !line.isEmpty()).count();
+                return (numLines == expectedNumLines) && res.v2.isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -253,9 +252,9 @@ public void testDescribeAllExistingGroups(String quorum, String groupProtocol) t
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
-                long numLines = Arrays.stream(res._1.trim().split("\n")).filter(s -> !s.isEmpty()).count();
-                return (numLines == expectedNumLines) && res._2.isEmpty();
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res.v1.trim().split("\n")).filter(s -> !s.isEmpty()).count();
+                return (numLines == expectedNumLines) && res.v2.isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -272,26 +271,28 @@ public void testDescribeOffsetsOfExistingGroup(String quorum, String groupProtoc
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
-            Option<String> state = groupOffsets._1;
-            Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>> assignments = groupOffsets._2;
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+            Optional<String> state = groupOffsets.v1;
+            Optional<Collection<PartitionAssignmentState>> assignments = groupOffsets.v2;
 
-            Function1<ConsumerGroupCommand.PartitionAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+            Predicate<PartitionAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
-            boolean res = state.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                assignments.isDefined() &&
-                assignments.get().count(isGrp) == 1;
+            boolean res = state.map(s -> s.contains("Stable")).orElse(false) &&
+                assignments.isPresent() &&
+                assignments.get().stream().filter(isGrp).count() == 1;
 
             if (!res)
                 return false;
 
-            @SuppressWarnings("cast")
-            ConsumerGroupCommand.PartitionAssignmentState partitionState =
-                (ConsumerGroupCommand.PartitionAssignmentState) assignments.get().filter(isGrp).head();
+            Optional<PartitionAssignmentState> maybePartitionState = assignments.get().stream().filter(isGrp).findFirst();
+            if (!maybePartitionState.isPresent())
+                return false;
+
+            PartitionAssignmentState partitionState = maybePartitionState.get();
 
-            return !partitionState.consumerId().map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
-                !partitionState.clientId().map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
-                !partitionState.host().map(h -> h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+            return !partitionState.consumerId.map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) &&
+                !partitionState.clientId.map(s0 -> s0.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) &&
+                !partitionState.host.map(h -> h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false);
         }, "Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for group " + GROUP + ".");
     }
 
@@ -306,32 +307,34 @@ public void testDescribeMembersOfExistingGroup(String quorum, String groupProtoc
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> groupMembers = service.collectGroupMembers(GROUP, false);
-            Option<String> state = groupMembers._1;
-            Option<Seq<ConsumerGroupCommand.MemberAssignmentState>> assignments = groupMembers._2;
+            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> groupMembers = service.collectGroupMembers(GROUP, false);
+            Optional<String> state = groupMembers.v1;
+            Optional<Collection<MemberAssignmentState>> assignments = groupMembers.v2;
 
-            Function1<ConsumerGroupCommand.MemberAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+            Predicate<MemberAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
-            boolean res = state.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                assignments.isDefined() &&
-                assignments.get().count(s -> Objects.equals(s.group(), GROUP)) == 1;
+            boolean res = state.map(s -> s.contains("Stable")).orElse(false) &&
+                assignments.isPresent() &&
+                assignments.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 1;
 
             if (!res)
                 return false;
 
-            @SuppressWarnings("cast")
-            ConsumerGroupCommand.MemberAssignmentState assignmentState =
-                (ConsumerGroupCommand.MemberAssignmentState) assignments.get().filter(isGrp).head();
+            Optional<MemberAssignmentState> maybeAssignmentState = assignments.get().stream().filter(isGrp).findFirst();
+            if (!maybeAssignmentState.isPresent())
+                return false;
+
+            MemberAssignmentState assignmentState = maybeAssignmentState.get();
 
-            return !Objects.equals(assignmentState.consumerId(), ConsumerGroupCommand.MISSING_COLUMN_VALUE()) &&
-                !Objects.equals(assignmentState.clientId(), ConsumerGroupCommand.MISSING_COLUMN_VALUE()) &&
-                !Objects.equals(assignmentState.host(), ConsumerGroupCommand.MISSING_COLUMN_VALUE());
+            return !Objects.equals(assignmentState.consumerId, ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
+                !Objects.equals(assignmentState.clientId, ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
+                !Objects.equals(assignmentState.host, ConsumerGroupCommand.MISSING_COLUMN_VALUE);
         }, "Expected a 'Stable' group status, rows and valid member information for group " + GROUP + ".");
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
 
-        if (res._2.isDefined()) {
-            assertTrue(res._2.get().size() == 1 && res._2.get().iterator().next().assignment().size() == 1,
+        if (res.v2.isPresent()) {
+            assertTrue(res.v2.get().size() == 1 && res.v2.get().iterator().next().assignment.size() == 1,
                 "Expected a topic partition assigned to the single group member for group " + GROUP);
         } else {
             fail("Expected partition assignments for members of group " + GROUP);
@@ -354,12 +357,12 @@ public void testDescribeStateOfExistingGroup(String quorum, String groupProtocol
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Stable") &&
-                state.numMembers() == 1 &&
-                Objects.equals(state.assignmentStrategy(), "range") &&
-                state.coordinator() != null &&
-                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Stable") &&
+                state.numMembers == 1 &&
+                Objects.equals(state.assignmentStrategy, "range") &&
+                state.coordinator != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator.id()) > 0;
         }, "Expected a 'Stable' group status, with one member and round robin assignment strategy for group " + GROUP + ".");
     }
 
@@ -381,12 +384,12 @@ public void testDescribeStateOfExistingGroupWithNonDefaultAssignor(String quorum
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Stable") &&
-                state.numMembers() == 1 &&
-                Objects.equals(state.assignmentStrategy(), expectedName) &&
-                state.coordinator() != null &&
-                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Stable") &&
+                state.numMembers == 1 &&
+                Objects.equals(state.assignmentStrategy, expectedName) &&
+                state.coordinator != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator.id()) > 0;
         }, "Expected a 'Stable' group status, with one member and " + expectedName + " assignment strategy for group " + GROUP + ".");
     }
 
@@ -404,14 +407,14 @@ public void testDescribeExistingGroupWithNoMembers(String quorum, String groupPr
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
-                return res._1.trim().split("\n").length == 2 && res._2.isEmpty();
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res.v1.trim().split("\n").length == 2 && res.v2.isEmpty();
             }, "Expected describe group results with one data row for describe type '" + String.join(" ", describeType) + "'");
 
             // stop the consumer so the group has no active member anymore
             executor.shutdown();
             TestUtils.waitForCondition(
-                () -> kafka.utils.TestUtils.grabConsoleError(describeGroups(service)).contains("Consumer group '" + group + "' has no active members."),
+                () -> ToolsTestUtils.grabConsoleError(describeGroups(service)).contains("Consumer group '" + group + "' has no active members."),
                 "Expected no active member in describe group results with describe type " + String.join(" ", describeType));
         }
     }
@@ -428,26 +431,25 @@ public void testDescribeOffsetsOfExistingGroupWithNoMembers(String quorum, Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
-                && res._2.map(c -> c.exists(assignment -> Objects.equals(assignment.group(), GROUP) && assignment.offset().isDefined())).getOrElse(() -> false);
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.v1.map(s -> s.contains("Stable")).orElse(false)
+                && res.v2.map(c -> c.stream().anyMatch(assignment -> Objects.equals(assignment.group, GROUP) && assignment.offset.isPresent())).orElse(false);
         }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
 
         // stop the consumer so the group has no active member anymore
         executor.shutdown();
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> offsets = service.collectGroupOffsets(GROUP);
-            Option<String> state = offsets._1;
-            Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>> assignments = offsets._2;
-            @SuppressWarnings("unchecked")
-            Seq<ConsumerGroupCommand.PartitionAssignmentState> testGroupAssignments = assignments.get().filter(a -> Objects.equals(a.group(), GROUP)).toSeq();
-            ConsumerGroupCommand.PartitionAssignmentState assignment = testGroupAssignments.head();
-            return state.map(s -> s.contains("Empty")).getOrElse(() -> false) &&
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> offsets = service.collectGroupOffsets(GROUP);
+            Optional<String> state = offsets.v1;
+            Optional<Collection<PartitionAssignmentState>> assignments = offsets.v2;
+            List<PartitionAssignmentState> testGroupAssignments = assignments.get().stream().filter(a -> Objects.equals(a.group, GROUP)).collect(Collectors.toList());
+            PartitionAssignmentState assignment = testGroupAssignments.get(0);
+            return state.map(s -> s.contains("Empty")).orElse(false) &&
                 testGroupAssignments.size() == 1 &&
-                assignment.consumerId().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) && // the member should be gone
-                assignment.clientId().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
-                assignment.host().map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+                assignment.consumerId.map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) && // the member should be gone
+                assignment.clientId.map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) &&
+                assignment.host.map(c -> c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false);
         }, "failed to collect group offsets");
     }
 
@@ -463,17 +465,17 @@ public void testDescribeMembersOfExistingGroupWithNoMembers(String quorum, Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
-                && res._2.map(c -> c.exists(m -> Objects.equals(m.group(), GROUP))).getOrElse(() -> false);
+            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.v1.map(s -> s.contains("Stable")).orElse(false)
+                && res.v2.map(c -> c.stream().anyMatch(m -> Objects.equals(m.group, GROUP))).orElse(false);
         }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
 
         // stop the consumer so the group has no active member anymore
         executor.shutdown();
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res._1.map(s -> s.contains("Empty")).getOrElse(() -> false) && res._2.isDefined() && res._2.get().isEmpty();
+            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.v1.map(s -> s.contains("Empty")).orElse(false) && res.v2.isPresent() && res.v2.get().isEmpty();
         }, "Expected no member in describe group members results for group '" + GROUP + "'");
     }
 
@@ -489,19 +491,19 @@ public void testDescribeStateOfExistingGroupWithNoMembers(String quorum, String
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Stable") &&
-                state.numMembers() == 1 &&
-                state.coordinator() != null &&
-                brokers().count(s -> s.config().brokerId() == state.coordinator().id()) > 0;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Stable") &&
+                state.numMembers == 1 &&
+                state.coordinator != null &&
+                brokers().count(s -> s.config().brokerId() == state.coordinator.id()) > 0;
         }, "Expected the group '" + GROUP + "' to initially become stable, and have a single member.");
 
         // stop the consumer so the group has no active member anymore
         executor.shutdown();
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Empty") && state.numMembers() == 0;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Empty") && state.numMembers == 0;
         }, "Expected the group '" + GROUP + "' to become empty after the only member leaving.");
     }
 
@@ -519,9 +521,9 @@ public void testDescribeWithConsumersWithoutAssignedPartitions(String quorum, St
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
                 int expectedNumRows = DESCRIBE_TYPE_MEMBERS.contains(describeType) ? 3 : 2;
-                return res._2.isEmpty() && res._1.trim().split("\n").length == expectedNumRows;
+                return res.v2.isEmpty() && res.v1.trim().split("\n").length == expectedNumRows;
             }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
         }
     }
@@ -538,11 +540,11 @@ public void testDescribeOffsetsWithConsumersWithoutAssignedPartitions(String quo
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                res._2.isDefined() &&
-                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 1 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.partition().isDefined()) == 1;
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.v1.map(s -> s.contains("Stable")).isPresent() &&
+                res.v2.isPresent() &&
+                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 1 &&
+                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 1;
         }, "Expected rows for consumers with no assigned partitions in describe group results");
     }
 
@@ -558,18 +560,18 @@ public void testDescribeMembersWithConsumersWithoutAssignedPartitions(String quo
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                res._2.isDefined() &&
-                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 1) == 1 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 0) == 1 &&
-                res._2.get().forall(s -> s.assignment().isEmpty());
+            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
+                res.v2.isPresent() &&
+                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 1 &&
+                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0).count() == 1 &&
+                res.v2.get().stream().allMatch(s -> s.assignment.isEmpty());
         }, "Expected rows for consumers with no assigned partitions in describe group results");
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
-        assertTrue(res._1.map(s -> s.contains("Stable")).getOrElse(() -> false)
-                && res._2.map(c -> c.exists(s -> !s.assignment().isEmpty())).getOrElse(() -> false),
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res.v1.map(s -> s.contains("Stable")).orElse(false)
+                && res.v2.map(c -> c.stream().anyMatch(s -> !s.assignment.isEmpty())).orElse(false),
             "Expected additional columns in verbose version of describe members");
     }
 
@@ -585,8 +587,8 @@ public void testDescribeStateWithConsumersWithoutAssignedPartitions(String quoru
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Stable") && state.numMembers() == 2;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Stable") && state.numMembers == 2;
         }, "Expected two consumers in describe group results");
     }
 
@@ -606,9 +608,9 @@ public void testDescribeWithMultiPartitionTopicAndMultipleConsumers(String quoru
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                scala.Tuple2<String, String> res = kafka.utils.TestUtils.grabConsoleOutputAndError(describeGroups(service));
+                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
                 int expectedNumRows = DESCRIBE_TYPE_STATE.contains(describeType) ? 2 : 3;
-                return res._2.isEmpty() && res._1.trim().split("\n").length == expectedNumRows;
+                return res.v2.isEmpty() && res.v1.trim().split("\n").length == expectedNumRows;
             }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
         }
     }
@@ -627,12 +629,12 @@ public void testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                res._2.isDefined() &&
-                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.partition().isDefined()) == 2 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && !x.partition().isDefined()) == 0;
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
+                res.v2.isPresent() &&
+                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 2 &&
+                res.v2.get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && !x.partition.isPresent());
         }, "Expected two rows (one row per consumer) in describe group results.");
     }
 
@@ -650,16 +652,16 @@ public void testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                res._2.isDefined() &&
-                res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 1) == 2 &&
-                res._2.get().count(x -> Objects.equals(x.group(), GROUP) && x.numPartitions() == 0) == 0;
+            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
+                res.v2.isPresent() &&
+                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 2 &&
+                res.v2.get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0);
         }, "Expected two rows (one row per consumer) in describe group members results.");
 
-        scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
-        assertTrue(res._1.map(s -> s.contains("Stable")).getOrElse(() -> false) && res._2.map(s -> s.count(x -> x.assignment().isEmpty())).getOrElse(() -> 0) == 0,
+        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res.v1.map(s -> s.contains("Stable")).orElse(false) && res.v2.map(s -> s.stream().filter(x -> x.assignment.isEmpty()).count()).orElse(0L) == 0,
             "Expected additional columns in verbose version of describe members");
     }
 
@@ -677,8 +679,8 @@ public void testDescribeStateWithMultiPartitionTopicAndMultipleConsumers(String
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            ConsumerGroupCommand.GroupState state = service.collectGroupState(GROUP);
-            return Objects.equals(state.state(), "Stable") && Objects.equals(state.group(), GROUP) && state.numMembers() == 2;
+            GroupState state = service.collectGroupState(GROUP);
+            return Objects.equals(state.state, "Stable") && Objects.equals(state.group, GROUP) && state.numMembers == 2;
         }, "Expected a stable group with two members in describe group state result.");
     }
 
@@ -696,9 +698,9 @@ public void testDescribeSimpleConsumerGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res._1.map(s -> s.contains("Empty")).getOrElse(() -> false)
-                && res._2.isDefined() && res._2.get().count(s -> Objects.equals(s.group(), GROUP)) == 2;
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.v1.map(s -> s.contains("Empty")).orElse(false)
+                && res.v2.isPresent() && res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2;
         }, "Expected a stable group with two members in describe group state result.");
     }
 
@@ -796,32 +798,33 @@ public void testDescribeNonOffsetCommitGroup(String quorum, String groupProtocol
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            scala.Tuple2<Option<String>, Option<Seq<ConsumerGroupCommand.PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
 
-            Function1<ConsumerGroupCommand.PartitionAssignmentState, Object> isGrp = s -> Objects.equals(s.group(), GROUP);
+            Predicate<PartitionAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
-            boolean res = groupOffsets._1.map(s -> s.contains("Stable")).getOrElse(() -> false) &&
-                groupOffsets._2.isDefined() &&
-                groupOffsets._2.get().count(isGrp) == 1;
+            boolean res = groupOffsets.v1.map(s -> s.contains("Stable")).orElse(false) &&
+                groupOffsets.v2.isPresent() &&
+                groupOffsets.v2.get().stream().filter(isGrp).count() == 1;
 
             if (!res)
                 return false;
 
-            @SuppressWarnings("cast")
-            ConsumerGroupCommand.PartitionAssignmentState assignmentState =
-                (ConsumerGroupCommand.PartitionAssignmentState) groupOffsets._2.get().filter(isGrp).head();
+            Optional<PartitionAssignmentState> maybeAssignmentState = groupOffsets.v2.get().stream().filter(isGrp).findFirst();
+            if (!maybeAssignmentState.isPresent())
+                return false;
+
+            PartitionAssignmentState assignmentState = maybeAssignmentState.get();
 
-            return assignmentState.consumerId().map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
-                assignmentState.clientId().map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false) &&
-                assignmentState.host().map(h -> !h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE())).getOrElse(() -> false);
+            return assignmentState.consumerId.map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) &&
+                assignmentState.clientId.map(c -> !c.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false) &&
+                assignmentState.host.map(h -> !h.trim().equals(ConsumerGroupCommand.MISSING_COLUMN_VALUE)).orElse(false);
         }, "Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for non-offset-committing group " + GROUP + ".");
     }
 
-    private Function0<BoxedUnit> describeGroups(ConsumerGroupCommand.ConsumerGroupService service) {
+    private Runnable describeGroups(ConsumerGroupCommand.ConsumerGroupService service) {
         return () -> {
             try {
                 service.describeGroups();
-                return null;
             } catch (Exception e) {
                 throw new RuntimeException(e);
             }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
index ba5ebd254fc93..6fd2928da9f2f 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ListConsumerGroupTest.java
@@ -17,16 +17,17 @@
 package org.apache.kafka.tools.consumer.group;
 
 import joptsimple.OptionException;
-import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.admin.ConsumerGroupListing;
 import org.apache.kafka.common.ConsumerGroupState;
 import org.apache.kafka.common.GroupType;
 import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.Test;
+import org.apache.kafka.tools.ToolsTestUtils;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashSet;
@@ -58,11 +59,11 @@ public void testListConsumerGroupsWithoutFilters(String quorum, String groupProt
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--list"};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        scala.collection.Set<String> expectedGroups = set(Arrays.asList(GROUP, simpleGroup, PROTOCOL_GROUP));
-        final AtomicReference<scala.collection.Set> foundGroups = new AtomicReference<>();
+        Set<String> expectedGroups = set(Arrays.asList(GROUP, simpleGroup, PROTOCOL_GROUP));
+        final AtomicReference<Set> foundGroups = new AtomicReference<>();
 
         TestUtils.waitForCondition(() -> {
-            foundGroups.set(service.listConsumerGroups().toSet());
+            foundGroups.set(set(service.listConsumerGroups()));
             return Objects.equals(expectedGroups, foundGroups.get());
         }, "Expected --list to show groups " + expectedGroups + ", but found " + foundGroups.get() + ".");
     }
@@ -272,7 +273,7 @@ public void testListConsumerGroupsWithTypesConsumerProtocol(String quorum, Strin
 
     @Test
     public void testConsumerGroupStatesFromString() {
-        scala.collection.Set<ConsumerGroupState> result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable");
+        Set<ConsumerGroupState> result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable");
         assertEquals(set(Collections.singleton(ConsumerGroupState.STABLE)), result);
 
         result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable, PreparingRebalance");
@@ -299,7 +300,7 @@ public void testConsumerGroupStatesFromString() {
 
     @Test
     public void testConsumerGroupTypesFromString() {
-        scala.collection.Set<GroupType> result = ConsumerGroupCommand.consumerGroupTypesFromString("consumer");
+        Set<GroupType> result = ConsumerGroupCommand.consumerGroupTypesFromString("consumer");
         assertEquals(set(Collections.singleton(GroupType.CONSUMER)), result);
 
         result = ConsumerGroupCommand.consumerGroupTypesFromString("consumer, classic");
@@ -475,9 +476,9 @@ private static void assertGroupListing(
         Set<ConsumerGroupState> stateFilterSet,
         Set<ConsumerGroupListing> expectedListing
     ) throws Exception {
-        final AtomicReference<scala.collection.Set> foundListing = new AtomicReference<>();
+        final AtomicReference<Set<ConsumerGroupListing>> foundListing = new AtomicReference<>();
         TestUtils.waitForCondition(() -> {
-            foundListing.set(service.listConsumerGroupsWithFilters(set(typeFilterSet), set(stateFilterSet)).toSet());
+            foundListing.set(set(service.listConsumerGroupsWithFilters(set(typeFilterSet), set(stateFilterSet))));
             return Objects.equals(set(expectedListing), foundListing.get());
         }, () -> "Expected to show groups " + expectedListing + ", but found " + foundListing.get() + ".");
     }
@@ -498,7 +499,7 @@ private static void validateListOutput(
     ) throws InterruptedException {
         final AtomicReference<String> out = new AtomicReference<>("");
         TestUtils.waitForCondition(() -> {
-            String output = runAndGrabConsoleOutput(args);
+            String output = ToolsTestUtils.grabConsoleOutput(() -> ConsumerGroupCommand.main(args.toArray(new String[0])));
             out.set(output);
 
             int index = 0;
@@ -522,12 +523,7 @@ private static void validateListOutput(
         }, () -> String.format("Expected header=%s and groups=%s, but found:%n%s", expectedHeader, expectedRows, out.get()));
     }
 
-    private static String runAndGrabConsoleOutput(
-        List<String> args
-    ) {
-        return kafka.utils.TestUtils.grabConsoleOutput(() -> {
-            ConsumerGroupCommand.main(args.toArray(new String[0]));
-            return null;
-        });
+    public static <T> Set<T> set(Collection<T> set) {
+        return new HashSet<>(set);
     }
 }
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java
index dca263452231f..4d805ce68f881 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ResetConsumerGroupOffsetTest.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.tools.consumer.group;
 
 import joptsimple.OptionException;
-import kafka.admin.ConsumerGroupCommand;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.GroupProtocol;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -25,7 +24,6 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.Test;
-import scala.Option;
 
 import java.io.BufferedWriter;
 import java.io.File;
@@ -98,10 +96,10 @@ public void testResetOffsetsNotExistingGroup() throws Exception {
         ConsumerGroupCommand.ConsumerGroupService consumerGroupCommand = getConsumerGroupService(args);
         // Make sure we got a coordinator
         TestUtils.waitForCondition(
-            () -> Objects.equals(consumerGroupCommand.collectGroupState(group).coordinator().host(), "localhost"),
+            () -> Objects.equals(consumerGroupCommand.collectGroupState(group).coordinator.host(), "localhost"),
             "Can't find a coordinator");
-        Option<scala.collection.Map<TopicPartition, OffsetAndMetadata>> resetOffsets = consumerGroupCommand.resetOffsets().get(group);
-        assertTrue(resetOffsets.isDefined() && resetOffsets.get().isEmpty());
+        Map<TopicPartition, OffsetAndMetadata> resetOffsets = consumerGroupCommand.resetOffsets().get(group);
+        assertTrue(resetOffsets.isEmpty());
         assertTrue(committedOffsets(TOPIC, group).isEmpty());
     }
 
@@ -211,7 +209,7 @@ public void testResetOffsetsByDurationToEarliest() throws Exception {
     }
 
     @Test
-    public void testResetOffsetsByDurationFallbackToLatestWhenNoRecords() throws Exception {
+    public void testResetOffsetsByDurationFallbackToLatestWhenNoRecords() {
         String topic = "foo2";
         String[] args = buildArgsForGroup(GROUP, "--topic", topic, "--by-duration", "PT1M", "--execute");
         createTopic(topic, 1, 1, new Properties(), listenerName(), new Properties());
@@ -326,7 +324,7 @@ public void testResetOffsetsToEarliestOnTopics() throws Exception {
         TopicPartition tp1 = new TopicPartition(topic1, 0);
         TopicPartition tp2 = new TopicPartition(topic2, 0);
 
-        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(consumerGroupCommand.resetOffsets().get(GROUP));
+        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(resetOffsets(consumerGroupCommand).get(GROUP));
         Map<TopicPartition, Long> expMap = new HashMap<>();
         expMap.put(tp1, 0L);
         expMap.put(tp2, 0L);
@@ -357,7 +355,7 @@ public void testResetOffsetsToEarliestOnTopicsAndPartitions() throws Exception {
 
         TopicPartition tp1 = new TopicPartition(topic1, 1);
         TopicPartition tp2 = new TopicPartition(topic2, 1);
-        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(consumerGroupCommand.resetOffsets().get(GROUP));
+        Map<TopicPartition, Long> allResetOffsets = toOffsetMap(resetOffsets(consumerGroupCommand).get(GROUP));
         Map<TopicPartition, Long> expMap = new HashMap<>();
         expMap.put(tp1, 0L);
         expMap.put(tp2, 0L);
@@ -386,7 +384,7 @@ public void testResetOffsetsExportImportPlanSingleGroupArg() throws Exception {
 
         File file = TestUtils.tempFile("reset", ".csv");
 
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
         BufferedWriter bw = new BufferedWriter(new FileWriter(file));
         bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets));
         bw.close();
@@ -398,7 +396,7 @@ public void testResetOffsetsExportImportPlanSingleGroupArg() throws Exception {
 
         String[] cgcArgsExec = buildArgsForGroup(GROUP, "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
         ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec);
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
         assertEquals(exp1, toOffsetMap(importedOffsets.get(GROUP)));
 
         adminZkClient().deleteTopic(topic);
@@ -430,7 +428,7 @@ public void testResetOffsetsExportImportPlan() throws Exception {
 
         File file = TestUtils.tempFile("reset", ".csv");
 
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> exportedOffsets = consumerGroupCommand.resetOffsets();
         BufferedWriter bw = new BufferedWriter(new FileWriter(file));
         bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets));
         bw.close();
@@ -447,14 +445,14 @@ public void testResetOffsetsExportImportPlan() throws Exception {
         // Multiple --group's offset import
         String[] cgcArgsExec = buildArgsForGroups(Arrays.asList(group1, group2), "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
         ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec);
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> importedOffsets = consumerGroupCommandExec.resetOffsets();
         assertEquals(exp1, toOffsetMap(importedOffsets.get(group1)));
         assertEquals(exp2, toOffsetMap(importedOffsets.get(group2)));
 
         // Single --group offset import using "group,topic,partition,offset" csv format
         String[] cgcArgsExec2 = buildArgsForGroup(group1, "--all-topics", "--from-file", file.getCanonicalPath(), "--dry-run");
         ConsumerGroupCommand.ConsumerGroupService consumerGroupCommandExec2 = getConsumerGroupService(cgcArgsExec2);
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> importedOffsets2 = consumerGroupCommandExec2.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> importedOffsets2 = consumerGroupCommandExec2.resetOffsets();
         assertEquals(exp1, toOffsetMap(importedOffsets2.get(group1)));
 
         adminZkClient().deleteTopic(TOPIC);
@@ -503,9 +501,9 @@ private void awaitConsumerProgress(String topic,
 
     private void awaitConsumerGroupInactive(ConsumerGroupCommand.ConsumerGroupService consumerGroupService, String group) throws Exception {
         TestUtils.waitForCondition(() -> {
-            String state = consumerGroupService.collectGroupState(group).state();
+            String state = consumerGroupService.collectGroupState(group).state;
             return Objects.equals(state, "Empty") || Objects.equals(state, "Dead");
-        }, "Expected that consumer group is inactive. Actual state: " + consumerGroupService.collectGroupState(group).state());
+        }, "Expected that consumer group is inactive. Actual state: " + consumerGroupService.collectGroupState(group).state);
     }
 
     private void resetAndAssertOffsets(String[] args,
@@ -521,25 +519,17 @@ private void resetAndAssertOffsets(String[] args,
         Map<String, Map<TopicPartition, Long>> expectedOffsets = topics.stream().collect(Collectors.toMap(
             Function.identity(),
             topic -> Collections.singletonMap(new TopicPartition(topic, 0), expectedOffset)));
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> resetOffsetsResultByGroup = consumerGroupCommand.resetOffsets();
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> resetOffsetsResultByGroup = resetOffsets(consumerGroupCommand);
 
         try {
             for (final String topic : topics) {
-                resetOffsetsResultByGroup.foreach(entry -> {
-                    String group = entry._1;
-                    scala.collection.Map<TopicPartition, OffsetAndMetadata> partitionInfo = entry._2;
+                resetOffsetsResultByGroup.forEach((group, partitionInfo) -> {
                     Map<TopicPartition, Long> priorOffsets = committedOffsets(topic, group);
-                    Map<TopicPartition, Long> offsets = new HashMap<>();
-                    partitionInfo.foreach(partitionInfoEntry -> {
-                        TopicPartition tp = partitionInfoEntry._1;
-                        OffsetAndMetadata offsetAndMetadata = partitionInfoEntry._2;
-                        if (Objects.equals(tp.topic(), topic))
-                            offsets.put(tp, offsetAndMetadata.offset());
-                        return null;
-                    });
-                    assertEquals(expectedOffsets.get(topic), offsets);
+                    assertEquals(expectedOffsets.get(topic),
+                        partitionInfo.entrySet().stream()
+                            .filter(entry -> Objects.equals(entry.getKey().topic(), topic))
+                            .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().offset())));
                     assertEquals(dryRun ? priorOffsets : expectedOffsets.get(topic), committedOffsets(topic, group));
-                    return null;
                 });
             }
         } finally {
@@ -550,35 +540,22 @@ private void resetAndAssertOffsets(String[] args,
     private void resetAndAssertOffsetsCommitted(ConsumerGroupCommand.ConsumerGroupService consumerGroupService,
                                                 Map<TopicPartition, Long> expectedOffsets,
                                                 String topic) {
-        scala.collection.Map<String, scala.collection.Map<TopicPartition, OffsetAndMetadata>> allResetOffsets = consumerGroupService.resetOffsets();
-
-        allResetOffsets.foreach(entry -> {
-            String group = entry._1;
-            scala.collection.Map<TopicPartition, OffsetAndMetadata> offsetsInfo = entry._2;
-            offsetsInfo.foreach(offsetInfoEntry -> {
-                TopicPartition tp = offsetInfoEntry._1;
-                OffsetAndMetadata offsetMetadata = offsetInfoEntry._2;
+        Map<String, Map<TopicPartition, OffsetAndMetadata>> allResetOffsets = resetOffsets(consumerGroupService);
+
+        allResetOffsets.forEach((group, offsetsInfo) -> {
+            offsetsInfo.forEach((tp, offsetMetadata) -> {
                 assertEquals(offsetMetadata.offset(), expectedOffsets.get(tp));
                 assertEquals(expectedOffsets, committedOffsets(topic, group));
-                return null;
             });
-            return null;
         });
     }
 
-    Map<TopicPartition, Long> toOffsetMap(Option<scala.collection.Map<TopicPartition, OffsetAndMetadata>> map) {
-        assertTrue(map.isDefined());
-        Map<TopicPartition, Long> res = new HashMap<>();
-        map.foreach(m -> {
-            m.foreach(entry -> {
-                TopicPartition tp = entry._1;
-                OffsetAndMetadata offsetAndMetadata = entry._2;
-                res.put(tp, offsetAndMetadata.offset());
-                return null;
-            });
-            return null;
-        });
-        return res;
+    private Map<String, Map<TopicPartition, OffsetAndMetadata>> resetOffsets(ConsumerGroupCommand.ConsumerGroupService consumerGroupService) {
+        return consumerGroupService.resetOffsets();
+    }
+
+    Map<TopicPartition, Long> toOffsetMap(Map<TopicPartition, OffsetAndMetadata> map) {
+        return map.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().offset()));
     }
 
     private String[] addTo(String[] args, String...extra) {
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java
index 06b727a59a8fd..1ce1f909dd5e5 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/SaslClientsWithInvalidCredentialsTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.tools.consumer.group;
 
-import kafka.admin.ConsumerGroupCommand;
 import kafka.api.AbstractSaslTest;
 import kafka.api.Both$;
 import kafka.utils.JaasTestUtils;
@@ -36,7 +35,6 @@
 import scala.Some$;
 import scala.collection.JavaConverters;
 import scala.collection.Seq;
-import scala.collection.immutable.Map$;
 
 import java.io.File;
 import java.io.IOException;
@@ -164,8 +162,8 @@ private ConsumerGroupCommand.ConsumerGroupService prepareConsumerGroupService()
             "--describe",
             "--group", "test.group",
             "--command-config", propsFile.getAbsolutePath()};
-        ConsumerGroupCommand.ConsumerGroupCommandOptions opts = new ConsumerGroupCommand.ConsumerGroupCommandOptions(cgcArgs);
-        return new ConsumerGroupCommand.ConsumerGroupService(opts, Map$.MODULE$.empty());
+        ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(cgcArgs);
+        return new ConsumerGroupCommand.ConsumerGroupService(opts, Collections.emptyMap());
     }
 
     private void verifyAuthenticationException(Executable action) {

From 3e3c618bdc90ea225632a6aeecb0e65b5ac8294f Mon Sep 17 00:00:00 2001
From: Dongnuo Lyu <139248811+dongnuo123@users.noreply.github.com>
Date: Wed, 20 Mar 2024 03:49:11 -0400
Subject: [PATCH 178/258] KAFKA-16313: Offline group protocol migration
 (#15546)

This patch enables an empty classic group to be automatically converted to a new consumer group and vice versa.

Reviewers: David Jacot <djacot@confluent.io>
---
 .../group/GroupCoordinatorShard.java          |   2 +-
 .../group/GroupMetadataManager.java           | 147 +++++++++++++----
 .../group/GroupCoordinatorShardTest.java      |   8 +-
 .../group/GroupMetadataManagerTest.java       | 156 +++++++++++++++++-
 .../GroupMetadataManagerTestContext.java      |   4 +-
 5 files changed, 266 insertions(+), 51 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
index b55dd91cc1cf3..12c194c331b1f 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
@@ -388,7 +388,7 @@ public CoordinatorResult<DeleteGroupsResponseData.DeletableGroupResultCollection
             try {
                 groupMetadataManager.validateDeleteGroup(groupId);
                 numDeletedOffsets += offsetMetadataManager.deleteAllOffsets(groupId, records);
-                groupMetadataManager.deleteGroup(groupId, records);
+                groupMetadataManager.createGroupTombstoneRecords(groupId, records);
                 deletedGroups.add(groupId);
 
                 resultCollection.add(
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 0a789fa963089..9068ad17efc8f 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -588,7 +588,9 @@ public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(
      *
      * @param groupId           The group id.
      * @param createIfNotExists A boolean indicating whether the group should be
-     *                          created if it does not exist.
+     *                          created if it does not exist or is an empty classic group.
+     * @param records           The record list to which the group tombstones are written
+     *                          if the group is empty and is a classic group.
      *
      * @return A ConsumerGroup.
      * @throws GroupIdNotFoundException if the group does not exist and createIfNotExists is false or
@@ -598,7 +600,8 @@ public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(
      */
     ConsumerGroup getOrMaybeCreateConsumerGroup(
         String groupId,
-        boolean createIfNotExists
+        boolean createIfNotExists,
+        List<Record> records
     ) throws GroupIdNotFoundException {
         Group group = groups.get(groupId);
 
@@ -606,7 +609,7 @@ ConsumerGroup getOrMaybeCreateConsumerGroup(
             throw new GroupIdNotFoundException(String.format("Consumer group %s not found.", groupId));
         }
 
-        if (group == null) {
+        if (group == null || (createIfNotExists && maybeDeleteEmptyClassicGroup(group, records))) {
             return new ConsumerGroup(snapshotRegistry, groupId, metrics);
         } else {
             if (group.type() == CONSUMER) {
@@ -619,6 +622,40 @@ ConsumerGroup getOrMaybeCreateConsumerGroup(
         }
     }
 
+    /**
+     * Gets a consumer group by committed offset.
+     *
+     * @param groupId           The group id.
+     * @param committedOffset   A specified committed offset corresponding to this shard.
+     *
+     * @return A ConsumerGroup.
+     * @throws GroupIdNotFoundException if the group does not exist or is not a consumer group.
+     */
+    public ConsumerGroup consumerGroup(
+        String groupId,
+        long committedOffset
+    ) throws GroupIdNotFoundException {
+        Group group = group(groupId, committedOffset);
+
+        if (group.type() == CONSUMER) {
+            return (ConsumerGroup) group;
+        } else {
+            // We don't support upgrading/downgrading between protocols at the moment so
+            // we throw an exception if a group exists with the wrong type.
+            throw new GroupIdNotFoundException(String.format("Group %s is not a consumer group.",
+                groupId));
+        }
+    }
+
+    /**
+     * An overloaded method of {@link GroupMetadataManager#consumerGroup(String, long)}
+     */
+    ConsumerGroup consumerGroup(
+        String groupId
+    ) throws GroupIdNotFoundException {
+        return consumerGroup(groupId, Long.MAX_VALUE);
+    }
+
     /**
      * The method should be called on the replay path.
      * Gets or maybe creates a consumer group and updates the groups map if a new group is created.
@@ -723,31 +760,6 @@ public ClassicGroup classicGroup(
         }
     }
 
-    /**
-     * Gets a consumer group by committed offset.
-     *
-     * @param groupId           The group id.
-     * @param committedOffset   A specified committed offset corresponding to this shard.
-     *
-     * @return A ConsumerGroup.
-     * @throws GroupIdNotFoundException if the group does not exist or is not a consumer group.
-     */
-    public ConsumerGroup consumerGroup(
-        String groupId,
-        long committedOffset
-    ) throws GroupIdNotFoundException {
-        Group group = group(groupId, committedOffset);
-
-        if (group.type() == CONSUMER) {
-            return (ConsumerGroup) group;
-        } else {
-            // We don't support upgrading/downgrading between protocols at the moment so
-            // we throw an exception if a group exists with the wrong type.
-            throw new GroupIdNotFoundException(String.format("Group %s is not a consumer group.",
-                groupId));
-        }
-    }
-
     /**
      * Removes the group.
      *
@@ -1056,7 +1068,7 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
 
         // Get or create the consumer group.
         boolean createIfNotExists = memberEpoch == 0;
-        final ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, createIfNotExists);
+        final ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, createIfNotExists, records);
         throwIfConsumerGroupIsFull(group, memberId);
 
         // Get or create the member.
@@ -1324,7 +1336,7 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
         String memberId,
         int memberEpoch
     ) throws ApiException {
-        ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, false);
+        ConsumerGroup group = consumerGroup(groupId);
         List<Record> records;
         if (instanceId == null) {
             ConsumerGroupMember member = group.getOrMaybeCreateMember(memberId, false);
@@ -1449,7 +1461,7 @@ private void scheduleConsumerGroupSessionTimeout(
         String key = consumerGroupSessionTimeoutKey(groupId, memberId);
         timer.schedule(key, consumerGroupSessionTimeoutMs, TimeUnit.MILLISECONDS, true, () -> {
             try {
-                ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, false);
+                ConsumerGroup group = consumerGroup(groupId);
                 ConsumerGroupMember member = group.getOrMaybeCreateMember(memberId, false);
                 log.info("[GroupId {}] Member {} fenced from the group because its session expired.",
                     groupId, memberId);
@@ -1496,7 +1508,7 @@ private void scheduleConsumerGroupRebalanceTimeout(
         String key = consumerGroupRebalanceTimeoutKey(groupId, memberId);
         timer.schedule(key, rebalanceTimeoutMs, TimeUnit.MILLISECONDS, true, () -> {
             try {
-                ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, false);
+                ConsumerGroup group = consumerGroup(groupId);
                 ConsumerGroupMember member = group.getOrMaybeCreateMember(memberId, false);
 
                 if (member.memberEpoch() == memberEpoch) {
@@ -2000,6 +2012,7 @@ public CoordinatorResult<Void, Record> classicGroupJoin(
         CompletableFuture<JoinGroupResponseData> responseFuture
     ) {
         CoordinatorResult<Void, Record> result = EMPTY_RESULT;
+        List<Record> records = new ArrayList<>();
 
         String groupId = request.groupId();
         String memberId = request.memberId();
@@ -2017,6 +2030,7 @@ public CoordinatorResult<Void, Record> classicGroupJoin(
             // Group is created if it does not exist and the member id is UNKNOWN. if member
             // is specified but group does not exist, request is rejected with GROUP_ID_NOT_FOUND
             ClassicGroup group;
+            maybeDeleteEmptyConsumerGroup(groupId, records);
             boolean isNewGroup = !groups.containsKey(groupId);
             try {
                 group = getOrMaybeCreateClassicGroup(groupId, isUnknownMember);
@@ -2065,7 +2079,7 @@ public CoordinatorResult<Void, Record> classicGroupJoin(
                     }
                 });
 
-                List<Record> records = Collections.singletonList(
+                records.add(
                     RecordHelpers.newEmptyGroupMetadataRecord(group, metadataImage.features().metadataVersion())
                 );
 
@@ -3547,12 +3561,25 @@ private void removeCurrentMemberFromClassicGroup(
      * @param groupId The id of the group to be deleted. It has been checked in {@link GroupMetadataManager#validateDeleteGroup}.
      * @param records The record list to populate.
      */
-    public void deleteGroup(
+    public void createGroupTombstoneRecords(
         String groupId,
         List<Record> records
     ) {
         // At this point, we have already validated the group id, so we know that the group exists and that no exception will be thrown.
-        group(groupId).createGroupTombstoneRecords(records);
+        createGroupTombstoneRecords(group(groupId), records);
+    }
+
+    /**
+     * Populates the record list passed in with record to update the state machine.
+     *
+     * @param group The group to be deleted.
+     * @param records The record list to populate.
+     */
+    public void createGroupTombstoneRecords(
+        Group group,
+        List<Record> records
+    ) {
+        group.createGroupTombstoneRecords(records);
     }
 
     /**
@@ -3574,7 +3601,55 @@ void validateDeleteGroup(String groupId) throws ApiException {
     public void maybeDeleteGroup(String groupId, List<Record> records) {
         Group group = groups.get(groupId);
         if (group != null && group.isEmpty()) {
-            deleteGroup(groupId, records);
+            createGroupTombstoneRecords(groupId, records);
+        }
+    }
+
+    /**
+     * @return true if the group is an empty classic group.
+     */
+    private static boolean isEmptyClassicGroup(Group group) {
+        return group != null && group.type() == CLASSIC && group.isEmpty();
+    }
+
+    /**
+     * @return true if the group is an empty consumer group.
+     */
+    private static boolean isEmptyConsumerGroup(Group group) {
+        return group != null && group.type() == CONSUMER && group.isEmpty();
+    }
+
+    /**
+     * Write tombstones for the group if it's empty and is a classic group.
+     *
+     * @param group     The group to be deleted.
+     * @param records   The list of records to delete the group.
+     *
+     * @return true if the group is empty
+     */
+    private boolean maybeDeleteEmptyClassicGroup(Group group, List<Record> records) {
+        if (isEmptyClassicGroup(group)) {
+            // Delete the classic group by adding tombstones.
+            // There's no need to remove the group as the replay of tombstones removes it.
+            if (group != null) createGroupTombstoneRecords(group, records);
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Delete and write tombstones for the group if it's empty and is a consumer group.
+     *
+     * @param groupId The group id to be deleted.
+     * @param records The list of records to delete the group.
+     */
+    private void maybeDeleteEmptyConsumerGroup(String groupId, List<Record> records) {
+        Group group = groups.get(groupId, Long.MAX_VALUE);
+        if (isEmptyConsumerGroup(group)) {
+            // Add tombstones for the previous consumer group. The tombstones won't actually be
+            // replayed because its coordinator result has a non-null appendFuture.
+            createGroupTombstoneRecords(group, records);
+            removeGroup(groupId);
         }
     }
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
index 1e6930a8029d9..59868f36f10bd 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
@@ -225,14 +225,14 @@ public void testDeleteGroups() {
             List<Record> records = invocation.getArgument(1);
             records.add(RecordHelpers.newGroupMetadataTombstoneRecord(groupId));
             return null;
-        }).when(groupMetadataManager).deleteGroup(anyString(), anyList());
+        }).when(groupMetadataManager).createGroupTombstoneRecords(anyString(), anyList());
 
         CoordinatorResult<DeleteGroupsResponseData.DeletableGroupResultCollection, Record> coordinatorResult =
             coordinator.deleteGroups(context, groupIds);
 
         for (String groupId : groupIds) {
             verify(groupMetadataManager, times(1)).validateDeleteGroup(ArgumentMatchers.eq(groupId));
-            verify(groupMetadataManager, times(1)).deleteGroup(ArgumentMatchers.eq(groupId), anyList());
+            verify(groupMetadataManager, times(1)).createGroupTombstoneRecords(ArgumentMatchers.eq(groupId), anyList());
             verify(offsetMetadataManager, times(1)).deleteAllOffsets(ArgumentMatchers.eq(groupId), anyList());
         }
         assertEquals(expectedResult, coordinatorResult);
@@ -291,7 +291,7 @@ public void testDeleteGroupsInvalidGroupId() {
             List<Record> records = invocation.getArgument(1);
             records.add(RecordHelpers.newGroupMetadataTombstoneRecord(groupId));
             return null;
-        }).when(groupMetadataManager).deleteGroup(anyString(), anyList());
+        }).when(groupMetadataManager).createGroupTombstoneRecords(anyString(), anyList());
 
         CoordinatorResult<DeleteGroupsResponseData.DeletableGroupResultCollection, Record> coordinatorResult =
             coordinator.deleteGroups(context, groupIds);
@@ -299,7 +299,7 @@ public void testDeleteGroupsInvalidGroupId() {
         for (String groupId : groupIds) {
             verify(groupMetadataManager, times(1)).validateDeleteGroup(eq(groupId));
             if (!groupId.equals("group-id-2")) {
-                verify(groupMetadataManager, times(1)).deleteGroup(eq(groupId), anyList());
+                verify(groupMetadataManager, times(1)).createGroupTombstoneRecords(eq(groupId), anyList());
                 verify(offsetMetadataManager, times(1)).deleteAllOffsets(eq(groupId), anyList());
             }
         }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index 43703059915a6..dc21a2140d2de 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -434,7 +434,7 @@ public void testMemberJoinsEmptyConsumerGroup() {
         ));
 
         assertThrows(GroupIdNotFoundException.class, () ->
-            context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false));
+            context.groupMetadataManager.consumerGroup(groupId));
 
         CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result = context.consumerGroupHeartbeat(
             new ConsumerGroupHeartbeatRequestData()
@@ -2382,7 +2382,7 @@ public void testSubscriptionMetadataRefreshedAfterGroupIsLoaded() {
 
         // The metadata refresh flag should be true.
         ConsumerGroup consumerGroup = context.groupMetadataManager
-            .getOrMaybeCreateConsumerGroup(groupId, false);
+            .consumerGroup(groupId);
         assertTrue(consumerGroup.hasMetadataExpired(context.time.milliseconds()));
 
         // Prepare the assignment result.
@@ -2493,7 +2493,7 @@ public void testSubscriptionMetadataRefreshedAgainAfterWriteFailure() {
 
         // The metadata refresh flag should be true.
         ConsumerGroup consumerGroup = context.groupMetadataManager
-            .getOrMaybeCreateConsumerGroup(groupId, false);
+            .consumerGroup(groupId);
         assertTrue(consumerGroup.hasMetadataExpired(context.time.milliseconds()));
 
         // Prepare the assignment result.
@@ -2731,7 +2731,7 @@ public void testOnNewMetadataImage() {
 
         // Ensures that all refresh flags are set to the future.
         Arrays.asList("group1", "group2", "group3", "group4", "group5").forEach(groupId -> {
-            ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false);
+            ConsumerGroup group = context.groupMetadataManager.consumerGroup(groupId);
             group.setMetadataRefreshDeadline(context.time.milliseconds() + 5000L, 0);
             assertFalse(group.hasMetadataExpired(context.time.milliseconds()));
         });
@@ -2768,12 +2768,12 @@ public void testOnNewMetadataImage() {
 
         // Verify the groups.
         Arrays.asList("group1", "group2", "group3", "group4").forEach(groupId -> {
-            ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false);
+            ConsumerGroup group = context.groupMetadataManager.consumerGroup(groupId);
             assertTrue(group.hasMetadataExpired(context.time.milliseconds()));
         });
 
         Collections.singletonList("group5").forEach(groupId -> {
-            ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false);
+            ConsumerGroup group = context.groupMetadataManager.consumerGroup(groupId);
             assertFalse(group.hasMetadataExpired(context.time.milliseconds()));
         });
 
@@ -9167,7 +9167,7 @@ public void testClassicGroupDelete() {
 
         List<Record> expectedRecords = Collections.singletonList(RecordHelpers.newGroupMetadataTombstoneRecord("group-id"));
         List<Record> records = new ArrayList<>();
-        context.groupMetadataManager.deleteGroup("group-id", records);
+        context.groupMetadataManager.createGroupTombstoneRecords("group-id", records);
         assertEquals(expectedRecords, records);
     }
 
@@ -9205,7 +9205,7 @@ public void testConsumerGroupDelete() {
             RecordHelpers.newGroupEpochTombstoneRecord(groupId)
         );
         List<Record> records = new ArrayList<>();
-        context.groupMetadataManager.deleteGroup(groupId, records);
+        context.groupMetadataManager.createGroupTombstoneRecords("group-id", records);
         assertEquals(expectedRecords, records);
     }
 
@@ -9380,6 +9380,146 @@ public void testOnConsumerGroupStateTransitionOnLoading() {
         verify(context.metrics, times(1)).onConsumerGroupStateTransition(ConsumerGroup.ConsumerGroupState.EMPTY, null);
     }
 
+    @Test
+    public void testConsumerGroupHeartbeatWithNonEmptyClassicGroup() {
+        String classicGroupId = "classic-group-id";
+        String memberId = Uuid.randomUuid().toString();
+        MockPartitionAssignor assignor = new MockPartitionAssignor("range");
+        assignor.prepareGroupAssignment(new GroupAssignment(Collections.emptyMap()));
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withAssignors(Collections.singletonList(assignor))
+            .build();
+        ClassicGroup classicGroup = new ClassicGroup(
+            new LogContext(),
+            classicGroupId,
+            EMPTY,
+            context.time,
+            context.metrics
+        );
+        context.replay(RecordHelpers.newGroupMetadataRecord(classicGroup, classicGroup.groupAssignment(), MetadataVersion.latestTesting()));
+
+        context.groupMetadataManager.getOrMaybeCreateClassicGroup(classicGroupId, false).transitionTo(PREPARING_REBALANCE);
+        assertThrows(GroupIdNotFoundException.class, () ->
+            context.consumerGroupHeartbeat(
+                new ConsumerGroupHeartbeatRequestData()
+                    .setGroupId(classicGroupId)
+                    .setMemberId(memberId)
+                    .setMemberEpoch(0)
+                    .setServerAssignor("range")
+                    .setRebalanceTimeoutMs(5000)
+                    .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
+                    .setTopicPartitions(Collections.emptyList())));
+    }
+
+    @Test
+    public void testConsumerGroupHeartbeatWithEmptyClassicGroup() {
+        String classicGroupId = "classic-group-id";
+        String memberId = Uuid.randomUuid().toString();
+        MockPartitionAssignor assignor = new MockPartitionAssignor("range");
+        assignor.prepareGroupAssignment(new GroupAssignment(Collections.emptyMap()));
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withAssignors(Collections.singletonList(assignor))
+            .build();
+        ClassicGroup classicGroup = new ClassicGroup(
+            new LogContext(),
+            classicGroupId,
+            EMPTY,
+            context.time,
+            context.metrics
+        );
+        context.replay(RecordHelpers.newGroupMetadataRecord(classicGroup, classicGroup.groupAssignment(), MetadataVersion.latestTesting()));
+
+        CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> result = context.consumerGroupHeartbeat(
+            new ConsumerGroupHeartbeatRequestData()
+                .setGroupId(classicGroupId)
+                .setMemberId(memberId)
+                .setMemberEpoch(0)
+                .setServerAssignor("range")
+                .setRebalanceTimeoutMs(5000)
+                .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
+                .setTopicPartitions(Collections.emptyList()));
+
+        ConsumerGroupMember expectedMember = new ConsumerGroupMember.Builder(memberId)
+            .setState(MemberState.STABLE)
+            .setMemberEpoch(1)
+            .setPreviousMemberEpoch(0)
+            .setRebalanceTimeoutMs(5000)
+            .setClientId("client")
+            .setClientHost("localhost/127.0.0.1")
+            .setSubscribedTopicNames(Arrays.asList("foo", "bar"))
+            .setServerAssignorName("range")
+            .setAssignedPartitions(Collections.emptyMap())
+            .build();
+
+        assertEquals(Errors.NONE.code(), result.response().errorCode());
+        assertEquals(
+            Arrays.asList(
+                RecordHelpers.newGroupMetadataTombstoneRecord(classicGroupId),
+                RecordHelpers.newMemberSubscriptionRecord(classicGroupId, expectedMember),
+                RecordHelpers.newGroupEpochRecord(classicGroupId, 1),
+                RecordHelpers.newTargetAssignmentRecord(classicGroupId, memberId, Collections.emptyMap()),
+                RecordHelpers.newTargetAssignmentEpochRecord(classicGroupId, 1),
+                RecordHelpers.newCurrentAssignmentRecord(classicGroupId, expectedMember)
+            ),
+            result.records()
+        );
+        assertEquals(
+            Group.GroupType.CONSUMER,
+            context.groupMetadataManager.consumerGroup(classicGroupId).type()
+        );
+    }
+
+    @Test
+    public void testClassicGroupJoinWithNonEmptyConsumerGroup() throws Exception {
+        String consumerGroupId = "consumer-group-id";
+        String memberId = Uuid.randomUuid().toString();
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withConsumerGroup(new ConsumerGroupBuilder(consumerGroupId, 10)
+            .withMember(new ConsumerGroupMember.Builder(memberId)
+                .setState(MemberState.STABLE)
+                .setMemberEpoch(10)
+                .setPreviousMemberEpoch(10)
+                .build()))
+            .build();
+
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
+            .withGroupId(consumerGroupId)
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withDefaultProtocolTypeAndProtocols()
+            .build();
+
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
+        assertEquals(Errors.GROUP_ID_NOT_FOUND.code(), joinResult.joinFuture.get().errorCode());
+    }
+
+    @Test
+    public void testClassicGroupJoinWithEmptyConsumerGroup() throws Exception {
+        String consumerGroupId = "consumer-group-id";
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .withConsumerGroup(new ConsumerGroupBuilder(consumerGroupId, 10))
+            .build();
+
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
+            .withGroupId(consumerGroupId)
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withDefaultProtocolTypeAndProtocols()
+            .build();
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request, true);
+
+        List<Record> expectedRecords = Arrays.asList(
+            RecordHelpers.newTargetAssignmentEpochTombstoneRecord(consumerGroupId),
+            RecordHelpers.newGroupSubscriptionMetadataTombstoneRecord(consumerGroupId),
+            RecordHelpers.newGroupEpochTombstoneRecord(consumerGroupId)
+        );
+
+        assertEquals(Errors.MEMBER_ID_REQUIRED.code(), joinResult.joinFuture.get().errorCode());
+        assertEquals(expectedRecords, joinResult.records.subList(0, expectedRecords.size()));
+        assertEquals(
+            Group.GroupType.CLASSIC,
+            context.groupMetadataManager.getOrMaybeCreateClassicGroup(consumerGroupId, false).type()
+        );
+    }
+
     private static void checkJoinGroupResponse(
         JoinGroupResponseData expectedResponse,
         JoinGroupResponseData actualResponse,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
index a02359ffed17c..4a7cd8cae9ed0 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -483,7 +483,7 @@ public ConsumerGroup.ConsumerGroupState consumerGroupState(
         String groupId
     ) {
         return groupMetadataManager
-            .getOrMaybeCreateConsumerGroup(groupId, false)
+            .consumerGroup(groupId)
             .state();
     }
 
@@ -492,7 +492,7 @@ public MemberState consumerGroupMemberState(
         String memberId
     ) {
         return groupMetadataManager
-            .getOrMaybeCreateConsumerGroup(groupId, false)
+            .consumerGroup(groupId)
             .getOrMaybeCreateMember(memberId, false)
             .state();
     }

From f1d741a9c1057eb40a57d6ffbc81edb7e529e1f5 Mon Sep 17 00:00:00 2001
From: Chris Egerton <chrise@aiven.io>
Date: Wed, 20 Mar 2024 15:54:22 +0000
Subject: [PATCH 179/258] KAFKA-16392: Stop emitting warning log message when
 parsing source connector offsets with null partitions (#15562)

Reviewers: Yash Mayya <yash.mayya@gmail.com>
---
 .../org/apache/kafka/connect/storage/OffsetUtils.java  |  6 ++++--
 .../apache/kafka/connect/storage/OffsetUtilsTest.java  | 10 ++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetUtils.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetUtils.java
index 1d6632e60cbfd..6c98f42792123 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetUtils.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetUtils.java
@@ -115,8 +115,10 @@ public static void processPartitionKey(byte[] partitionKey, byte[] offsetValue,
         }
 
         if (!(keyList.get(1) instanceof Map)) {
-            log.warn("Ignoring offset partition key with an unexpected format for the second element in the partition key list. " +
-                    "Expected type: {}, actual type: {}", Map.class.getName(), className(keyList.get(1)));
+            if (keyList.get(1) != null) {
+                log.warn("Ignoring offset partition key with an unexpected format for the second element in the partition key list. " +
+                        "Expected type: {}, actual type: {}", Map.class.getName(), className(keyList.get(1)));
+            }
             return;
         }
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetUtilsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetUtilsTest.java
index 06fb51d3ca61b..ba81a700b440a 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetUtilsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetUtilsTest.java
@@ -134,6 +134,16 @@ public void testProcessPartitionKeyValidList() {
         }
     }
 
+    @Test
+    public void testProcessPartitionKeyNullPartition() {
+        try (LogCaptureAppender logCaptureAppender = LogCaptureAppender.createAndRegister(OffsetUtils.class)) {
+            Map<String, Set<Map<String, Object>>> connectorPartitions = new HashMap<>();
+            OffsetUtils.processPartitionKey(serializePartitionKey(Arrays.asList("connector-name", null)), new byte[0], CONVERTER, connectorPartitions);
+            assertEquals(Collections.emptyMap(), connectorPartitions);
+            assertEquals(0, logCaptureAppender.getMessages().size());
+        }
+    }
+
     private byte[] serializePartitionKey(Object key) {
         return CONVERTER.fromConnectData("", null, key);
     }

From cfa7d393608108be4757a1062226d40e9f8ed29a Mon Sep 17 00:00:00 2001
From: John Yu <54207775+chiacyu@users.noreply.github.com>
Date: Thu, 21 Mar 2024 10:57:59 +0800
Subject: [PATCH 180/258] MINOR : Removed the depreciated information about Zk
 to Kraft migration. (#15552)

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 docs/ops.html | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/docs/ops.html b/docs/ops.html
index 6e440b6f3b510..ce1e0b4109576 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -3799,12 +3799,6 @@ <h4 class="anchor-heading"><a id="kraft_missing" class="anchor-link"></a><a href
 
   <h4 class="anchor-heading"><a id="kraft_zk_migration" class="anchor-link"></a><a href="#kraft_zk_migration">ZooKeeper to KRaft Migration</a></h4>
 
-  <p>
-    <b>ZooKeeper to KRaft migration is considered an Early Access feature and is not recommended for production clusters.</b>
-    Please report issues with ZooKeeper to KRaft migration using the
-    <a href="https://issues.apache.org/jira/projects/KAFKA" target="_blank">project JIRA</a> and the "kraft" component.
-  </p>
-
   <h3>Terminology</h3>
   <ul>
     <li>Brokers that are in <b>ZK mode</b> store their metadata in Apache ZooKepeer. This is the old mode of handling metadata.</li>

From 319a2d15149f3a011e9200fcc03d1862e92fecfc Mon Sep 17 00:00:00 2001
From: Manikumar Reddy <manikumar.reddy@gmail.com>
Date: Thu, 21 Mar 2024 16:02:53 +0530
Subject: [PATCH 181/258] MINOR: Update muckrake version mapping for 3.7
 (#1118)

---
 Jenkinsfile | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Jenkinsfile b/Jenkinsfile
index ef0372294c765..183fe6e7b3204 100755
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -58,6 +58,7 @@ def job = {
             "3.4": "7.4.x",
             "3.5": "7.5.x",
             "3.6": "7.6.x",
+            "3.7": "7.7.x",
             "trunk": "master",
             "master": "master"
     ]

From b7e42d114100d2c61eccbedd9639ecc101056953 Mon Sep 17 00:00:00 2001
From: Manikumar Reddy <manikumar.reddy@gmail.com>
Date: Thu, 21 Mar 2024 16:44:52 +0530
Subject: [PATCH 182/258] MINOR: Bump master to 7.8.0-0 (#1107)

---
 gradle.properties                                               | 2 +-
 streams/quickstart/java/pom.xml                                 | 2 +-
 .../java/src/main/resources/archetype-resources/pom.xml         | 2 +-
 streams/quickstart/pom.xml                                      | 2 +-
 tests/kafkatest/__init__.py                                     | 2 +-
 tests/kafkatest/version.py                                      | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/gradle.properties b/gradle.properties
index 220028256798b..0f3ec02e41758 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -23,7 +23,7 @@ group=org.apache.kafka
 #  - streams/quickstart/pom.xml
 #  - streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
 #  - streams/quickstart/java/pom.xml
-version=7.7.0-0-ccs
+version=7.8.0-0-ccs
 scalaVersion=2.13.12
 # Adding swaggerVersion in gradle.properties to have a single version in place for swagger
 # New version of Swagger 2.2.14 requires minimum JDK 11.
diff --git a/streams/quickstart/java/pom.xml b/streams/quickstart/java/pom.xml
index ab59dc7032ad1..568063999bbd6 100644
--- a/streams/quickstart/java/pom.xml
+++ b/streams/quickstart/java/pom.xml
@@ -26,7 +26,7 @@
     <parent>
         <groupId>org.apache.kafka</groupId>
         <artifactId>streams-quickstart</artifactId>
-        <version>7.7.0-0-ccs</version>
+        <version>7.8.0-0-ccs</version>
         <relativePath>..</relativePath>
     </parent>
 
diff --git a/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml b/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
index 5bd26489a9ede..673a206d6d788 100644
--- a/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
+++ b/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
@@ -29,7 +29,7 @@
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
-        <kafka.version>7.7.0-0-ccs</kafka.version>
+        <kafka.version>7.8.0-0-ccs</kafka.version>
         <slf4j.version>1.7.36</slf4j.version>
     </properties>
 
diff --git a/streams/quickstart/pom.xml b/streams/quickstart/pom.xml
index 6a7b665a7f547..2d2c10edc6666 100644
--- a/streams/quickstart/pom.xml
+++ b/streams/quickstart/pom.xml
@@ -22,7 +22,7 @@
     <groupId>org.apache.kafka</groupId>
     <artifactId>streams-quickstart</artifactId>
     <packaging>pom</packaging>
-    <version>7.7.0-0-ccs</version>
+    <version>7.8.0-0-ccs</version>
 
     <name>Kafka Streams :: Quickstart</name>
 
diff --git a/tests/kafkatest/__init__.py b/tests/kafkatest/__init__.py
index 40f6ba829d048..0ea89a9810c24 100644
--- a/tests/kafkatest/__init__.py
+++ b/tests/kafkatest/__init__.py
@@ -22,4 +22,4 @@
 # Instead, in development branches, the version should have a suffix of the form ".devN"
 #
 # For example, when Kafka is at version 1.0.0-0, this should be something like "1.0.0-0.dev0"
-__version__ = '7.7.0-0.dev0'
+__version__ = '7.8.0-0.dev0'
diff --git a/tests/kafkatest/version.py b/tests/kafkatest/version.py
index b5ece3a3fbb13..da89ab882543e 100644
--- a/tests/kafkatest/version.py
+++ b/tests/kafkatest/version.py
@@ -122,7 +122,7 @@ def get_version(node=None):
         return DEV_BRANCH
 
 DEV_BRANCH = KafkaVersion("dev")
-DEV_VERSION = KafkaVersion("7.7.0-0")
+DEV_VERSION = KafkaVersion("7.8.0-0-ccs")
 
 LATEST_METADATA_VERSION = "3.8"
 

From a41c10fd49841381b5207c184a385622094ed440 Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Thu, 21 Mar 2024 19:20:05 +0800
Subject: [PATCH 183/258] KAFKA-16318 : add javafoc for kafka metric (#15483)

Add the javadoc for KafkaMetric

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/common/metrics/KafkaMetric.java     | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
index 6d395fa3fbcd0..805ae2a52781b 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
@@ -29,6 +29,14 @@ public final class KafkaMetric implements Metric {
     private MetricConfig config;
 
     // public for testing
+    /**
+     * Create a metric to monitor an object that implements MetricValueProvider.
+     * @param lock The lock used to prevent race condition
+     * @param metricName The name of the metric
+     * @param valueProvider The metric value provider associated with this metric
+     * @param config The configuration of the metric
+     * @param time The time instance to use with the metrics
+     */
     public KafkaMetric(Object lock, MetricName metricName, MetricValueProvider<?> valueProvider,
             MetricConfig config, Time time) {
         this.metricName = metricName;
@@ -40,15 +48,29 @@ public KafkaMetric(Object lock, MetricName metricName, MetricValueProvider<?> va
         this.time = time;
     }
 
+    /**
+     * Get the configuration of this metric.
+     * This is supposed to be used by server only.
+     * @return Return the config of this metric
+     */
     public MetricConfig config() {
         return this.config;
     }
 
+    /**
+     * Get the metric name
+     * @return Return the name of this metric
+     */
     @Override
     public MetricName metricName() {
         return this.metricName;
     }
 
+    /**
+     * Take the metric and return the value, which could be a {@link Measurable} or a {@link Gauge}
+     * @return Return the metric value
+     * @throws IllegalStateException if the underlying metric is not a {@link Measurable} or a {@link Gauge}.
+     */
     @Override
     public Object metricValue() {
         long now = time.milliseconds();
@@ -62,6 +84,11 @@ else if (this.metricValueProvider instanceof Gauge)
         }
     }
 
+    /**
+     * Get the underlying metric provider, which should be a {@link Measurable}
+     * @return Return the metric provider
+     * @throws IllegalStateException if the underlying metric is not a {@link Measurable}.
+     */
     public Measurable measurable() {
         if (this.metricValueProvider instanceof Measurable)
             return (Measurable) metricValueProvider;
@@ -69,6 +96,11 @@ public Measurable measurable() {
             throw new IllegalStateException("Not a measurable: " + this.metricValueProvider.getClass());
     }
 
+    /**
+     * Take the metric and return the value, where the underlying metric provider should be a {@link Measurable}
+     * @param timeMs The time that this metric is taken
+     * @return Return the metric value if it's measurable, otherwise 0
+     */
     double measurableValue(long timeMs) {
         synchronized (this.lock) {
             if (this.metricValueProvider instanceof Measurable)
@@ -78,6 +110,11 @@ public Measurable measurable() {
         }
     }
 
+    /**
+     * Set the metric config.
+     * This is supposed to be used by server only.
+     * @param config configuration for this metrics
+     */
     public void config(MetricConfig config) {
         synchronized (lock) {
             this.config = config;

From 03f7b5aa3abe9ee3b0474508a7797e0b992ec1f7 Mon Sep 17 00:00:00 2001
From: Alyssa Huang <ahuang@confluent.io>
Date: Thu, 21 Mar 2024 07:38:42 -0700
Subject: [PATCH 184/258] KAFKA-16206: Fix unnecessary topic config deletion
 during ZK migration (#14206)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Ron Dagostino <rndgstn@gmail.com>
---
 .../migration/ZkConfigMigrationClient.scala   |  4 +-
 .../kafka/image/ConfigurationsDelta.java      | 11 ++---
 .../migration/KRaftMigrationZkWriter.java     |  5 ---
 .../kafka/image/ConfigurationsImageTest.java  | 23 +++++++----
 .../apache/kafka/image/TopicsImageTest.java   | 41 ++++++++++++++-----
 .../migration/KRaftMigrationDriverTest.java   | 27 ++++++------
 6 files changed, 69 insertions(+), 42 deletions(-)

diff --git a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
index 846a599875bb7..dffdc138aa138 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
@@ -228,8 +228,8 @@ class ZkConfigMigrationClient(
       val (migrationZkVersion, responses) = zkClient.retryMigrationRequestsUntilConnected(requests, state)
 
       if (responses.head.resultCode.equals(Code.NONODE)) {
-        // Not fatal.
-        error(s"Did not delete $configResource since the node did not exist.")
+        // Not fatal. This is expected in the case this is a topic config and we delete the topic
+        debug(s"Did not delete $configResource since the node did not exist.")
         state
       } else if (responses.head.resultCode.equals(Code.OK)) {
         // Write the notification znode if our update was successful
diff --git a/metadata/src/main/java/org/apache/kafka/image/ConfigurationsDelta.java b/metadata/src/main/java/org/apache/kafka/image/ConfigurationsDelta.java
index eab0505e332f3..9d794bfddc210 100644
--- a/metadata/src/main/java/org/apache/kafka/image/ConfigurationsDelta.java
+++ b/metadata/src/main/java/org/apache/kafka/image/ConfigurationsDelta.java
@@ -71,11 +71,12 @@ public void replay(ConfigRecord record) {
     public void replay(RemoveTopicRecord record, String topicName) {
         ConfigResource resource =
             new ConfigResource(Type.TOPIC, topicName);
-        ConfigurationImage configImage = image.resourceData().getOrDefault(resource,
-                new ConfigurationImage(resource, Collections.emptyMap()));
-        ConfigurationDelta delta = changes.computeIfAbsent(resource,
-            __ -> new ConfigurationDelta(configImage));
-        delta.deleteAll();
+        if (image.resourceData().containsKey(resource)) {
+            ConfigurationImage configImage = image.resourceData().get(resource);
+            ConfigurationDelta delta = changes.computeIfAbsent(resource,
+                __ -> new ConfigurationDelta(configImage));
+            delta.deleteAll();
+        }
     }
 
     public ConfigurationsImage apply() {
diff --git a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
index 8a7e148d579e0..6870ad8f0cbd2 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/migration/KRaftMigrationZkWriter.java
@@ -250,11 +250,6 @@ public void visitPartition(TopicIdPartition topicIdPartition, PartitionRegistrat
                 migrationState -> migrationClient.topicClient().deleteTopic(topicName, migrationState)
             );
             ConfigResource resource = new ConfigResource(ConfigResource.Type.TOPIC, topicName);
-            operationConsumer.accept(
-                UPDATE_TOPIC_CONFIG,
-                "Updating Configs for Topic " + topicName + ", ID " + topicId,
-                migrationState -> migrationClient.configClient().deleteConfigs(resource, migrationState)
-            );
         });
 
         newPartitions.forEach((topicId, partitionMap) -> {
diff --git a/metadata/src/test/java/org/apache/kafka/image/ConfigurationsImageTest.java b/metadata/src/test/java/org/apache/kafka/image/ConfigurationsImageTest.java
index 9b7cd39dcd6e5..995a44b3f8f84 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ConfigurationsImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ConfigurationsImageTest.java
@@ -27,6 +27,7 @@
 import org.junit.jupiter.api.Timeout;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -61,26 +62,34 @@ public class ConfigurationsImageTest {
         IMAGE1 = new ConfigurationsImage(map1);
 
         DELTA1_RECORDS = new ArrayList<>();
+        // remove configs
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new ConfigRecord().setResourceType(BROKER.id()).
             setResourceName("0").setName("foo").setValue(null),
             CONFIG_RECORD.highestSupportedVersion()));
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ConfigRecord().setResourceType(BROKER.id()).
+            setResourceName("0").setName("baz").setValue(null),
+            CONFIG_RECORD.highestSupportedVersion()));
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ConfigRecord().setResourceType(BROKER.id()).
+            setResourceName("1").setName("foobar").setValue(null),
+            CONFIG_RECORD.highestSupportedVersion()));
+        // add new config to b1
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new ConfigRecord().setResourceType(BROKER.id()).
             setResourceName("1").setName("barfoo").setValue("bazfoo"),
             CONFIG_RECORD.highestSupportedVersion()));
+        // add new config to b2
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ConfigRecord().setResourceType(BROKER.id()).
+            setResourceName("2").setName("foo").setValue("bar"),
+            CONFIG_RECORD.highestSupportedVersion()));
 
         DELTA1 = new ConfigurationsDelta(IMAGE1);
         RecordTestUtils.replayAll(DELTA1, DELTA1_RECORDS);
 
         Map<ConfigResource, ConfigurationImage> map2 = new HashMap<>();
-        Map<String, String> broker0Map2 = new HashMap<>();
-        broker0Map2.put("baz", "quux");
-        map2.put(new ConfigResource(BROKER, "0"),
-            new ConfigurationImage(new ConfigResource(BROKER, "0"), broker0Map2));
-        Map<String, String> broker1Map2 = new HashMap<>();
-        broker1Map2.put("foobar", "foobaz");
-        broker1Map2.put("barfoo", "bazfoo");
+        Map<String, String> broker1Map2 = Collections.singletonMap("barfoo", "bazfoo");
         map2.put(new ConfigResource(BROKER, "1"),
             new ConfigurationImage(new ConfigResource(BROKER, "1"), broker1Map2));
+        Map<String, String> broker2Map = Collections.singletonMap("foo", "bar");
+        map2.put(new ConfigResource(BROKER, "2"), new ConfigurationImage(new ConfigResource(BROKER, "2"), broker2Map));
         IMAGE2 = new ConfigurationsImage(map2);
     }
 
diff --git a/metadata/src/test/java/org/apache/kafka/image/TopicsImageTest.java b/metadata/src/test/java/org/apache/kafka/image/TopicsImageTest.java
index 84a092ccd3932..eabb63ff858e2 100644
--- a/metadata/src/test/java/org/apache/kafka/image/TopicsImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/TopicsImageTest.java
@@ -94,10 +94,16 @@ private static ImmutableMap<String, TopicImage> newTopicsByNameMap(Collection<To
 
     public static final Uuid FOO_UUID = Uuid.fromString("ThIaNwRnSM2Nt9Mx1v0RvA");
 
+    private static final Uuid FOO_UUID2 = Uuid.fromString("9d3lha5qv8DoIl93jf8pbX");
+
     private static final Uuid BAR_UUID = Uuid.fromString("f62ptyETTjet8SL5ZeREiw");
 
     private static final Uuid BAZ_UUID = Uuid.fromString("tgHBnRglT5W_RlENnuG5vg");
 
+    private static final Uuid BAM_UUID = Uuid.fromString("b66ybsWIQoygs01vdjH07A");
+
+    private static final Uuid BAM_UUID2 = Uuid.fromString("yd6Sq3a9aK1G8snlKv7ag5");
+
     static {
         TOPIC_IMAGES1 = Arrays.asList(
             newTopicImage("foo", FOO_UUID,
@@ -118,16 +124,20 @@ private static ImmutableMap<String, TopicImage> newTopicsByNameMap(Collection<To
         IMAGE1 = new TopicsImage(newTopicsByIdMap(TOPIC_IMAGES1), newTopicsByNameMap(TOPIC_IMAGES1));
 
         DELTA1_RECORDS = new ArrayList<>();
+        // remove topic
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new RemoveTopicRecord().
             setTopicId(FOO_UUID),
             REMOVE_TOPIC_RECORD.highestSupportedVersion()));
+        // change topic
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new PartitionChangeRecord().
             setTopicId(BAR_UUID).
             setPartitionId(0).setLeader(1),
             PARTITION_CHANGE_RECORD.highestSupportedVersion()));
+        // add topic
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new TopicRecord().
             setName("baz").setTopicId(BAZ_UUID),
             TOPIC_RECORD.highestSupportedVersion()));
+        // add partition record for new topic
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new PartitionRecord().
             setPartitionId(0).
             setTopicId(BAZ_UUID).
@@ -138,11 +148,23 @@ private static ImmutableMap<String, TopicImage> newTopicsByNameMap(Collection<To
             setLeader(3).
             setLeaderEpoch(2).
             setPartitionEpoch(1), PARTITION_RECORD.highestSupportedVersion()));
+        // re-add topic with different topic id
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new TopicRecord().
+            setName("foo").setTopicId(FOO_UUID2),
+            TOPIC_RECORD.highestSupportedVersion()));
+        // add then remove topic
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new TopicRecord().
+            setName("bam").setTopicId(BAM_UUID),
+            TOPIC_RECORD.highestSupportedVersion()));
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new RemoveTopicRecord().
+            setTopicId(BAM_UUID),
+            REMOVE_TOPIC_RECORD.highestSupportedVersion()));
 
         DELTA1 = new TopicsDelta(IMAGE1);
         RecordTestUtils.replayAll(DELTA1, DELTA1_RECORDS);
 
         List<TopicImage> topics2 = Arrays.asList(
+            newTopicImage("foo", FOO_UUID2),
             newTopicImage("bar", BAR_UUID,
                 new PartitionRegistration.Builder().setReplicas(new int[] {0, 1, 2, 3, 4}).
                     setDirectories(DirectoryId.migratingArray(5)).
@@ -188,22 +210,22 @@ private PartitionRegistration newPartition(int[] replicas) {
     public void testBasicLocalChanges() {
         int localId = 3;
         /* Changes already include in DELTA1_RECORDS and IMAGE1:
-         * foo - topic id deleted
+         * foo - topic id deleted then recreated with different topic id
          * bar-0 - stay as follower with different partition epoch
          * baz-0 - new topic to leader
+         * bam - topic id created then deleted
          */
         List<ApiMessageAndVersion> topicRecords = new ArrayList<>(DELTA1_RECORDS);
 
-        // Create a new foo topic with a different id
-        Uuid newFooId = Uuid.fromString("b66ybsWIQoygs01vdjH07A");
+        // Create a new bam topic with a different id
         topicRecords.add(
             new ApiMessageAndVersion(
-                new TopicRecord().setName("foo") .setTopicId(newFooId),
+                new TopicRecord().setName("bam").setTopicId(BAM_UUID2),
                 TOPIC_RECORD.highestSupportedVersion()
             )
         );
-        topicRecords.add(newPartitionRecord(newFooId, 0, Arrays.asList(0, 1, 2)));
-        topicRecords.add(newPartitionRecord(newFooId, 1, Arrays.asList(0, 1, localId)));
+        topicRecords.add(newPartitionRecord(BAM_UUID2, 0, Arrays.asList(0, 1, 2)));
+        topicRecords.add(newPartitionRecord(BAM_UUID2, 1, Arrays.asList(0, 1, localId)));
 
         // baz-1 - new partition to follower
         topicRecords.add(
@@ -224,10 +246,6 @@ public void testBasicLocalChanges() {
         RecordTestUtils.replayAll(delta, topicRecords);
 
         LocalReplicaChanges changes = delta.localChanges(localId);
-        assertEquals(
-            new HashSet<>(Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("foo", 1))),
-            changes.deletes()
-        );
         assertEquals(
             new HashSet<>(Arrays.asList(new TopicPartition("baz", 0))),
             changes.electedLeaders().keySet()
@@ -238,7 +256,8 @@ public void testBasicLocalChanges() {
         );
         assertEquals(
             new HashSet<>(
-                Arrays.asList(new TopicPartition("baz", 1), new TopicPartition("bar", 0), new TopicPartition("foo", 1))
+                Arrays.asList(new TopicPartition("baz", 1), new TopicPartition("bar", 0),
+                    new TopicPartition("bam", 1))
             ),
             changes.followers().keySet()
         );
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationDriverTest.java b/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationDriverTest.java
index 67df561209634..dea5e62db969d 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationDriverTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/migration/KRaftMigrationDriverTest.java
@@ -567,7 +567,7 @@ public void testTopicDualWriteSnapshot() throws Exception {
             TestUtils.waitForCondition(() -> driver.migrationState().get(1, TimeUnit.MINUTES).equals(MigrationDriverState.DUAL_WRITE),
                 "Waiting for KRaftMigrationDriver to enter ZK_MIGRATION state");
 
-            // Modify topics in a KRaft snapshot -- delete foo, modify bar, add baz
+            // Modify topics in a KRaft snapshot -- delete foo, modify bar, add baz, add new foo, add bam, delete bam
             provenance = new MetadataProvenance(200, 1, 1);
             delta = new MetadataDelta(image);
             RecordTestUtils.replayAll(delta, DELTA1_RECORDS);
@@ -577,10 +577,11 @@ public void testTopicDualWriteSnapshot() throws Exception {
 
             assertEquals(1, topicClient.deletedTopics.size());
             assertEquals("foo", topicClient.deletedTopics.get(0));
-            assertEquals(1, topicClient.createdTopics.size());
-            assertEquals("baz", topicClient.createdTopics.get(0));
+            assertEquals(2, topicClient.createdTopics.size());
+            assertTrue(topicClient.createdTopics.contains("foo"));
+            assertTrue(topicClient.createdTopics.contains("baz"));
             assertTrue(topicClient.updatedTopicPartitions.get("bar").contains(0));
-            assertEquals(new ConfigResource(ConfigResource.Type.TOPIC, "foo"), configClient.deletedResources.get(0));
+            assertEquals(0, configClient.deletedResources.size());
         });
     }
 
@@ -621,7 +622,7 @@ public void testTopicDualWriteDelta() throws Exception {
             TestUtils.waitForCondition(() -> driver.migrationState().get(1, TimeUnit.MINUTES).equals(MigrationDriverState.DUAL_WRITE),
                     "Waiting for KRaftMigrationDriver to enter DUAL_WRITE state");
 
-            // Modify topics in a KRaft snapshot -- delete foo, modify bar, add baz
+            // Modify topics in a KRaft snapshot -- delete foo, modify bar, add baz, add new foo, add bam, delete bam
             provenance = new MetadataProvenance(200, 1, 1);
             delta = new MetadataDelta(image);
             RecordTestUtils.replayAll(delta, DELTA1_RECORDS);
@@ -631,10 +632,11 @@ public void testTopicDualWriteDelta() throws Exception {
 
             assertEquals(1, topicClient.deletedTopics.size());
             assertEquals("foo", topicClient.deletedTopics.get(0));
-            assertEquals(1, topicClient.createdTopics.size());
-            assertEquals("baz", topicClient.createdTopics.get(0));
+            assertEquals(2, topicClient.createdTopics.size());
+            assertTrue(topicClient.createdTopics.contains("foo"));
+            assertTrue(topicClient.createdTopics.contains("baz"));
             assertTrue(topicClient.updatedTopicPartitions.get("bar").contains(0));
-            assertEquals(new ConfigResource(ConfigResource.Type.TOPIC, "foo"), configClient.deletedResources.get(0));
+            assertEquals(0, configClient.deletedResources.size());
         });
     }
 
@@ -730,7 +732,7 @@ public void testControllerFailover() throws Exception {
             migrationClient.setMigrationRecoveryState(
                 ZkMigrationLeadershipState.EMPTY.withKRaftMetadataOffsetAndEpoch(100, 1));
 
-            // Modify topics in a KRaft -- delete foo, modify bar, add baz
+            // Modify topics in a KRaft -- delete foo, modify bar, add baz, add new foo, add bam, delete bam
             provenance = new MetadataProvenance(200, 1, 1);
             delta = new MetadataDelta(image);
             RecordTestUtils.replayAll(delta, DELTA1_RECORDS);
@@ -746,10 +748,11 @@ public void testControllerFailover() throws Exception {
                 "");
             assertEquals(1, topicClient.deletedTopics.size());
             assertEquals("foo", topicClient.deletedTopics.get(0));
-            assertEquals(1, topicClient.createdTopics.size());
-            assertEquals("baz", topicClient.createdTopics.get(0));
+            assertEquals(2, topicClient.createdTopics.size());
+            assertTrue(topicClient.createdTopics.contains("foo"));
+            assertTrue(topicClient.createdTopics.contains("baz"));
             assertTrue(topicClient.updatedTopicPartitions.get("bar").contains(0));
-            assertEquals(new ConfigResource(ConfigResource.Type.TOPIC, "foo"), configClient.deletedResources.get(0));
+            assertEquals(0, configClient.deletedResources.size());
         });
     }
 

From c5804e7649265df29ce0b85333043383fc8779aa Mon Sep 17 00:00:00 2001
From: Philip Nee <pnee@confluent.io>
Date: Thu, 21 Mar 2024 08:23:51 -0700
Subject: [PATCH 185/258] KAFKA-16273: Update consumer_bench_test.py to use
 consumer group protocol (#15548)

Adding this as part of the greater effort to modify the system tests to incorporate the use of consumer group protocol from KIP-848. Following is the test results and the tests using protocol = consumer are expected to fail:

================================================================================
SESSION REPORT (ALL TESTS)
ducktape version: 0.11.4
session_id:       2024-03-16--002
run time:         76 minutes 36.150 seconds
tests run:        28
passed:           25
flaky:            0
failed:           3
ignored:          0
================================================================================

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>, Kirk True <ktrue@confluent.io>
---
 .../tests/core/consume_bench_test.py          | 72 ++++++++++++++-----
 1 file changed, 53 insertions(+), 19 deletions(-)

diff --git a/tests/kafkatest/tests/core/consume_bench_test.py b/tests/kafkatest/tests/core/consume_bench_test.py
index 07633c5533072..0575b5d8f2172 100644
--- a/tests/kafkatest/tests/core/consume_bench_test.py
+++ b/tests/kafkatest/tests/core/consume_bench_test.py
@@ -17,7 +17,7 @@
 from ducktape.mark import matrix
 from ducktape.mark.resource import cluster
 from ducktape.tests.test import Test
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.trogdor.produce_bench_workload import ProduceBenchWorkloadService, ProduceBenchWorkloadSpec
 from kafkatest.services.trogdor.consume_bench_workload import ConsumeBenchWorkloadService, ConsumeBenchWorkloadSpec
 from kafkatest.services.trogdor.task_spec import TaskSpec
@@ -82,9 +82,18 @@ def produce_messages(self, topics, max_messages=10000):
             ["consume_bench_topic[0-5]:[0-4]"] # manual topic assignment
         ],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        topics=[
+            ["consume_bench_topic[0-5]"], # topic subscription
+            ["consume_bench_topic[0-5]:[0-4]"] # manual topic assignment
+        ],
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_consume_bench(self, topics, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_consume_bench(self, topics, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Runs a ConsumeBench workload to consume messages
         """
@@ -94,7 +103,7 @@ def test_consume_bench(self, topics, metadata_quorum=quorum.zk, use_new_coordina
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=10000,
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 active_topics=topics)
@@ -111,9 +120,14 @@ def test_consume_bench(self, topics, metadata_quorum=quorum.zk, use_new_coordina
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
     )
-    def test_single_partition(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
+    )
+    def test_single_partition(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Run a ConsumeBench against a single partition
         """
@@ -124,7 +138,7 @@ def test_single_partition(self, metadata_quorum=quorum.zk, use_new_coordinator=F
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=2500,
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 active_topics=["consume_bench_topic:1"])
@@ -141,9 +155,14 @@ def test_single_partition(self, metadata_quorum=quorum.zk, use_new_coordinator=F
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_multiple_consumers_random_group_topics(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_multiple_consumers_random_group_topics(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Runs multiple consumers group to read messages from topics.
         Since a consumerGroup isn't specified, each consumer should read from all topics independently
@@ -154,7 +173,7 @@ def test_multiple_consumers_random_group_topics(self, metadata_quorum=quorum.zk,
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=5000, # all should read exactly 5k messages
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 threads_per_worker=5,
@@ -172,9 +191,14 @@ def test_multiple_consumers_random_group_topics(self, metadata_quorum=quorum.zk,
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
     )
-    def test_two_consumers_specified_group_topics(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
+    )
+    def test_two_consumers_specified_group_topics(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Runs two consumers in the same consumer group to read messages from topics.
         Since a consumerGroup is specified, each consumer should dynamically get assigned a partition from group
@@ -185,7 +209,7 @@ def test_two_consumers_specified_group_topics(self, metadata_quorum=quorum.zk, u
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=2000, # both should read at least 2k messages
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 threads_per_worker=2,
@@ -204,9 +228,14 @@ def test_two_consumers_specified_group_topics(self, metadata_quorum=quorum.zk, u
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_multiple_consumers_random_group_partitions(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_multiple_consumers_random_group_partitions(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Runs multiple consumers in to read messages from specific partitions.
         Since a consumerGroup isn't specified, each consumer will get assigned a random group
@@ -218,7 +247,7 @@ def test_multiple_consumers_random_group_partitions(self, metadata_quorum=quorum
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=2000,
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 threads_per_worker=4,
@@ -236,9 +265,14 @@ def test_multiple_consumers_random_group_partitions(self, metadata_quorum=quorum
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        metadata_quorum=[quorum.isolated_kraft],
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_multiple_consumers_specified_group_partitions_should_raise(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_multiple_consumers_specified_group_partitions_should_raise(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Runs multiple consumers in the same group to read messages from specific partitions.
         It is an invalid configuration to provide a consumer group and specific partitions.
@@ -250,7 +284,7 @@ def test_multiple_consumers_specified_group_partitions_should_raise(self, metada
                                                 self.consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=1000,
                                                 max_messages=2000,
-                                                consumer_conf={},
+                                                consumer_conf=consumer_group.maybe_set_group_protocol(group_protocol),
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 threads_per_worker=4,

From 7e85d7d32e8cb4f3fa9b01fd4197ee43d64ba6d0 Mon Sep 17 00:00:00 2001
From: Alyssa Huang <ahuang@confluent.io>
Date: Thu, 21 Mar 2024 11:06:40 -0700
Subject: [PATCH 186/258] MINOR: KRaft upgrade tests should only use latest
 stable mv (#15566)

This should help us avoid testing MVs before they are usable (stable).
We revert back from testing 3.8 in this case since 3.7 is the current stable version.

Reviewers: Proven Provenzano <pprovenzano@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 tests/kafkatest/tests/core/kraft_upgrade_test.py | 6 +++---
 tests/kafkatest/version.py                       | 3 ++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/kafkatest/tests/core/kraft_upgrade_test.py b/tests/kafkatest/tests/core/kraft_upgrade_test.py
index f6a291869568e..3f3c4a81b103a 100644
--- a/tests/kafkatest/tests/core/kraft_upgrade_test.py
+++ b/tests/kafkatest/tests/core/kraft_upgrade_test.py
@@ -23,7 +23,7 @@
 from kafkatest.tests.produce_consume_validate import ProduceConsumeValidateTest
 from kafkatest.utils import is_int
 from kafkatest.version import LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, \
-    DEV_BRANCH, KafkaVersion, LATEST_METADATA_VERSION
+    DEV_BRANCH, KafkaVersion, LATEST_STABLE_METADATA_VERSION
 
 #
 # Test upgrading between different KRaft versions.
@@ -71,8 +71,8 @@ def perform_version_change(self, from_kafka_version):
             self.kafka.start_node(node)
             self.wait_until_rejoin()
             self.logger.info("Successfully restarted broker node %s" % node.account.hostname)
-        self.logger.info("Changing metadata.version to %s" % LATEST_METADATA_VERSION)
-        self.kafka.upgrade_metadata_version(LATEST_METADATA_VERSION)
+        self.logger.info("Changing metadata.version to %s" % LATEST_STABLE_METADATA_VERSION)
+        self.kafka.upgrade_metadata_version(LATEST_STABLE_METADATA_VERSION)
 
     def run_upgrade(self, from_kafka_version):
         """Test upgrade of Kafka broker cluster from various versions to the current version
diff --git a/tests/kafkatest/version.py b/tests/kafkatest/version.py
index 3e3da017bc4a0..9f21c120fce14 100644
--- a/tests/kafkatest/version.py
+++ b/tests/kafkatest/version.py
@@ -124,7 +124,8 @@ def get_version(node=None):
 DEV_BRANCH = KafkaVersion("dev")
 DEV_VERSION = KafkaVersion("3.8.0-SNAPSHOT")
 
-LATEST_METADATA_VERSION = "3.8"
+# This should match the LATEST_PRODUCTION version defined in MetadataVersion.java
+LATEST_STABLE_METADATA_VERSION = "3.7"
 
 # 0.8.2.x versions
 V_0_8_2_1 = KafkaVersion("0.8.2.1")

From 997ca14f8057462e52ea00721c9cc16469ac5108 Mon Sep 17 00:00:00 2001
From: Christo Lolov <lolovc@amazon.com>
Date: Fri, 22 Mar 2024 09:43:53 +0000
Subject: [PATCH 187/258] KAFKA-14133: Move stateDirectory mock in
 TaskManagerTest to Mockito (#15254)

This pull requests migrates the StateDirectory mock in TaskManagerTest from EasyMock to Mockito.
The change is restricted to a single mock to minimize the scope and make it easier for review.

Reviewers: Ismael Juma <ismael@juma.me.uk>, Bruno Cadonna <cadonna@apache.org>
---
 .../processor/internals/TaskManagerTest.java  | 1070 ++++++++---------
 1 file changed, 517 insertions(+), 553 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index 36d2a3e378633..64ad0d1caf96e 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -58,9 +58,6 @@
 import java.time.Duration;
 import java.util.ArrayList;
 
-import org.easymock.EasyMockRunner;
-import org.easymock.Mock;
-import org.easymock.MockType;
 import org.hamcrest.Matchers;
 import org.junit.Before;
 import org.junit.Rule;
@@ -69,10 +66,10 @@
 import org.junit.runner.RunWith;
 import org.mockito.Answers;
 import org.mockito.InOrder;
-import org.mockito.Mockito;
+import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoJUnitRunner;
 import org.mockito.junit.MockitoRule;
-import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
 
 import java.io.File;
@@ -107,11 +104,6 @@
 import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask;
 import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask;
 import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask;
-import static org.easymock.EasyMock.eq;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.expectLastCall;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.verify;
 import static org.hamcrest.CoreMatchers.hasItem;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.empty;
@@ -125,20 +117,25 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.mock;
 
-@MockitoSettings(strictness = Strictness.STRICT_STUBS)
-@RunWith(EasyMockRunner.class)
+@RunWith(MockitoJUnitRunner.StrictStubs.class)
 public class TaskManagerTest {
 
     private final String topic1 = "topic1";
@@ -188,26 +185,26 @@ public class TaskManagerTest {
 
     final java.util.function.Consumer<Set<TopicPartition>> noOpResetter = partitions -> { };
 
-    @org.mockito.Mock
+    @Mock
     private InternalTopologyBuilder topologyBuilder;
-    @Mock(type = MockType.DEFAULT)
+    @Mock
     private StateDirectory stateDirectory;
-    @org.mockito.Mock
+    @Mock
     private ChangelogReader changeLogReader;
-    @org.mockito.Mock
+    @Mock
     private Consumer<byte[], byte[]> consumer;
-    @org.mockito.Mock
+    @Mock
     private ActiveTaskCreator activeTaskCreator;
-    @org.mockito.Mock
+    @Mock
     private StandbyTaskCreator standbyTaskCreator;
-    @org.mockito.Mock
+    @Mock
     private Admin adminClient;
-    @org.mockito.Mock
+    @Mock
     private ProcessorStateManager stateManager;
-    @org.mockito.Mock(answer = Answers.RETURNS_DEEP_STUBS)
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private ProcessorStateManager.StateStoreMetadata stateStore;
-    final StateUpdater stateUpdater = Mockito.mock(StateUpdater.class);
-    final DefaultTaskManager schedulingTaskManager = Mockito.mock(DefaultTaskManager.class);
+    final StateUpdater stateUpdater = mock(StateUpdater.class);
+    final DefaultTaskManager schedulingTaskManager = mock(DefaultTaskManager.class);
 
     private TaskManager taskManager;
     private TopologyMetadata topologyMetadata;
@@ -267,7 +264,7 @@ public void shouldClassifyExistingTasksWithoutStateUpdater() {
 
         taskManager.handleAssignment(activeTasks, standbyTasks);
 
-        Mockito.verifyNoInteractions(stateUpdater);
+        verifyNoInteractions(stateUpdater);
     }
 
     @Test
@@ -276,7 +273,7 @@ public void shouldNotUpdateExistingStandbyTaskIfStandbyIsReassignedWithSameInput
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
         updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(standbyTask, taskId03Partitions);
-        Mockito.verify(standbyTask, never()).updateInputPartitions(Mockito.eq(taskId03Partitions), Mockito.any());
+        verify(standbyTask, never()).updateInputPartitions(eq(taskId03Partitions), any());
     }
 
     @Test
@@ -285,12 +282,12 @@ public void shouldUpdateExistingStandbyTaskIfStandbyIsReassignedWithDifferentInp
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
         updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(standbyTask, taskId04Partitions);
-        Mockito.verify(standbyTask).updateInputPartitions(Mockito.eq(taskId04Partitions), Mockito.any());
+        verify(standbyTask).updateInputPartitions(eq(taskId04Partitions), any());
     }
 
     private void updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(final Task standbyTask,
                                                                                    final Set<TopicPartition> newInputPartition) {
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         when(tasks.allTasks()).thenReturn(mkSet(standbyTask));
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, false);
 
@@ -299,7 +296,7 @@ private void updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(f
             mkMap(mkEntry(standbyTask.id(), newInputPartition))
         );
 
-        Mockito.verify(standbyTask).resume();
+        verify(standbyTask).resume();
     }
 
     @Test
@@ -307,7 +304,7 @@ public void shouldLockAllTasksOnCorruptionWithProcessingThreads() {
         final StreamTask activeTask1 = statefulTask(taskId00, taskId00ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId00Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
         when(tasks.activeTaskIds()).thenReturn(mkSet(taskId00, taskId01));
         when(tasks.task(taskId00)).thenReturn(activeTask1);
@@ -316,9 +313,9 @@ public void shouldLockAllTasksOnCorruptionWithProcessingThreads() {
 
         taskManager.handleCorruption(mkSet(taskId00));
 
-        Mockito.verify(consumer).assignment();
-        Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
-        Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
+        verify(consumer).assignment();
+        verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
     }
 
     @Test
@@ -329,20 +326,20 @@ public void shouldLockCommitableTasksOnCorruptionWithProcessingThreads() {
         final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
         final KafkaFuture<Void> mockFuture = KafkaFuture.completedFuture(null);
         when(schedulingTaskManager.lockTasks(any())).thenReturn(mockFuture);
 
         taskManager.commit(mkSet(activeTask1, activeTask2));
 
-        Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
-        Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
     }
 
     @Test
     public void shouldLockActiveOnHandleAssignmentWithProcessingThreads() {
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
         when(tasks.allTaskIds()).thenReturn(mkSet(taskId00, taskId01));
         final KafkaFuture<Void> mockFuture = KafkaFuture.completedFuture(null);
@@ -353,8 +350,8 @@ public void shouldLockActiveOnHandleAssignmentWithProcessingThreads() {
                 mkMap(mkEntry(taskId01, taskId01Partitions))
         );
 
-        Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
-        Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
     }
 
     @Test
@@ -365,7 +362,7 @@ public void shouldLockAffectedTasksOnHandleRevocation() {
         final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTask1, activeTask2));
         final KafkaFuture<Void> mockFuture = KafkaFuture.completedFuture(null);
@@ -373,8 +370,8 @@ public void shouldLockAffectedTasksOnHandleRevocation() {
 
         taskManager.handleRevocation(taskId01Partitions);
 
-        Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
-        Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01));
+        verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01));
     }
 
     @Test
@@ -385,7 +382,7 @@ public void shouldLockTasksOnClose() {
         final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTask1, activeTask2));
         final KafkaFuture<Void> mockFuture = KafkaFuture.completedFuture(null);
@@ -393,8 +390,8 @@ public void shouldLockTasksOnClose() {
 
         taskManager.closeAndCleanUpTasks(mkSet(activeTask1), mkSet(), false);
 
-        Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00));
-        Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00));
+        verify(schedulingTaskManager).lockTasks(mkSet(taskId00));
+        verify(schedulingTaskManager).unlockTasks(mkSet(taskId00));
     }
 
     @Test
@@ -405,14 +402,14 @@ public void shouldResumePollingForPartitionsWithAvailableSpaceForAllActiveTasks(
         final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.activeTasks()).thenReturn(mkSet(activeTask1, activeTask2));
 
         taskManager.resumePollingForPartitionsWithAvailableSpace();
 
-        Mockito.verify(activeTask1).resumePollingForPartitionsWithAvailableSpace();
-        Mockito.verify(activeTask2).resumePollingForPartitionsWithAvailableSpace();
+        verify(activeTask1).resumePollingForPartitionsWithAvailableSpace();
+        verify(activeTask2).resumePollingForPartitionsWithAvailableSpace();
     }
 
     @Test
@@ -423,14 +420,14 @@ public void shouldUpdateLagForAllActiveTasks() {
         final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.activeTasks()).thenReturn(mkSet(activeTask1, activeTask2));
 
         taskManager.updateLags();
 
-        Mockito.verify(activeTask1).updateLags();
-        Mockito.verify(activeTask2).updateLags();
+        verify(activeTask1).updateLags();
+        verify(activeTask2).updateLags();
     }
 
     @Test
@@ -438,7 +435,7 @@ public void shouldPrepareActiveTaskInStateUpdaterToBeRecycled() {
         final StreamTask activeTaskToRecycle = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToRecycle));
 
@@ -447,9 +444,9 @@ public void shouldPrepareActiveTaskInStateUpdaterToBeRecycled() {
             mkMap(mkEntry(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions()))
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(tasks).addPendingTaskToRecycle(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(tasks).addPendingTaskToRecycle(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -457,7 +454,7 @@ public void shouldPrepareStandbyTaskInStateUpdaterToBeRecycled() {
         final StandbyTask standbyTaskToRecycle = standbyTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToRecycle));
 
@@ -466,10 +463,10 @@ public void shouldPrepareStandbyTaskInStateUpdaterToBeRecycled() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater).remove(standbyTaskToRecycle.id());
-        Mockito.verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater).remove(standbyTaskToRecycle.id());
+        verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -477,16 +474,16 @@ public void shouldRemoveUnusedActiveTaskFromStateUpdater() {
         final StreamTask activeTaskToClose = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToClose));
 
         taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap());
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater).remove(activeTaskToClose.id());
-        Mockito.verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater).remove(activeTaskToClose.id());
+        verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -494,16 +491,16 @@ public void shouldRemoveUnusedStandbyTaskFromStateUpdater() {
         final StandbyTask standbyTaskToClose = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToClose));
 
         taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap());
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater).remove(standbyTaskToClose.id());
-        Mockito.verify(tasks).addPendingTaskToCloseClean(standbyTaskToClose.id());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater).remove(standbyTaskToClose.id());
+        verify(tasks).addPendingTaskToCloseClean(standbyTaskToClose.id());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -512,7 +509,7 @@ public void shouldUpdateInputPartitionOfActiveTaskInStateUpdater() {
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId02Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToUpdateInputPartitions));
 
@@ -521,10 +518,10 @@ public void shouldUpdateInputPartitionOfActiveTaskInStateUpdater() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater).remove(activeTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks).addPendingTaskToUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater).remove(activeTaskToUpdateInputPartitions.id());
+        verify(tasks).addPendingTaskToUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions);
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -533,7 +530,7 @@ public void shouldCloseReviveAndUpdateInputPartitionOfActiveTaskInStateUpdater()
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId02Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToUpdateInputPartitions));
         when(tasks.removePendingTaskToCloseClean(activeTaskToUpdateInputPartitions.id())).thenReturn(true);
@@ -543,11 +540,11 @@ public void shouldCloseReviveAndUpdateInputPartitionOfActiveTaskInStateUpdater()
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater, never()).remove(activeTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks).removePendingTaskToCloseClean(activeTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks).addPendingTaskToCloseReviveAndUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater, never()).remove(activeTaskToUpdateInputPartitions.id());
+        verify(tasks).removePendingTaskToCloseClean(activeTaskToUpdateInputPartitions.id());
+        verify(tasks).addPendingTaskToCloseReviveAndUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions);
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -555,7 +552,7 @@ public void shouldKeepReassignedActiveTaskInStateUpdater() {
         final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedActiveTask));
 
@@ -564,8 +561,8 @@ public void shouldKeepReassignedActiveTaskInStateUpdater() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -573,7 +570,7 @@ public void shouldMoveReassignedSuspendedActiveTaskToStateUpdater() {
         final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.SUSPENDED)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(reassignedActiveTask));
 
@@ -582,10 +579,10 @@ public void shouldMoveReassignedSuspendedActiveTaskToStateUpdater() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(tasks).removeTask(reassignedActiveTask);
-        Mockito.verify(stateUpdater).add(reassignedActiveTask);
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(tasks).removeTask(reassignedActiveTask);
+        verify(stateUpdater).add(reassignedActiveTask);
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -593,7 +590,7 @@ public void shouldRemoveReassignedRevokedActiveTaskInStateUpdaterFromPendingTask
         final StreamTask reassignedRevokedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedRevokedActiveTask));
 
@@ -602,9 +599,9 @@ public void shouldRemoveReassignedRevokedActiveTaskInStateUpdaterFromPendingTask
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(tasks).removePendingActiveTaskToSuspend(reassignedRevokedActiveTask.id());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(tasks).removePendingActiveTaskToSuspend(reassignedRevokedActiveTask.id());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -612,7 +609,7 @@ public void shouldRemoveReassignedLostTaskInStateUpdaterFromPendingTaskToCloseCl
         final StreamTask reassignedLostTask = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedLostTask));
         when(tasks.removePendingTaskToCloseClean(reassignedLostTask.id())).thenReturn(true);
@@ -622,10 +619,10 @@ public void shouldRemoveReassignedLostTaskInStateUpdaterFromPendingTaskToCloseCl
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(tasks).removePendingTaskToCloseClean(reassignedLostTask.id());
-        Mockito.verify(tasks).addPendingTaskToAddBack(reassignedLostTask.id());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(tasks).removePendingTaskToCloseClean(reassignedLostTask.id());
+        verify(tasks).addPendingTaskToAddBack(reassignedLostTask.id());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -633,7 +630,7 @@ public void shouldRemoveReassignedTaskInStateUpdaterFromPendingSuspend() {
         final StreamTask reassignedTask = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RESTORING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedTask));
         when(tasks.removePendingActiveTaskToSuspend(reassignedTask.id())).thenReturn(true);
@@ -643,10 +640,10 @@ public void shouldRemoveReassignedTaskInStateUpdaterFromPendingSuspend() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(tasks).removePendingActiveTaskToSuspend(reassignedTask.id());
-        Mockito.verify(tasks).addPendingTaskToAddBack(reassignedTask.id());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(tasks).removePendingActiveTaskToSuspend(reassignedTask.id());
+        verify(tasks).addPendingTaskToAddBack(reassignedTask.id());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -654,7 +651,7 @@ public void shouldReAddStandbyTaskFromPendingRecycle() {
         final StandbyTask reassignedStandbyTask = standbyTask(taskId01, taskId01ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId01Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedStandbyTask));
         when(tasks.removePendingTaskToRecycle(reassignedStandbyTask.id())).thenReturn(taskId01Partitions);
@@ -664,10 +661,10 @@ public void shouldReAddStandbyTaskFromPendingRecycle() {
             mkMap(mkEntry(reassignedStandbyTask.id(), reassignedStandbyTask.inputPartitions()))
         );
 
-        Mockito.verify(tasks).removePendingTaskToRecycle(reassignedStandbyTask.id());
-        Mockito.verify(tasks).addPendingTaskToAddBack(reassignedStandbyTask.id());
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(tasks).removePendingTaskToRecycle(reassignedStandbyTask.id());
+        verify(tasks).addPendingTaskToAddBack(reassignedStandbyTask.id());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -676,7 +673,7 @@ public void shouldNeverUpdateInputPartitionsOfStandbyTaskInStateUpdater() {
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId03Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToUpdateInputPartitions));
 
@@ -685,11 +682,11 @@ public void shouldNeverUpdateInputPartitionsOfStandbyTaskInStateUpdater() {
             mkMap(mkEntry(standbyTaskToUpdateInputPartitions.id(), newInputPartitions))
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks, never())
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id());
+        verify(tasks, never())
             .addPendingTaskToUpdateInputPartitions(standbyTaskToUpdateInputPartitions.id(), newInputPartitions);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -698,7 +695,7 @@ public void shouldNeverCloseReviveAndUpdateInputPartitionsOfStandbyTaskInStateUp
                 .inState(State.RUNNING)
                 .withInputPartitions(taskId02Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId03Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToUpdateInputPartitions));
 
@@ -707,12 +704,12 @@ public void shouldNeverCloseReviveAndUpdateInputPartitionsOfStandbyTaskInStateUp
                 mkMap(mkEntry(standbyTaskToUpdateInputPartitions.id(), newInputPartitions))
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks, never()).removePendingTaskToCloseClean(standbyTaskToUpdateInputPartitions.id());
-        Mockito.verify(tasks, never())
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id());
+        verify(tasks, never()).removePendingTaskToCloseClean(standbyTaskToUpdateInputPartitions.id());
+        verify(tasks, never())
                 .addPendingTaskToCloseReviveAndUpdateInputPartitions(standbyTaskToUpdateInputPartitions.id(), newInputPartitions);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -720,7 +717,7 @@ public void shouldKeepReassignedStandbyTaskInStateUpdater() {
         final StandbyTask reassignedStandbyTask = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedStandbyTask));
 
@@ -729,8 +726,8 @@ public void shouldKeepReassignedStandbyTaskInStateUpdater() {
             mkMap(mkEntry(reassignedStandbyTask.id(), reassignedStandbyTask.inputPartitions()))
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -741,7 +738,7 @@ public void shouldAssignMultipleTasksInStateUpdater() {
         final StandbyTask standbyTaskToRecycle = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToClose, standbyTaskToRecycle));
 
@@ -750,12 +747,12 @@ public void shouldAssignMultipleTasksInStateUpdater() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(stateUpdater).remove(activeTaskToClose.id());
-        Mockito.verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id());
-        Mockito.verify(stateUpdater).remove(standbyTaskToRecycle.id());
-        Mockito.verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(stateUpdater).remove(activeTaskToClose.id());
+        verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id());
+        verify(stateUpdater).remove(standbyTaskToRecycle.id());
+        verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -766,7 +763,7 @@ public void shouldReturnStateUpdaterTasksInAllTasks() {
         final StandbyTask standbyTask = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
 
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask));
@@ -782,10 +779,9 @@ public void shouldNotReturnStateUpdaterTasksInOwnedTasks() {
         final StandbyTask standbyTask = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
 
-        when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask));
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId03, activeTask)));
         assertEquals(taskManager.allOwnedTasks(), mkMap(mkEntry(taskId03, activeTask)));
     }
@@ -795,7 +791,7 @@ public void shouldCreateActiveTaskDuringAssignment() {
         final StreamTask activeTaskToBeCreated = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.CREATED)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         final Set<Task> createdTasks = mkSet(activeTaskToBeCreated);
         final Map<TaskId, Set<TopicPartition>> tasksToBeCreated = mkMap(
@@ -804,8 +800,8 @@ public void shouldCreateActiveTaskDuringAssignment() {
 
         taskManager.handleAssignment(tasksToBeCreated, Collections.emptyMap());
 
-        Mockito.verify(tasks).addPendingTasksToInit(createdTasks);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(tasks).addPendingTasksToInit(createdTasks);
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -813,7 +809,7 @@ public void shouldCreateStandbyTaskDuringAssignment() {
         final StandbyTask standbyTaskToBeCreated = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.CREATED)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         final Set<Task> createdTasks = mkSet(standbyTaskToBeCreated);
         when(standbyTaskCreator.createTasks(mkMap(
@@ -825,8 +821,8 @@ public void shouldCreateStandbyTaskDuringAssignment() {
             mkMap(mkEntry(standbyTaskToBeCreated.id(), standbyTaskToBeCreated.inputPartitions()))
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(tasks).addPendingTasksToInit(createdTasks);
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(tasks).addPendingTasksToInit(createdTasks);
     }
 
     @Test
@@ -845,12 +841,12 @@ public void shouldAddRecycledStandbyTaskfromActiveToPendingTasksToInitWithStateU
 
         taskManager.handleAssignment(emptyMap(), mkMap(mkEntry(taskId01, taskId01Partitions)));
 
-        Mockito.verify(activeTaskToRecycle).prepareCommit();
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToRecycle.id());
-        Mockito.verify(tasks).addPendingTasksToInit(mkSet(standbyTask));
-        Mockito.verify(tasks).removeTask(activeTaskToRecycle);
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskToRecycle).prepareCommit();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToRecycle.id());
+        verify(tasks).addPendingTasksToInit(mkSet(standbyTask));
+        verify(tasks).removeTask(activeTaskToRecycle);
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -869,11 +865,11 @@ public void shouldAddRecycledStandbyTaskfromActiveToTaskRegistryWithStateUpdater
 
         taskManager.handleAssignment(emptyMap(), mkMap(mkEntry(taskId01, taskId01Partitions)));
 
-        Mockito.verify(activeTaskToRecycle).prepareCommit();
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToRecycle.id());
-        Mockito.verify(tasks).replaceActiveWithStandby(standbyTask);
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskToRecycle).prepareCommit();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToRecycle.id());
+        verify(tasks).replaceActiveWithStandby(standbyTask);
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -881,7 +877,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToRecycleIsFoundInTasksRegis
         final StandbyTask standbyTaskToRecycle = standbyTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToRecycle));
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
 
@@ -895,7 +891,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToRecycleIsFoundInTasksRegis
 
         assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater, " +
             "but standby task " + taskId03 + " is managed by the stream thread");
-        Mockito.verifyNoInteractions(activeTaskCreator);
+        verifyNoInteractions(activeTaskCreator);
     }
 
     @Test
@@ -903,18 +899,18 @@ public void shouldAssignActiveTaskInTasksRegistryToBeClosedCleanlyWithStateUpdat
         final StreamTask activeTaskToClose = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTaskToClose));
 
         taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap());
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToClose.id());
-        Mockito.verify(activeTaskToClose).prepareCommit();
-        Mockito.verify(activeTaskToClose).closeClean();
-        Mockito.verify(tasks).removeTask(activeTaskToClose);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(activeTaskToClose.id());
+        verify(activeTaskToClose).prepareCommit();
+        verify(activeTaskToClose).closeClean();
+        verify(tasks).removeTask(activeTaskToClose);
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -922,7 +918,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToCloseIsFoundInTasksRegistr
         final StandbyTask standbyTaskToClose = standbyTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToClose));
 
@@ -933,7 +929,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToCloseIsFoundInTasksRegistr
 
         assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater, " +
             "but standby task " + taskId03 + " is managed by the stream thread");
-        Mockito.verifyNoInteractions(activeTaskCreator);
+        verifyNoInteractions(activeTaskCreator);
     }
 
     @Test
@@ -942,7 +938,7 @@ public void shouldAssignActiveTaskInTasksRegistryToUpdateInputPartitionsWithStat
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId02Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTaskToUpdateInputPartitions));
         when(tasks.updateActiveTaskInputPartitions(activeTaskToUpdateInputPartitions, newInputPartitions)).thenReturn(true);
@@ -952,9 +948,9 @@ public void shouldAssignActiveTaskInTasksRegistryToUpdateInputPartitionsWithStat
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(activeTaskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(newInputPartitions), any());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(activeTaskToUpdateInputPartitions).updateInputPartitions(eq(newInputPartitions), any());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -962,7 +958,7 @@ public void shouldResumeActiveRunningTaskInTasksRegistryWithStateUpdaterEnabled(
         final StreamTask activeTaskToResume = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTaskToResume));
 
@@ -971,8 +967,8 @@ public void shouldResumeActiveRunningTaskInTasksRegistryWithStateUpdaterEnabled(
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -980,7 +976,7 @@ public void shouldResumeActiveSuspendedTaskInTasksRegistryAndAddToStateUpdater()
         final StreamTask activeTaskToResume = statefulTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.SUSPENDED)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTaskToResume));
 
@@ -989,11 +985,11 @@ public void shouldResumeActiveSuspendedTaskInTasksRegistryAndAddToStateUpdater()
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
-        Mockito.verify(activeTaskToResume).resume();
-        Mockito.verify(stateUpdater).add(activeTaskToResume);
-        Mockito.verify(tasks).removeTask(activeTaskToResume);
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap());
+        verify(activeTaskToResume).resume();
+        verify(stateUpdater).add(activeTaskToResume);
+        verify(tasks).removeTask(activeTaskToResume);
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -1002,7 +998,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToUpdateInputPartitionsIsFou
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
         final Set<TopicPartition> newInputPartitions = taskId03Partitions;
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToUpdateInputPartitions));
 
@@ -1016,7 +1012,7 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToUpdateInputPartitionsIsFou
 
         assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater, " +
             "but standby task " + taskId02 + " is managed by the stream thread");
-        Mockito.verifyNoInteractions(activeTaskCreator);
+        verifyNoInteractions(activeTaskCreator);
     }
 
     @Test
@@ -1027,7 +1023,7 @@ public void shouldAssignMultipleTasksInTasksRegistryWithStateUpdaterEnabled() {
         final StreamTask activeTaskToCreate = statefulTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.CREATED)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(activeTaskToClose));
 
@@ -1036,12 +1032,12 @@ public void shouldAssignMultipleTasksInTasksRegistryWithStateUpdaterEnabled() {
             Collections.emptyMap()
         );
 
-        Mockito.verify(activeTaskCreator).createTasks(
+        verify(activeTaskCreator).createTasks(
             consumer,
             mkMap(mkEntry(activeTaskToCreate.id(), activeTaskToCreate.inputPartitions()))
         );
-        Mockito.verify(activeTaskToClose).closeClean();
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verify(activeTaskToClose).closeClean();
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
     }
 
     @Test
@@ -1058,10 +1054,10 @@ public void shouldAddTasksToStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(task00).initializeIfNeeded();
-        Mockito.verify(task01).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(task00);
-        Mockito.verify(stateUpdater).add(task01);
+        verify(task00).initializeIfNeeded();
+        verify(task01).initializeIfNeeded();
+        verify(stateUpdater).add(task00);
+        verify(stateUpdater).add(task01);
     }
 
     @Test
@@ -1080,13 +1076,13 @@ public void shouldRetryInitializationWhenLockExceptionInStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(task00).initializeIfNeeded();
-        Mockito.verify(task01).initializeIfNeeded();
-        Mockito.verify(tasks).addPendingTasksToInit(
-            Mockito.argThat(tasksToInit -> tasksToInit.contains(task00) && !tasksToInit.contains(task01))
+        verify(task00).initializeIfNeeded();
+        verify(task01).initializeIfNeeded();
+        verify(tasks).addPendingTasksToInit(
+            argThat(tasksToInit -> tasksToInit.contains(task00) && !tasksToInit.contains(task01))
         );
-        Mockito.verify(stateUpdater, never()).add(task00);
-        Mockito.verify(stateUpdater).add(task01);
+        verify(stateUpdater, never()).add(task00);
+        verify(stateUpdater).add(task01);
     }
 
     @Test
@@ -1117,13 +1113,13 @@ public void shouldRetryInitializationWhenLockExceptionAfterRecyclingInStateUpdat
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(task00Converted).initializeIfNeeded();
-        Mockito.verify(task01Converted).initializeIfNeeded();
-        Mockito.verify(tasks).addPendingTasksToInit(
-            Mockito.argThat(tasksToInit -> tasksToInit.contains(task00Converted) && !tasksToInit.contains(task01Converted))
+        verify(task00Converted).initializeIfNeeded();
+        verify(task01Converted).initializeIfNeeded();
+        verify(tasks).addPendingTasksToInit(
+            argThat(tasksToInit -> tasksToInit.contains(task00Converted) && !tasksToInit.contains(task01Converted))
         );
-        Mockito.verify(stateUpdater, never()).add(task00Converted);
-        Mockito.verify(stateUpdater).add(task01Converted);
+        verify(stateUpdater, never()).add(task00Converted);
+        verify(stateUpdater).add(task01Converted);
     }
 
     @Test
@@ -1151,13 +1147,13 @@ public void shouldRecycleTasksRemovedFromStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any());
-        Mockito.verify(task00).suspend();
-        Mockito.verify(task01).suspend();
-        Mockito.verify(task00Converted).initializeIfNeeded();
-        Mockito.verify(task01Converted).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(task00Converted);
-        Mockito.verify(stateUpdater).add(task01Converted);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(task00).suspend();
+        verify(task01).suspend();
+        verify(task00Converted).initializeIfNeeded();
+        verify(task01Converted).initializeIfNeeded();
+        verify(stateUpdater).add(task00Converted);
+        verify(stateUpdater).add(task01Converted);
     }
 
     @Test
@@ -1178,11 +1174,11 @@ public void shouldCloseTasksRemovedFromStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any());
-        Mockito.verify(task00).suspend();
-        Mockito.verify(task00).closeClean();
-        Mockito.verify(task01).suspend();
-        Mockito.verify(task01).closeClean();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(task00).suspend();
+        verify(task00).closeClean();
+        verify(task01).suspend();
+        verify(task01).closeClean();
     }
 
     @Test
@@ -1204,14 +1200,14 @@ public void shouldUpdateInputPartitionsOfTasksRemovedFromStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTask).updateInputPartitions(Mockito.eq(taskId02Partitions), anyMap());
-        Mockito.verify(activeTask, never()).closeDirty();
-        Mockito.verify(activeTask, never()).closeClean();
-        Mockito.verify(stateUpdater).add(activeTask);
-        Mockito.verify(standbyTask).updateInputPartitions(Mockito.eq(taskId03Partitions), anyMap());
-        Mockito.verify(standbyTask, never()).closeDirty();
-        Mockito.verify(standbyTask, never()).closeClean();
-        Mockito.verify(stateUpdater).add(standbyTask);
+        verify(activeTask).updateInputPartitions(eq(taskId02Partitions), anyMap());
+        verify(activeTask, never()).closeDirty();
+        verify(activeTask, never()).closeClean();
+        verify(stateUpdater).add(activeTask);
+        verify(standbyTask).updateInputPartitions(eq(taskId03Partitions), anyMap());
+        verify(standbyTask, never()).closeDirty();
+        verify(standbyTask, never()).closeClean();
+        verify(stateUpdater).add(standbyTask);
     }
 
     @Test
@@ -1228,12 +1224,12 @@ public void shouldCloseReviveAndUpdateInputPartitionsOfTasksRemovedFromStateUpda
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTask).closeClean();
-        Mockito.verify(activeTask).revive();
-        Mockito.verify(activeTask).updateInputPartitions(Mockito.eq(taskId02Partitions), anyMap());
-        Mockito.verify(activeTask).initializeIfNeeded();
-        Mockito.verify(activeTask, never()).closeDirty();
-        Mockito.verify(stateUpdater).add(activeTask);
+        verify(activeTask).closeClean();
+        verify(activeTask).revive();
+        verify(activeTask).updateInputPartitions(eq(taskId02Partitions), anyMap());
+        verify(activeTask).initializeIfNeeded();
+        verify(activeTask, never()).closeDirty();
+        verify(stateUpdater).add(activeTask);
     }
 
     @Test
@@ -1252,9 +1248,9 @@ public void shouldSuspendRevokedTaskRemovedFromStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(statefulTask).suspend();
-        Mockito.verify(tasks).addTask(statefulTask);
-        Mockito.verifyNoInteractions(consumer);
+        verify(statefulTask).suspend();
+        verify(tasks).addTask(statefulTask);
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1301,21 +1297,21 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter -> { });
 
-        Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
-        Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
-        Mockito.verify(convertedTask0).initializeIfNeeded();
-        Mockito.verify(convertedTask1).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(convertedTask0);
-        Mockito.verify(stateUpdater).add(convertedTask1);
-        Mockito.verify(taskToClose).closeClean();
-        Mockito.verify(taskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId04Partitions), anyMap());
-        Mockito.verify(stateUpdater).add(taskToUpdateInputPartitions);
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).closeClean();
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).revive();
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId05Partitions), anyMap());
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions);
-        Mockito.verifyNoInteractions(consumer);
+        verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(convertedTask0).initializeIfNeeded();
+        verify(convertedTask1).initializeIfNeeded();
+        verify(stateUpdater).add(convertedTask0);
+        verify(stateUpdater).add(convertedTask1);
+        verify(taskToClose).closeClean();
+        verify(taskToUpdateInputPartitions).updateInputPartitions(eq(taskId04Partitions), anyMap());
+        verify(stateUpdater).add(taskToUpdateInputPartitions);
+        verify(taskToCloseReviveAndUpdateInputPartitions).closeClean();
+        verify(taskToCloseReviveAndUpdateInputPartitions).revive();
+        verify(taskToCloseReviveAndUpdateInputPartitions).updateInputPartitions(eq(taskId05Partitions), anyMap());
+        verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded();
+        verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions);
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1364,8 +1360,8 @@ public void shouldAddActiveTaskWithRevokedInputPartitionsInStateUpdaterToPending
 
         taskManager.handleRevocation(task.inputPartitions());
 
-        Mockito.verify(tasks).addPendingActiveTaskToSuspend(task.id());
-        Mockito.verify(stateUpdater, never()).remove(task.id());
+        verify(tasks).addPendingActiveTaskToSuspend(task.id());
+        verify(stateUpdater, never()).remove(task.id());
     }
 
     public void shouldAddMultipleActiveTasksWithRevokedInputPartitionsInStateUpdaterToPendingTasksToSuspend() {
@@ -1380,8 +1376,8 @@ public void shouldAddMultipleActiveTasksWithRevokedInputPartitionsInStateUpdater
 
         taskManager.handleRevocation(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
-        Mockito.verify(tasks).addPendingActiveTaskToSuspend(task1.id());
-        Mockito.verify(tasks).addPendingActiveTaskToSuspend(task2.id());
+        verify(tasks).addPendingActiveTaskToSuspend(task1.id());
+        verify(tasks).addPendingActiveTaskToSuspend(task2.id());
     }
 
     @Test
@@ -1394,8 +1390,8 @@ public void shouldNotAddActiveTaskWithoutRevokedInputPartitionsInStateUpdaterToP
 
         taskManager.handleRevocation(taskId01Partitions);
 
-        Mockito.verify(stateUpdater, never()).remove(task.id());
-        Mockito.verify(tasks, never()).addPendingActiveTaskToSuspend(task.id());
+        verify(stateUpdater, never()).remove(task.id());
+        verify(tasks, never()).addPendingActiveTaskToSuspend(task.id());
     }
 
     @Test
@@ -1408,8 +1404,8 @@ public void shouldNotRevokeStandbyTaskInStateUpdaterOnRevocation() {
 
         taskManager.handleRevocation(taskId00Partitions);
 
-        Mockito.verify(stateUpdater, never()).remove(task.id());
-        Mockito.verify(tasks, never()).addPendingActiveTaskToSuspend(task.id());
+        verify(stateUpdater, never()).remove(task.id());
+        verify(tasks, never()).addPendingActiveTaskToSuspend(task.id());
     }
 
     @Test
@@ -1428,12 +1424,12 @@ public void shouldRemoveAllActiveTasksFromStateUpdaterOnPartitionLost() {
 
         taskManager.handleLostAll();
 
-        Mockito.verify(stateUpdater).remove(task1.id());
-        Mockito.verify(stateUpdater, never()).remove(task2.id());
-        Mockito.verify(stateUpdater).remove(task3.id());
-        Mockito.verify(tasks).addPendingTaskToCloseClean(task1.id());
-        Mockito.verify(tasks, never()).addPendingTaskToCloseClean(task2.id());
-        Mockito.verify(tasks).addPendingTaskToCloseClean(task3.id());
+        verify(stateUpdater).remove(task1.id());
+        verify(stateUpdater, never()).remove(task2.id());
+        verify(stateUpdater).remove(task3.id());
+        verify(tasks).addPendingTaskToCloseClean(task1.id());
+        verify(tasks, never()).addPendingTaskToCloseClean(task2.id());
+        verify(tasks).addPendingTaskToCloseClean(task3.id());
     }
 
     private TaskManager setupForRevocationAndLost(final Set<Task> tasksInStateUpdater,
@@ -1454,10 +1450,10 @@ public void shouldTransitRestoredTaskToRunning() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(task).completeRestoration(noOpResetter);
-        Mockito.verify(task).clearTaskTimeout();
-        Mockito.verify(tasks).addTask(task);
-        Mockito.verify(consumer).resume(task.inputPartitions());
+        verify(task).completeRestoration(noOpResetter);
+        verify(task).clearTaskTimeout();
+        verify(tasks).addTask(task);
+        verify(consumer).resume(task.inputPartitions());
     }
 
     @Test
@@ -1472,10 +1468,10 @@ public void shouldHandleTimeoutExceptionInTransitRestoredTaskToRunning() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(task).maybeInitTaskTimeoutOrThrow(anyLong(), Mockito.eq(timeoutException));
-        Mockito.verify(tasks, never()).addTask(task);
-        Mockito.verify(task, never()).clearTaskTimeout();
-        Mockito.verifyNoInteractions(consumer);
+        verify(task).maybeInitTaskTimeoutOrThrow(anyLong(), eq(timeoutException));
+        verify(tasks, never()).addTask(task);
+        verify(task, never()).clearTaskTimeout();
+        verifyNoInteractions(consumer);
     }
 
     private TaskManager setUpTransitionToRunningOfRestoredTask(final StreamTask statefulTask,
@@ -1514,10 +1510,10 @@ public void shouldRecycleRestoredTask() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
-        Mockito.verify(statefulTask).suspend();
-        Mockito.verify(standbyTask).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(standbyTask);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
+        verify(statefulTask).suspend();
+        verify(standbyTask).initializeIfNeeded();
+        verify(stateUpdater).add(standbyTask);
     }
 
     @Test
@@ -1534,8 +1530,8 @@ public void shouldHandleExceptionThrownDuringConversionInRecycleRestoredTask() {
             () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
         );
 
-        Mockito.verify(stateUpdater, never()).add(any());
-        Mockito.verify(statefulTask).closeDirty();
+        verify(stateUpdater, never()).add(any());
+        verify(statefulTask).closeDirty();
     }
 
     @Test
@@ -1556,8 +1552,8 @@ public void shouldHandleExceptionThrownDuringTaskInitInRecycleRestoredTask() {
             () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
         );
 
-        Mockito.verify(stateUpdater, never()).add(any());
-        Mockito.verify(standbyTask).closeDirty();
+        verify(stateUpdater, never()).add(any());
+        verify(standbyTask).closeDirty();
     }
 
     private TaskManager setUpRecycleRestoredTask(final StreamTask statefulTask) {
@@ -1579,11 +1575,11 @@ public void shouldCloseCleanRestoredTask() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
-        Mockito.verify(statefulTask).suspend();
-        Mockito.verify(statefulTask).closeClean();
-        Mockito.verify(statefulTask, never()).closeDirty();
-        Mockito.verify(tasks, never()).removeTask(statefulTask);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
+        verify(statefulTask).suspend();
+        verify(statefulTask).closeClean();
+        verify(statefulTask, never()).closeDirty();
+        verify(tasks, never()).removeTask(statefulTask);
     }
 
     @Test
@@ -1600,9 +1596,9 @@ public void shouldHandleExceptionThrownDuringCloseInCloseCleanRestoredTask() {
             () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
         );
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
-        Mockito.verify(statefulTask).closeDirty();
-        Mockito.verify(tasks, never()).removeTask(statefulTask);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask.id());
+        verify(statefulTask).closeDirty();
+        verify(tasks, never()).removeTask(statefulTask);
     }
 
     @Test
@@ -1621,8 +1617,8 @@ public void shouldHandleExceptionThrownDuringClosingTaskProducerInCloseCleanRest
             () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
         );
 
-        Mockito.verify(statefulTask, never()).closeDirty();
-        Mockito.verify(tasks, never()).removeTask(statefulTask);
+        verify(statefulTask, never()).closeDirty();
+        verify(tasks, never()).removeTask(statefulTask);
     }
 
     private TaskManager setUpCloseCleanRestoredTask(final StreamTask statefulTask,
@@ -1649,8 +1645,8 @@ public void shouldAddBackRestoredTask() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(stateUpdater).add(statefulTask);
-        Mockito.verify(tasks, never()).removeTask(statefulTask);
+        verify(stateUpdater).add(statefulTask);
+        verify(tasks, never()).removeTask(statefulTask);
     }
 
     @Test
@@ -1668,11 +1664,11 @@ public void shouldUpdateInputPartitionsOfRestoredTask() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(consumer).resume(statefulTask.inputPartitions());
-        Mockito.verify(statefulTask).updateInputPartitions(Mockito.eq(taskId01Partitions), anyMap());
-        Mockito.verify(statefulTask).completeRestoration(noOpResetter);
-        Mockito.verify(statefulTask).clearTaskTimeout();
-        Mockito.verify(tasks).addTask(statefulTask);
+        verify(consumer).resume(statefulTask.inputPartitions());
+        verify(statefulTask).updateInputPartitions(eq(taskId01Partitions), anyMap());
+        verify(statefulTask).completeRestoration(noOpResetter);
+        verify(statefulTask).clearTaskTimeout();
+        verify(tasks).addTask(statefulTask);
     }
 
     @Test
@@ -1689,11 +1685,11 @@ public void shouldCloseReviveAndUpdateInputPartitionsOfRestoredTask() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(statefulTask).updateInputPartitions(Mockito.eq(taskId01Partitions), anyMap());
-        Mockito.verify(statefulTask).closeClean();
-        Mockito.verify(statefulTask).revive();
-        Mockito.verify(statefulTask).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(statefulTask);
+        verify(statefulTask).updateInputPartitions(eq(taskId01Partitions), anyMap());
+        verify(statefulTask).closeClean();
+        verify(statefulTask).revive();
+        verify(statefulTask).initializeIfNeeded();
+        verify(stateUpdater).add(statefulTask);
     }
 
     @Test
@@ -1712,9 +1708,9 @@ public void shouldSuspendRestoredTaskIfRevoked() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(statefulTask).suspend();
-        Mockito.verify(tasks).addTask(statefulTask);
-        Mockito.verifyNoInteractions(consumer);
+        verify(statefulTask).suspend();
+        verify(tasks).addTask(statefulTask);
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -1776,16 +1772,16 @@ public void shouldHandleMultipleRestoredTasks() {
 
         taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);
 
-        Mockito.verify(tasks).addTask(taskToTransitToRunning);
-        Mockito.verify(stateUpdater).add(recycledStandbyTask);
-        Mockito.verify(stateUpdater).add(recycledStandbyTask);
-        Mockito.verify(taskToCloseClean).closeClean();
-        Mockito.verify(stateUpdater).add(taskToAddBack);
-        Mockito.verify(taskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId05Partitions), anyMap());
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).closeClean();
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).revive();
-        Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded();
-        Mockito.verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions);
+        verify(tasks).addTask(taskToTransitToRunning);
+        verify(stateUpdater).add(recycledStandbyTask);
+        verify(stateUpdater).add(recycledStandbyTask);
+        verify(taskToCloseClean).closeClean();
+        verify(stateUpdater).add(taskToAddBack);
+        verify(taskToUpdateInputPartitions).updateInputPartitions(eq(taskId05Partitions), anyMap());
+        verify(taskToCloseReviveAndUpdateInputPartitions).closeClean();
+        verify(taskToCloseReviveAndUpdateInputPartitions).revive();
+        verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded();
+        verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions);
     }
 
     @Test
@@ -1892,9 +1888,9 @@ public void shouldRethrowTaskCorruptedExceptionFromInitialization() {
             () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
         );
 
-        Mockito.verify(tasks).addTask(statefulTask0);
-        Mockito.verify(tasks).addTask(statefulTask1);
-        Mockito.verify(stateUpdater).add(statefulTask2);
+        verify(tasks).addTask(statefulTask0);
+        verify(tasks).addTask(statefulTask1);
+        verify(stateUpdater).add(statefulTask2);
         assertEquals(mkSet(taskId00, taskId01), thrown.corruptedTasks());
         assertEquals("Tasks [0_1, 0_0] are corrupted and hence need to be re-initialized", thrown.getMessage());
     }
@@ -1913,19 +1909,17 @@ public void shouldAddSubscribedTopicsFromAssignmentToTopologyMetadata() {
 
         taskManager.handleAssignment(activeTasksAssignment, standbyTasksAssignment);
 
-        Mockito.verify(topologyBuilder).addSubscribedTopicsFromAssignment(Mockito.eq(mkSet(t1p1, t1p2, t2p2)), Mockito.anyString());
-        Mockito.verify(topologyBuilder, never()).addSubscribedTopicsFromAssignment(Mockito.eq(mkSet(t1p3, t1p4)), Mockito.anyString());
-        Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(activeTasksAssignment));
+        verify(topologyBuilder).addSubscribedTopicsFromAssignment(eq(mkSet(t1p1, t1p2, t2p2)), anyString());
+        verify(topologyBuilder, never()).addSubscribedTopicsFromAssignment(eq(mkSet(t1p3, t1p4)), anyString());
+        verify(activeTaskCreator).createTasks(any(), eq(activeTasksAssignment));
     }
 
     @Test
     public void shouldNotLockAnythingIfStateDirIsEmpty() {
-        expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()).once();
+        when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>());
 
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
-        verify(stateDirectory);
         assertTrue(taskManager.lockedTaskDirectories().isEmpty());
     }
 
@@ -1940,10 +1934,8 @@ public void shouldTryToLockValidTaskDirsAtRebalanceStart() throws Exception {
             taskId10.toString(),
             "dummy"
         );
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
-        verify(stateDirectory);
         assertThat(taskManager.lockedTaskDirectories(), is(singleton(taskId01)));
     }
 
@@ -1951,14 +1943,12 @@ public void shouldTryToLockValidTaskDirsAtRebalanceStart() throws Exception {
     public void shouldUnlockEmptyDirsAtRebalanceStart() throws Exception {
         expectLockObtainedFor(taskId01, taskId10);
         expectDirectoryNotEmpty(taskId01);
-        expect(stateDirectory.directoryForTaskIsEmpty(taskId10)).andReturn(true);
-        expectUnlockFor(taskId10);
+        when(stateDirectory.directoryForTaskIsEmpty(taskId10)).thenReturn(true);
 
         makeTaskFolders(taskId01.toString(), taskId10.toString());
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
-        verify(stateDirectory);
+        verify(stateDirectory).unlock(taskId10);
         assertThat(taskManager.lockedTaskDirectories(), is(singleton(taskId01)));
     }
 
@@ -1969,7 +1959,7 @@ public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() {
 
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(consumer).pause(assigned);
+        verify(consumer).pause(assigned);
     }
 
     @Test
@@ -1977,7 +1967,7 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
         final StreamTask statefulTask0 = statefulTask(taskId00, taskId00ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId00Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasks()).thenReturn(mkSet(statefulTask0));
         final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
@@ -1985,21 +1975,19 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
 
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(consumer).pause(mkSet(t1p1));
+        verify(consumer).pause(mkSet(t1p1));
     }
 
     @Test
     public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception {
         expectLockObtainedFor(taskId00, taskId01, taskId02);
         expectDirectoryNotEmpty(taskId00, taskId01, taskId02);
-        expectUnlockFor(taskId02);
 
         makeTaskFolders(
             taskId00.toString(),  // active task
             taskId01.toString(),  // standby task
             taskId02.toString()   // unassigned but able to lock
         );
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
@@ -2008,9 +1996,9 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception
 
         taskManager.handleRebalanceComplete();
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01)));
-        verify(stateDirectory);
 
-        Mockito.verify(consumer).pause(assignment);
+        verify(stateDirectory).unlock(taskId02);
+        verify(consumer).pause(assignment);
     }
 
     @Test
@@ -2027,21 +2015,19 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
         final StandbyTask unassignedStandbyTask = standbyTask(taskId03, taskId03ChangelogPartitions)
             .inState(State.CREATED)
             .withInputPartitions(taskId03Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId00, runningStatefulTask)));
         when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask, restoringStatefulTask));
         when(tasks.allTasks()).thenReturn(mkSet(runningStatefulTask));
         expectLockObtainedFor(taskId00, taskId01, taskId02, taskId03);
         expectDirectoryNotEmpty(taskId00, taskId01, taskId02, taskId03);
-        expectUnlockFor(taskId03);
         makeTaskFolders(
             taskId00.toString(),
             taskId01.toString(),
             taskId02.toString(),
             taskId03.toString()
         );
-        replay(stateDirectory);
 
         final Set<TopicPartition> assigned = mkSet(t1p0, t1p1, t1p2);
         when(consumer.assignment()).thenReturn(assigned);
@@ -2049,8 +2035,8 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
         taskManager.handleRebalanceStart(singleton("topic"));
         taskManager.handleRebalanceComplete();
 
-        Mockito.verify(consumer).pause(mkSet(t1p1, t1p2));
-        verify(stateDirectory);
+        verify(consumer).pause(mkSet(t1p1, t1p2));
+        verify(stateDirectory).unlock(taskId03);
         assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
     }
 
@@ -2086,10 +2072,9 @@ public void shouldComputeOffsetSumForRestoringActiveTaskWithStateUpdater() throw
         makeTaskFolders(taskId00.toString());
         final Map<TopicPartition, Long> changelogOffsetInCheckpoint = mkMap(mkEntry(t1p0changelog, 24L));
         writeCheckpointFile(taskId00, changelogOffsetInCheckpoint);
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(restoringStatefulTask));
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertThat(taskManager.getTaskOffsetSums(), is(mkMap(mkEntry(taskId00, changelogOffset))));
@@ -2105,10 +2090,9 @@ public void shouldComputeOffsetSumForRestoringStandbyTaskWithStateUpdater() thro
         makeTaskFolders(taskId00.toString());
         final Map<TopicPartition, Long> changelogOffsetInCheckpoint = mkMap(mkEntry(t1p0changelog, 24L));
         writeCheckpointFile(taskId00, changelogOffsetInCheckpoint);
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(stateUpdater.getTasks()).thenReturn(mkSet(restoringStandbyTask));
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertThat(taskManager.getTaskOffsetSums(), is(mkMap(mkEntry(taskId00, changelogOffset))));
@@ -2131,7 +2115,7 @@ public void shouldComputeOffsetSumForRunningStatefulTaskAndRestoringTaskWithStat
             .thenReturn(mkMap(mkEntry(t1p1changelog, changelogOffsetOfRestoringStatefulTask)));
         when(restoringStandbyTask.changelogOffsets())
             .thenReturn(mkMap(mkEntry(t1p2changelog, changelogOffsetOfRestoringStandbyTask)));
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId00, runningStatefulTask)));
         when(stateUpdater.getTasks()).thenReturn(mkSet(restoringStandbyTask, restoringStatefulTask));
@@ -2162,7 +2146,6 @@ private void computeOffsetSumAndVerify(final Map<TopicPartition, Long> changelog
         expectLockObtainedFor(taskId00);
         expectDirectoryNotEmpty(taskId00);
         makeTaskFolders(taskId00.toString());
-        replay(stateDirectory);
 
         taskManager.handleRebalanceStart(singleton("topic"));
         final StateMachineTask restoringTask = handleAssignment(
@@ -2186,7 +2169,6 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception {
         expectLockObtainedFor(taskId00);
         expectDirectoryNotEmpty(taskId00);
         makeTaskFolders(taskId00.toString());
-        replay(stateDirectory);
 
         taskManager.handleRebalanceStart(singleton("topic"));
         final StateMachineTask restoringTask = handleAssignment(
@@ -2211,7 +2193,6 @@ public void shouldComputeOffsetSumForUnassignedTaskWeCanLock() throws Exception
         makeTaskFolders(taskId00.toString());
         writeCheckpointFile(taskId00, changelogOffsets);
 
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums));
@@ -2228,11 +2209,10 @@ public void shouldComputeOffsetSumFromCheckpointFileForUninitializedTask() throw
         expectLockObtainedFor(taskId00);
         makeTaskFolders(taskId00.toString());
         writeCheckpointFile(taskId00, changelogOffsets);
-        replay(stateDirectory);
 
         taskManager.handleRebalanceStart(singleton("topic"));
         final StateMachineTask uninitializedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singleton(uninitializedTask));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singleton(uninitializedTask));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
 
@@ -2252,13 +2232,12 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep
         expectLockObtainedFor(taskId00);
         makeTaskFolders(taskId00.toString());
         writeCheckpointFile(taskId00, changelogOffsets);
-        replay(stateDirectory);
 
         final StateMachineTask closedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
         taskManager.handleRebalanceStart(singleton("topic"));
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singleton(closedTask));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singleton(closedTask));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
 
@@ -2273,7 +2252,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep
     public void shouldNotReportOffsetSumsForTaskWeCantLock() throws Exception {
         expectLockFailedFor(taskId00);
         makeTaskFolders(taskId00.toString());
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
         assertTrue(taskManager.lockedTaskDirectories().isEmpty());
 
@@ -2285,12 +2263,10 @@ public void shouldNotReportOffsetSumsAndReleaseLockForUnassignedTaskWithoutCheck
         expectLockObtainedFor(taskId00);
         makeTaskFolders(taskId00.toString());
         expectDirectoryNotEmpty(taskId00);
-        expect(stateDirectory.checkpointFileFor(taskId00)).andReturn(getCheckpointFile(taskId00));
-        replay(stateDirectory);
+        when(stateDirectory.checkpointFileFor(taskId00)).thenReturn(getCheckpointFile(taskId00));
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertTrue(taskManager.getTaskOffsetSums().isEmpty());
-        verify(stateDirectory);
     }
 
     @Test
@@ -2306,7 +2282,6 @@ public void shouldPinOffsetSumToLongMaxValueInCaseOfOverflow() throws Exception
         expectLockObtainedFor(taskId00);
         makeTaskFolders(taskId00.toString());
         writeCheckpointFile(taskId00, changelogOffsets);
-        replay(stateDirectory);
         taskManager.handleRebalanceStart(singleton("topic"));
 
         assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums));
@@ -2321,7 +2296,7 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() {
         // first `handleAssignment`
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -2334,7 +2309,7 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() {
         assertThat(task00.state(), is(Task.State.CLOSED));
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
     }
 
     @Test
@@ -2346,7 +2321,7 @@ public void closeClean() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         taskManager.handleRevocation(taskId00Partitions);
@@ -2362,7 +2337,7 @@ public void closeClean() {
             is("Encounter unexpected fatal error for task 0_0")
         );
         assertThat(thrown.getCause().getMessage(), is("KABOOM!"));
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
     }
 
     @Test
@@ -2372,18 +2347,20 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
 
         // `handleAssignment`
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
-        makeTaskFolders(taskId00.toString(), taskId01.toString());
+        final ArrayList<TaskDirectory> taskFolders = new ArrayList<>(2);
+        taskFolders.add(new TaskDirectory(testFolder.newFolder(taskId00.toString()), null));
+        taskFolders.add(new TaskDirectory(testFolder.newFolder(taskId01.toString()), null));
+
+        when(stateDirectory.listNonEmptyTaskDirectories())
+            .thenReturn(taskFolders)
+            .thenReturn(new ArrayList<>());
+
         expectLockObtainedFor(taskId00, taskId01);
         expectDirectoryNotEmpty(taskId00, taskId01);
 
-        // The second attempt will return empty tasks.
-        makeTaskFolders();
-        expectLockObtainedFor();
-        replay(stateDirectory);
-
         taskManager.handleRebalanceStart(emptySet());
         assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01)));
 
@@ -2406,7 +2383,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
         taskManager.handleRebalanceStart(emptySet());
 
         assertThat(taskManager.lockedTaskDirectories(), is(emptySet()));
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
     }
 
     @Test
@@ -2415,7 +2392,7 @@ public void shouldReInitializeThreadProducerOnHandleLostAllIfEosV2Enabled() {
 
         taskManager.handleLostAll();
 
-        Mockito.verify(activeTaskCreator).reInitializeThreadProducer();
+        verify(activeTaskCreator).reInitializeThreadProducer();
     }
 
     @Test
@@ -2426,7 +2403,7 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() {
 
         // `handleAssignment`
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         // `handleAssignment`
         doThrow(new RuntimeException("KABOOM!")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
@@ -2458,7 +2435,7 @@ public void shouldReAddRevivedTasksToStateUpdater() {
         final StandbyTask corruptedStandbyTask = standbyTask(taskId02, taskId02ChangelogPartitions)
             .inState(State.RUNNING)
             .withInputPartitions(taskId02Partitions).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
         when(tasks.task(taskId03)).thenReturn(corruptedActiveTask);
         when(tasks.task(taskId02)).thenReturn(corruptedStandbyTask);
@@ -2471,16 +2448,16 @@ public void shouldReAddRevivedTasksToStateUpdater() {
         final InOrder standbyTaskOrder = inOrder(corruptedStandbyTask);
         standbyTaskOrder.verify(corruptedStandbyTask).closeDirty();
         standbyTaskOrder.verify(corruptedStandbyTask).revive();
-        Mockito.verify(tasks).removeTask(corruptedActiveTask);
-        Mockito.verify(tasks).removeTask(corruptedStandbyTask);
-        Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedActiveTask));
-        Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedStandbyTask));
-        Mockito.verify(consumer).assignment();
+        verify(tasks).removeTask(corruptedActiveTask);
+        verify(tasks).removeTask(corruptedStandbyTask);
+        verify(tasks).addPendingTasksToInit(mkSet(corruptedActiveTask));
+        verify(tasks).addPendingTasksToInit(mkSet(corruptedStandbyTask));
+        verify(consumer).assignment();
     }
 
     @Test
     public void shouldReviveCorruptTasks() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final AtomicBoolean enforcedCheckpoint = new AtomicBoolean(false);
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
@@ -2497,7 +2474,7 @@ public void postCommit(final boolean enforceCheckpoint) {
         when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(taskId00Partitions);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
@@ -2513,12 +2490,12 @@ public void postCommit(final boolean enforceCheckpoint) {
         assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
 
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
     public void shouldReviveCorruptTasksEvenIfTheyCannotCloseClean() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
             @Override
@@ -2531,7 +2508,7 @@ public void suspend() {
         when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(taskId00Partitions);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
@@ -2545,12 +2522,12 @@ public void suspend() {
         assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
 
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
     public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final StateMachineTask nonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
@@ -2559,7 +2536,7 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         firstAssignment.putAll(taskId01Assignment);
 
         // `handleAssignment`
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
+        when(activeTaskCreator.createTasks(any(), eq(firstAssignment)))
             .thenReturn(asList(corruptedTask, nonCorruptedTask));
 
         when(consumer.assignment())
@@ -2580,13 +2557,13 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
         // check that we should not commit empty map either
-        Mockito.verify(consumer, never()).commitSync(emptyMap());
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(consumer, never()).commitSync(emptyMap());
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
     public void shouldNotCommitNonRunningNonCorruptedTasks() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final StateMachineTask nonRunningNonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager);
@@ -2597,7 +2574,7 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         assignment.putAll(taskId01Assignment);
 
         // `handleAssignment`
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(assignment)))
             .thenReturn(asList(corruptedTask, nonRunningNonCorruptedTask));
         when(consumer.assignment()).thenReturn(taskId00Partitions);
 
@@ -2611,7 +2588,7 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
         assertFalse(nonRunningNonCorruptedTask.commitPrepared);
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
@@ -2625,21 +2602,20 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndNotCommitRunningSt
         final StreamTask corruptedTask = statefulTask(taskId02, taskId02ChangelogPartitions)
             .withInputPartitions(taskId02Partitions)
             .inState(State.RUNNING).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId02, corruptedTask)));
         when(tasks.task(taskId02)).thenReturn(corruptedTask);
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
-        when(stateUpdater.getTasks()).thenReturn(mkSet(activeRestoringTask, standbyTask));
         when(consumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
         taskManager.handleCorruption(mkSet(taskId02));
 
-        Mockito.verify(activeRestoringTask, never()).commitNeeded();
-        Mockito.verify(activeRestoringTask, never()).prepareCommit();
-        Mockito.verify(activeRestoringTask, never()).postCommit(Mockito.anyBoolean());
-        Mockito.verify(standbyTask, never()).commitNeeded();
-        Mockito.verify(standbyTask, never()).prepareCommit();
-        Mockito.verify(standbyTask, never()).postCommit(Mockito.anyBoolean());
+        verify(activeRestoringTask, never()).commitNeeded();
+        verify(activeRestoringTask, never()).prepareCommit();
+        verify(activeRestoringTask, never()).postCommit(anyBoolean());
+        verify(standbyTask, never()).commitNeeded();
+        verify(standbyTask, never()).prepareCommit();
+        verify(standbyTask, never()).postCommit(anyBoolean());
     }
 
     @Test
@@ -2654,7 +2630,7 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndCommitRunningStand
         final StreamTask corruptedTask = statefulTask(taskId02, taskId02ChangelogPartitions)
             .withInputPartitions(taskId02Partitions)
             .inState(State.RUNNING).build();
-        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TasksRegistry tasks = mock(TasksRegistry.class);
         when(tasks.allTasksPerId()).thenReturn(mkMap(
             mkEntry(taskId00, activeRestoringTask),
             mkEntry(taskId01, standbyTask),
@@ -2666,16 +2642,16 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndCommitRunningStand
 
         taskManager.handleCorruption(mkSet(taskId02));
 
-        Mockito.verify(activeRestoringTask, never()).commitNeeded();
-        Mockito.verify(activeRestoringTask, never()).prepareCommit();
-        Mockito.verify(activeRestoringTask, never()).postCommit(Mockito.anyBoolean());
-        Mockito.verify(standbyTask).prepareCommit();
-        Mockito.verify(standbyTask).postCommit(Mockito.anyBoolean());
+        verify(activeRestoringTask, never()).commitNeeded();
+        verify(activeRestoringTask, never()).prepareCommit();
+        verify(activeRestoringTask, never()).postCommit(anyBoolean());
+        verify(standbyTask).prepareCommit();
+        verify(standbyTask).postCommit(anyBoolean());
     }
 
     @Test
     public void shouldCleanAndReviveCorruptedStandbyTasksBeforeCommittingNonCorruptedTasks() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask corruptedStandby = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
         final StateMachineTask runningNonCorruptedActive = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) {
@@ -2686,7 +2662,7 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
         };
 
         // handleAssignment
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId01Assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(taskId01Assignment)))
             .thenReturn(singleton(runningNonCorruptedActive));
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singleton(corruptedStandby));
 
@@ -2707,13 +2683,13 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
 
         assertThat(corruptedStandby.commitPrepared, is(true));
         assertThat(corruptedStandby.state(), is(Task.State.CREATED));
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
     public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
-        expect(stateDirectory.listNonEmptyTaskDirectories()).andStubReturn(new ArrayList<>());
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
+        when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>());
 
         final StateMachineTask corruptedActive = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
@@ -2726,15 +2702,13 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
         final Map<TaskId, Set<TopicPartition>> firstAssignement = new HashMap<>();
         firstAssignement.putAll(taskId00Assignment);
         firstAssignement.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignement)))
+        when(activeTaskCreator.createTasks(any(), eq(firstAssignement)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
         when(consumer.assignment())
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));
 
-        replay(stateDirectory);
-
         uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);
 
         taskManager.handleAssignment(firstAssignement, emptyMap());
@@ -2759,7 +2733,7 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
 
     @Test
     public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCommitWithALOS() {
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask corruptedActive = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final StateMachineTask uncorruptedActive = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) {
@@ -2775,7 +2749,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>();
         firstAssignment.putAll(taskId00Assignment);
         firstAssignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
+        when(activeTaskCreator.createTasks(any(), eq(firstAssignment)))
             .thenReturn(asList(corruptedActive, uncorruptedActive));
 
         when(consumer.assignment())
@@ -2812,7 +2786,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
 
         assertThat(corruptedActive.state(), is(Task.State.CREATED));
         assertThat(uncorruptedActive.state(), is(Task.State.CREATED));
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
     }
 
     @Test
@@ -2820,7 +2794,7 @@ public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCo
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false);
         final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.threadProducer()).thenReturn(producer);
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final AtomicBoolean corruptedTaskChangelogMarkedAsCorrupted = new AtomicBoolean(false);
         final StateMachineTask corruptedActiveTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
@@ -2846,7 +2820,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         final Map<TaskId, Set<TopicPartition>> firstAssignment = new HashMap<>();
         firstAssignment.putAll(taskId00Assignment);
         firstAssignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
+        when(activeTaskCreator.createTasks(any(), eq(firstAssignment)))
             .thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask));
 
         when(consumer.assignment())
@@ -2892,8 +2866,8 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         assertThat(uncorruptedActiveTask.state(), is(Task.State.CREATED));
         assertThat(corruptedTaskChangelogMarkedAsCorrupted.get(), is(true));
         assertThat(uncorruptedTaskChangelogMarkedAsCorrupted.get(), is(true));
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions);
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions);
     }
 
     @Test
@@ -2929,7 +2903,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded));
 
         doThrow(new TimeoutException()).when(consumer).commitSync(expectedCommittedOffsets);
@@ -2952,7 +2926,7 @@ public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCo
         final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false);
         final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.threadProducer()).thenReturn(producer);
-        final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
+        final ProcessorStateManager stateManager = mock(ProcessorStateManager.class);
 
         final StateMachineTask revokedActiveTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         final Map<TopicPartition, OffsetAndMetadata> revokedActiveTaskOffsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
@@ -2987,7 +2961,7 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
             .thenReturn(assignment)
             .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded));
 
         final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId");
@@ -3012,8 +2986,8 @@ public void markChangelogAsCorrupted(final Collection<TopicPartition> partitions
         assertThat(revokedActiveTask.state(), is(State.SUSPENDED));
         assertThat(unrevokedActiveTask.state(), is(State.CREATED));
         assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING));
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions);
-        Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions);
+        verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions);
     }
 
     @Test
@@ -3039,7 +3013,7 @@ public void shouldAddNonResumedSuspendedTasks() {
         final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));
 
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
@@ -3053,9 +3027,9 @@ public void shouldAddNonResumedSuspendedTasks() {
         assertThat(task01.state(), is(Task.State.RUNNING));
 
         // expect these calls twice (because we're going to tryToCompleteRestoration twice)
-        Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
-        Mockito.verify(consumer, times(2)).assignment();
-        Mockito.verify(consumer, times(2)).resume(assignment);
+        verify(activeTaskCreator).createTasks(any(), eq(emptyMap()));
+        verify(consumer, times(2)).assignment();
+        verify(consumer, times(2)).resume(assignment);
     }
 
     @Test
@@ -3063,7 +3037,7 @@ public void shouldUpdateInputPartitionsAfterRebalance() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3076,9 +3050,9 @@ public void shouldUpdateInputPartitionsAfterRebalance() {
         assertThat(task00.state(), is(Task.State.RUNNING));
         assertEquals(newPartitionsSet, task00.inputPartitions());
         // expect these calls twice (because we're going to tryToCompleteRestoration twice)
-        Mockito.verify(consumer, times(2)).resume(assignment);
-        Mockito.verify(consumer, times(2)).assignment();
-        Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
+        verify(consumer, times(2)).resume(assignment);
+        verify(consumer, times(2)).assignment();
+        verify(activeTaskCreator).createTasks(any(), eq(emptyMap()));
     }
 
     @Test
@@ -3086,7 +3060,7 @@ public void shouldAddNewActiveTasks() {
         final Map<TaskId, Set<TopicPartition>> assignment = taskId00Assignment;
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(assignment, emptyMap());
 
@@ -3097,9 +3071,9 @@ public void shouldAddNewActiveTasks() {
         assertThat(task00.state(), is(Task.State.RUNNING));
         assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(consumer).assignment();
-        Mockito.verify(consumer).resume(Mockito.eq(emptySet()));
+        verify(changeLogReader).enforceRestoreActive();
+        verify(consumer).assignment();
+        verify(consumer).resume(eq(emptySet()));
     }
 
     @Test
@@ -3121,7 +3095,7 @@ public void initializeIfNeeded() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01));
 
         taskManager.handleAssignment(assignment, emptyMap());
 
@@ -3137,8 +3111,8 @@ public void initializeIfNeeded() {
             Matchers.equalTo(mkMap(mkEntry(taskId00, task00), mkEntry(taskId01, task01)))
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verifyNoInteractions(consumer);
+        verify(changeLogReader).enforceRestoreActive();
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -3153,7 +3127,7 @@ public void completeRestoration(final java.util.function.Consumer<Set<TopicParti
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(assignment, emptyMap());
 
@@ -3167,8 +3141,8 @@ public void completeRestoration(final java.util.function.Consumer<Set<TopicParti
             Matchers.equalTo(mkMap(mkEntry(taskId00, task00)))
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verifyNoInteractions(consumer);
+        verify(changeLogReader).enforceRestoreActive();
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -3178,7 +3152,7 @@ public void shouldSuspendActiveTasksDuringRevocation() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -3224,7 +3198,7 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
         );
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
 
         when(activeTaskCreator.threadProducer()).thenReturn(producer);
@@ -3253,7 +3227,7 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo
         assertThat(task02.commitPrepared, is(false));
         assertThat(task10.commitPrepared, is(false));
 
-        Mockito.verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata);
+        verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata);
     }
 
     @Test
@@ -3289,7 +3263,7 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         );
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task10));
@@ -3310,7 +3284,7 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
         assertThat(task02.commitPrepared, is(false));
         assertThat(task10.commitPrepared, is(false));
 
-        Mockito.verify(consumer).commitSync(expectedCommittedOffsets);
+        verify(consumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -3327,7 +3301,7 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() {
 
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
@@ -3355,7 +3329,7 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() {
 
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(singleton(task00));
         when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10));
 
         taskManager.handleAssignment(assignmentActive, assignmentStandby);
@@ -3372,7 +3346,7 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() {
     public void shouldNotCommitCreatedTasksOnRevocationOrClosure() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(task00.state(), is(Task.State.CREATED));
@@ -3382,7 +3356,7 @@ public void shouldNotCommitCreatedTasksOnRevocationOrClosure() {
 
         taskManager.handleAssignment(emptyMap(), emptyMap());
         assertThat(task00.state(), is(Task.State.CLOSED));
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(eq(taskId00));
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
     }
 
     @Test
@@ -3396,7 +3370,7 @@ public void suspend() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
@@ -3463,7 +3437,7 @@ public void closeDirty() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(assignment)))
             .thenReturn(asList(task00, task01, task02, task03));
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3491,8 +3465,8 @@ public void closeDirty() {
             )
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(changeLogReader).completedChangelogs();
+        verify(changeLogReader).enforceRestoreActive();
+        verify(changeLogReader).completedChangelogs();
 
         final RuntimeException exception = assertThrows(
             RuntimeException.class,
@@ -3509,9 +3483,9 @@ public void closeDirty() {
         assertThat(task03.state(), is(Task.State.CLOSED));
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(activeTaskCreator, times(4)).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(activeTaskCreator, times(4)).closeAndRemoveTaskProducerIfNeeded(any());
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
     }
 
     @Test
@@ -3529,7 +3503,7 @@ public Set<TopicPartition> changelogPartitions() {
         final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
         task00.setCommittableOffsetsAndMetadata(offsets);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00));
         doThrow(new RuntimeException("whatever"))
             .when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
 
@@ -3549,8 +3523,8 @@ public Set<TopicPartition> changelogPartitions() {
             )
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(changeLogReader).completedChangelogs();
+        verify(changeLogReader).enforceRestoreActive();
+        verify(changeLogReader).completedChangelogs();
 
         final RuntimeException exception = assertThrows(RuntimeException.class, () -> taskManager.shutdown(true));
 
@@ -3558,9 +3532,9 @@ public Set<TopicPartition> changelogPartitions() {
         assertThat(exception.getCause().getMessage(), is("whatever"));
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
     }
 
     @Test
@@ -3576,7 +3550,7 @@ public Set<TopicPartition> changelogPartitions() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00));
         doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeThreadProducerIfNeeded();
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3595,8 +3569,8 @@ public Set<TopicPartition> changelogPartitions() {
             )
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(changeLogReader).completedChangelogs();
+        verify(changeLogReader).enforceRestoreActive();
+        verify(changeLogReader).completedChangelogs();
 
         final RuntimeException exception = assertThrows(RuntimeException.class, () -> taskManager.shutdown(true));
 
@@ -3605,7 +3579,7 @@ public Set<TopicPartition> changelogPartitions() {
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
     }
 
     @Test
@@ -3666,7 +3640,7 @@ public void suspend() {
         assertThat(task01.state(), is(Task.State.SUSPENDED));
         assertThat(task02.state(), is(Task.State.SUSPENDED));
 
-        Mockito.verifyNoInteractions(activeTaskCreator);
+        verifyNoInteractions(activeTaskCreator);
     }
 
     @Test
@@ -3698,8 +3672,8 @@ public void suspend() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01, task02));
-        doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(Mockito.any());
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01, task02));
+        doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any());
         doThrow(new RuntimeException("whatever all")).when(activeTaskCreator).closeThreadProducerIfNeeded();
 
         taskManager.handleAssignment(assignment, emptyMap());
@@ -3724,8 +3698,8 @@ public void suspend() {
             )
         );
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(changeLogReader).enforceRestoreActive();
-        Mockito.verify(changeLogReader).completedChangelogs();
+        verify(changeLogReader).enforceRestoreActive();
+        verify(changeLogReader).completedChangelogs();
 
         taskManager.shutdown(false);
 
@@ -3734,9 +3708,9 @@ public void suspend() {
         assertThat(task02.state(), is(Task.State.CLOSED));
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
-        Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
+        verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
     }
 
     @Test
@@ -3760,10 +3734,10 @@ public void shouldCloseStandbyTasksOnShutdown() {
         assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap());
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         // the active task creator should also get closed (so that it closes the thread producer if applicable)
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
         // `tryToCompleteRestoration`
-        Mockito.verify(consumer).assignment();
-        Mockito.verify(consumer).resume(Mockito.eq(emptySet()));
+        verify(consumer).assignment();
+        verify(consumer).resume(eq(emptySet()));
     }
 
     @Test
@@ -3782,12 +3756,12 @@ public void shouldShutDownStateUpdaterAndCloseFailedTasksDirty() {
 
         taskManager.shutdown(true);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(failedStatefulTask.id());
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
-        Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
-        Mockito.verify(failedStatefulTask).prepareCommit();
-        Mockito.verify(failedStatefulTask).suspend();
-        Mockito.verify(failedStatefulTask).closeDirty();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(failedStatefulTask.id());
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
+        verify(failedStatefulTask).prepareCommit();
+        verify(failedStatefulTask).suspend();
+        verify(failedStatefulTask).closeDirty();
     }
 
     @Test
@@ -3797,7 +3771,7 @@ public void shouldShutdownSchedulingTaskManager() {
 
         taskManager.shutdown(true);
 
-        Mockito.verify(schedulingTaskManager).shutdown(Duration.ofMillis(Long.MAX_VALUE));
+        verify(schedulingTaskManager).shutdown(Duration.ofMillis(Long.MAX_VALUE));
     }
 
     @Test
@@ -3815,13 +3789,13 @@ public void shouldShutDownStateUpdaterAndAddRestoredTasksToTaskRegistry() {
 
         taskManager.shutdown(true);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask1.id());
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask2.id());
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
-        Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
-        Mockito.verify(tasks).addActiveTasks(restoredTasks);
-        Mockito.verify(statefulTask1).closeClean();
-        Mockito.verify(statefulTask2).closeClean();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask1.id());
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask2.id());
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
+        verify(tasks).addActiveTasks(restoredTasks);
+        verify(statefulTask1).closeClean();
+        verify(statefulTask2).closeClean();
     }
 
     @Test
@@ -3838,13 +3812,13 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() {
 
         taskManager.shutdown(true);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(removedStatefulTask.id());
-        Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded();
-        Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
-        Mockito.verify(tasks).addActiveTasks(mkSet(removedStatefulTask));
-        Mockito.verify(tasks).addStandbyTasks(mkSet(removedStandbyTask));
-        Mockito.verify(removedStatefulTask).closeClean();
-        Mockito.verify(removedStandbyTask).closeClean();
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(removedStatefulTask.id());
+        verify(activeTaskCreator).closeThreadProducerIfNeeded();
+        verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE));
+        verify(tasks).addActiveTasks(mkSet(removedStatefulTask));
+        verify(tasks).addStandbyTasks(mkSet(removedStandbyTask));
+        verify(removedStatefulTask).closeClean();
+        verify(removedStandbyTask).closeClean();
     }
 
     @Test
@@ -3852,7 +3826,7 @@ public void shouldInitializeNewActiveTasks() {
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);
         when(consumer.assignment()).thenReturn(assignment);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
@@ -3862,7 +3836,7 @@ public void shouldInitializeNewActiveTasks() {
         assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00)));
         assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
         // verifies that we actually resume the assignment at the end of restoration.
-        Mockito.verify(consumer).resume(assignment);
+        verify(consumer).resume(assignment);
     }
 
     @Test
@@ -3883,14 +3857,13 @@ public void shouldInitialiseNewStandbyTasks() {
     @Test
     public void shouldHandleRebalanceEvents() {
         when(consumer.assignment()).thenReturn(assignment);
-        expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>());
-        replay(stateDirectory);
+        when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>());
         assertThat(taskManager.rebalanceInProgress(), is(false));
         taskManager.handleRebalanceStart(emptySet());
         assertThat(taskManager.rebalanceInProgress(), is(true));
         taskManager.handleRebalanceComplete();
         assertThat(taskManager.rebalanceInProgress(), is(false));
-        Mockito.verify(consumer).pause(assignment);
+        verify(consumer).pause(assignment);
     }
 
     @Test
@@ -3901,7 +3874,7 @@ public void shouldCommitActiveAndStandbyTasks() {
         final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
             .thenReturn(singletonList(task01));
@@ -3919,7 +3892,7 @@ public void shouldCommitActiveAndStandbyTasks() {
         assertThat(task00.commitNeeded, is(false));
         assertThat(task01.commitNeeded, is(false));
 
-        Mockito.verify(consumer).commitSync(offsets);
+        verify(consumer).commitSync(offsets);
     }
 
     @Test
@@ -3943,7 +3916,7 @@ public void shouldCommitProvidedTasksIfNeeded() {
         );
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(Arrays.asList(task00, task01, task02));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(Arrays.asList(task03, task04, task05));
@@ -3995,13 +3968,11 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
         expectDirectoryNotEmpty(taskId00, taskId01);
         expectLockObtainedFor(taskId00, taskId01);
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
         when(standbyTaskCreator.createTasks(taskId01Assignment))
             .thenReturn(singletonList(task01));
 
-        replay(stateDirectory);
-
         taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
 
@@ -4034,12 +4005,12 @@ public void shouldCommitViaConsumerIfEosDisabled() {
 
         taskManager.commitAll();
 
-        Mockito.verify(consumer).commitSync(offsets);
+        verify(consumer).commitSync(offsets);
     }
 
     @Test
     public void shouldCommitViaProducerIfEosAlphaEnabled() {
-        final StreamsProducer producer = Mockito.mock(StreamsProducer.class);
+        final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.streamsProducerForTask(any(TaskId.class)))
             .thenReturn(producer);
 
@@ -4048,14 +4019,14 @@ public void shouldCommitViaProducerIfEosAlphaEnabled() {
 
         shouldCommitViaProducerIfEosEnabled(ProcessingMode.EXACTLY_ONCE_ALPHA, offsetsT01, offsetsT02);
 
-        Mockito.verify(producer).commitTransaction(offsetsT01, new ConsumerGroupMetadata("appId"));
-        Mockito.verify(producer).commitTransaction(offsetsT02, new ConsumerGroupMetadata("appId"));
-        Mockito.verifyNoMoreInteractions(producer);
+        verify(producer).commitTransaction(offsetsT01, new ConsumerGroupMetadata("appId"));
+        verify(producer).commitTransaction(offsetsT02, new ConsumerGroupMetadata("appId"));
+        verifyNoMoreInteractions(producer);
     }
 
     @Test
     public void shouldCommitViaProducerIfEosV2Enabled() {
-        final StreamsProducer producer = Mockito.mock(StreamsProducer.class);
+        final StreamsProducer producer = mock(StreamsProducer.class);
         when(activeTaskCreator.threadProducer()).thenReturn(producer);
 
         final Map<TopicPartition, OffsetAndMetadata> offsetsT01 = singletonMap(t1p1, new OffsetAndMetadata(0L, null));
@@ -4066,8 +4037,8 @@ public void shouldCommitViaProducerIfEosV2Enabled() {
 
         shouldCommitViaProducerIfEosEnabled(ProcessingMode.EXACTLY_ONCE_V2, offsetsT01, offsetsT02);
 
-        Mockito.verify(producer).commitTransaction(allOffsets, new ConsumerGroupMetadata("appId"));
-        Mockito.verifyNoMoreInteractions(producer);
+        verify(producer).commitTransaction(allOffsets, new ConsumerGroupMetadata("appId"));
+        verifyNoMoreInteractions(producer);
     }
 
     private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processingMode,
@@ -4099,7 +4070,7 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4144,7 +4115,7 @@ public void shouldSendPurgeData() {
         when(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(17L))))
             .thenReturn(new DeleteRecordsResult(singletonMap(t1p1, completedFuture())));
 
-        final InOrder inOrder = Mockito.inOrder(adminClient);
+        final InOrder inOrder = inOrder(adminClient);
 
         final Map<TopicPartition, Long> purgableOffsets = new HashMap<>();
         final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
@@ -4155,7 +4126,7 @@ public Map<TopicPartition, Long> purgeableOffsets() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4188,7 +4159,7 @@ public Map<TopicPartition, Long> purgeableOffsets() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4257,7 +4228,7 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
         );
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive)))
+        when(activeTaskCreator.createTasks(any(), eq(assignmentActive)))
             .thenReturn(asList(task00, task01, task02, task03));
         when(standbyTaskCreator.createTasks(assignmentStandby))
             .thenReturn(singletonList(task04));
@@ -4286,7 +4257,7 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() {
 
         assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3));
 
-        Mockito.verify(consumer).commitSync(expectedCommittedOffsets);
+        verify(consumer).commitSync(expectedCommittedOffsets);
     }
 
     @Test
@@ -4299,7 +4270,7 @@ public void shouldProcessActiveTasks() {
         firstAssignment.put(taskId01, taskId01Partitions);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment)))
+        when(activeTaskCreator.createTasks(any(), eq(firstAssignment)))
             .thenReturn(Arrays.asList(task00, task01));
 
         taskManager.handleAssignment(firstAssignment, emptyMap());
@@ -4411,7 +4382,7 @@ public boolean process(final long wallClockTime) {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4434,7 +4405,7 @@ public boolean process(final long wallClockTime) {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment)))
             .thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
@@ -4461,7 +4432,7 @@ public boolean maybePunctuateStreamTime() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4481,7 +4452,7 @@ public boolean maybePunctuateStreamTime() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4506,7 +4477,7 @@ public boolean maybePunctuateSystemTime() {
         };
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
@@ -4526,12 +4497,12 @@ public Set<TopicPartition> changelogPartitions() {
             }
         };
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         taskManager.handleAssignment(taskId00Assignment, emptyMap());
         assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(false));
         assertThat(task00.state(), is(Task.State.RESTORING));
-        Mockito.verifyNoInteractions(consumer);
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -4541,7 +4512,7 @@ public void shouldHaveRemainingPartitionsUncleared() {
         task00.setCommittableOffsetsAndMetadata(offsets);
 
         when(consumer.assignment()).thenReturn(assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));
 
         try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) {
             appender.setClassLoggerToDebug(TaskManager.class);
@@ -4695,7 +4666,7 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
         allActiveTasks.addAll(restoringTasks);
 
         when(standbyTaskCreator.createTasks(standbyAssignment)).thenReturn(standbyTasks);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(allActiveTasksAssignment))).thenReturn(allActiveTasks);
+        when(activeTaskCreator.createTasks(any(), eq(allActiveTasksAssignment))).thenReturn(allActiveTasks);
 
         lenient().when(consumer.assignment()).thenReturn(assignment);
 
@@ -4720,28 +4691,21 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
         return allTasks;
     }
 
-    private void expectLockObtainedFor(final TaskId... tasks) throws Exception {
+    private void expectLockObtainedFor(final TaskId... tasks) {
         for (final TaskId task : tasks) {
-            expect(stateDirectory.lock(task)).andReturn(true).once();
+            when(stateDirectory.lock(task)).thenReturn(true);
         }
     }
 
-    private void expectLockFailedFor(final TaskId... tasks) throws Exception {
+    private void expectLockFailedFor(final TaskId... tasks) {
         for (final TaskId task : tasks) {
-            expect(stateDirectory.lock(task)).andReturn(false).once();
-        }
-    }
-
-    private void expectUnlockFor(final TaskId... tasks) throws Exception {
-        for (final TaskId task : tasks) {
-            stateDirectory.unlock(task);
-            expectLastCall();
+            when(stateDirectory.lock(task)).thenReturn(false);
         }
     }
 
     private void expectDirectoryNotEmpty(final TaskId... tasks) {
         for (final TaskId taskId : tasks) {
-            expect(stateDirectory.directoryForTaskIsEmpty(taskId)).andReturn(false);
+            when(stateDirectory.directoryForTaskIsEmpty(taskId)).thenReturn(false);
         }
     }
 
@@ -4790,7 +4754,7 @@ public void shouldNotFailForTimeoutExceptionOnConsumerCommit() {
         assertNull(task00.timeout);
         assertNull(task01.timeout);
 
-        Mockito.verify(consumer, times(2)).commitSync(any(Map.class));
+        verify(consumer, times(2)).commitSync(any(Map.class));
     }
 
     @Test
@@ -4832,7 +4796,7 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() {
             equalTo(Collections.singleton(taskId00))
         );
 
-        Mockito.verify(consumer, times(2)).groupMetadata();
+        verify(consumer, times(2)).groupMetadata();
     }
 
     @Test
@@ -4867,7 +4831,7 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV
             equalTo(mkSet(taskId00, taskId01))
         );
 
-        Mockito.verify(consumer).groupMetadata();
+        verify(consumer).groupMetadata();
     }
 
     @Test
@@ -4922,7 +4886,7 @@ public void suspend() {
 
         final Map<TaskId, Set<TopicPartition>> assignment = new HashMap<>(taskId00Assignment);
         assignment.putAll(taskId01Assignment);
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01));
+        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01));
 
         taskManager.handleAssignment(assignment, Collections.emptyMap());
 
@@ -4933,29 +4897,29 @@ public void suspend() {
         assertThat(thrown.getCause().getMessage(), is("KABOOM!"));
         assertThat(task00.state(), is(Task.State.SUSPENDED));
         assertThat(task01.state(), is(Task.State.SUSPENDED));
-        Mockito.verifyNoInteractions(consumer);
+        verifyNoInteractions(consumer);
     }
 
     @Test
     public void shouldConvertActiveTaskToStandbyTask() {
-        final StreamTask activeTask = Mockito.mock(StreamTask.class);
+        final StreamTask activeTask = mock(StreamTask.class);
         when(activeTask.id()).thenReturn(taskId00);
         when(activeTask.inputPartitions()).thenReturn(taskId00Partitions);
         when(activeTask.isActive()).thenReturn(true);
 
-        final StandbyTask standbyTask = Mockito.mock(StandbyTask.class);
+        final StandbyTask standbyTask = mock(StandbyTask.class);
         when(standbyTask.id()).thenReturn(taskId00);
 
-        when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(activeTask));
-        when(standbyTaskCreator.createStandbyTaskFromActive(Mockito.any(), Mockito.eq(taskId00Partitions))).thenReturn(standbyTask);
+        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(activeTask));
+        when(standbyTaskCreator.createStandbyTaskFromActive(any(), eq(taskId00Partitions))).thenReturn(standbyTask);
 
         taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap());
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
 
-        Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
-        Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap()));
-        Mockito.verify(standbyTaskCreator, times(2)).createTasks(Collections.emptyMap());
-        Mockito.verifyNoInteractions(consumer);
+        verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
+        verify(activeTaskCreator).createTasks(any(), eq(emptyMap()));
+        verify(standbyTaskCreator, times(2)).createTasks(Collections.emptyMap());
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -4969,15 +4933,15 @@ public void shouldConvertStandbyTaskToActiveTask() {
         when(activeTask.id()).thenReturn(taskId00);
         when(activeTask.inputPartitions()).thenReturn(taskId00Partitions);
         when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(standbyTask));
-        when(activeTaskCreator.createActiveTaskFromStandby(Mockito.eq(standbyTask), Mockito.eq(taskId00Partitions), any()))
+        when(activeTaskCreator.createActiveTaskFromStandby(eq(standbyTask), eq(taskId00Partitions), any()))
             .thenReturn(activeTask);
 
         taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment);
         taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap());
 
-        Mockito.verify(activeTaskCreator, times(2)).createTasks(any(), Mockito.eq(emptyMap()));
-        Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap());
-        Mockito.verifyNoInteractions(consumer);
+        verify(activeTaskCreator, times(2)).createTasks(any(), eq(emptyMap()));
+        verify(standbyTaskCreator).createTasks(Collections.emptyMap());
+        verifyNoInteractions(consumer);
     }
 
     @Test
@@ -5002,14 +4966,14 @@ private void makeTaskFolders(final String... names) throws Exception {
         for (int i = 0; i < names.length; ++i) {
             taskFolders.add(new TaskDirectory(testFolder.newFolder(names[i]), null));
         }
-        expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(taskFolders).once();
+        when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(taskFolders);
     }
 
     private void writeCheckpointFile(final TaskId task, final Map<TopicPartition, Long> offsets) throws Exception {
         final File checkpointFile = getCheckpointFile(task);
         Files.createFile(checkpointFile.toPath());
         new OffsetCheckpoint(checkpointFile).write(offsets);
-        expect(stateDirectory.checkpointFileFor(task)).andReturn(checkpointFile);
+        lenient().when(stateDirectory.checkpointFileFor(task)).thenReturn(checkpointFile);
         expectDirectoryNotEmpty(task);
     }
 

From fb03da0df47775457a2785251739f4ad6be5cd67 Mon Sep 17 00:00:00 2001
From: Philip Nee <pnee@confluent.io>
Date: Fri, 22 Mar 2024 03:17:49 -0700
Subject: [PATCH 188/258] KAFKA-16271: Upgrade consumer_rolling_upgrade_test.py
 (#15578)

Upgrading the test to use the consumer group protocol. The two tests are failing due to Mismatch Assignment

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../tests/client/consumer_rolling_upgrade_test.py     | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/kafkatest/tests/client/consumer_rolling_upgrade_test.py b/tests/kafkatest/tests/client/consumer_rolling_upgrade_test.py
index 9805d3656ee24..d3b342bbfae49 100644
--- a/tests/kafkatest/tests/client/consumer_rolling_upgrade_test.py
+++ b/tests/kafkatest/tests/client/consumer_rolling_upgrade_test.py
@@ -18,7 +18,7 @@
 
 
 from kafkatest.tests.verifiable_consumer_test import VerifiableConsumerTest
-from kafkatest.services.kafka import TopicPartition, quorum
+from kafkatest.services.kafka import TopicPartition, quorum, consumer_group
 
 class ConsumerRollingUpgradeTest(VerifiableConsumerTest):
     TOPIC = "test_topic"
@@ -56,7 +56,12 @@ def _verify_roundrobin_assignment(self, consumer):
         metadata_quorum=[quorum.isolated_kraft],
         use_new_coordinator=[True, False]
     )
-    def rolling_update_test(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    @matrix(
+        metadata_quorum=quorum.all_kraft,
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
+    )
+    def rolling_update_test(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify rolling updates of partition assignment strategies works correctly. In this
         test, we use a rolling restart to change the group's assignment strategy from "range" 
@@ -65,7 +70,7 @@ def rolling_update_test(self, metadata_quorum=quorum.zk, use_new_coordinator=Fal
         """
 
         # initialize the consumer using range assignment
-        consumer = self.setup_consumer(self.TOPIC, assignment_strategy=self.RANGE)
+        consumer = self.setup_consumer(self.TOPIC, assignment_strategy=self.RANGE, group_protocol=group_protocol)
 
         consumer.start()
         self.await_all_members(consumer)

From f66610095c6342260ea21f994db0570b1a3c5e51 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 22 Mar 2024 06:57:37 -0700
Subject: [PATCH 189/258] =?UTF-8?q?KAFKA-16274:=20Update=20replica=5Fscale?=
 =?UTF-8?q?=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20group=20protoco?=
 =?UTF-8?q?l=20config=20(#15577)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the methods involved.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 tests/kafkatest/tests/core/replica_scale_test.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/kafkatest/tests/core/replica_scale_test.py b/tests/kafkatest/tests/core/replica_scale_test.py
index 22555d202d3b5..8500c0bf302d7 100644
--- a/tests/kafkatest/tests/core/replica_scale_test.py
+++ b/tests/kafkatest/tests/core/replica_scale_test.py
@@ -20,7 +20,7 @@
 from kafkatest.services.trogdor.produce_bench_workload import ProduceBenchWorkloadService, ProduceBenchWorkloadSpec
 from kafkatest.services.trogdor.consume_bench_workload import ConsumeBenchWorkloadService, ConsumeBenchWorkloadSpec
 from kafkatest.services.trogdor.task_spec import TaskSpec
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.trogdor.trogdor import TrogdorService
 from kafkatest.services.zookeeper import ZookeeperService
 
@@ -52,7 +52,7 @@ def teardown(self):
         topic_count=[50],
         partition_count=[34],
         replication_factor=[3],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
@@ -60,10 +60,11 @@ def teardown(self):
         partition_count=[34],
         replication_factor=[3],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
     def test_produce_consume(self, topic_count, partition_count, replication_factor, 
-                             metadata_quorum=quorum.zk, use_new_coordinator=False):
+                             metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         topics_create_start_time = time.time()
         for i in range(topic_count):
             topic = "replicas_produce_consume_%d" % i
@@ -101,12 +102,13 @@ def test_produce_consume(self, topic_count, partition_count, replication_factor,
         produce_workload.wait_for_done(timeout_sec=600)
         print("Completed produce bench", flush=True)  # Force some stdout for Travis
 
+        consumer_conf = consumer_group.maybe_set_group_protocol(group_protocol)
         consume_spec = ConsumeBenchWorkloadSpec(0, TaskSpec.MAX_DURATION_MS,
                                                 consumer_workload_service.consumer_node,
                                                 consumer_workload_service.bootstrap_servers,
                                                 target_messages_per_sec=150000,
                                                 max_messages=3400000,
-                                                consumer_conf={},
+                                                consumer_conf=consumer_conf,
                                                 admin_client_conf={},
                                                 common_client_conf={},
                                                 active_topics=["replicas_produce_consume_[0-2]"])

From 159d25a7df25975694e2e0eb18a8feb125f7c39e Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Fri, 22 Mar 2024 06:58:40 -0700
Subject: [PATCH 190/258] =?UTF-8?q?KAFKA-16276:=20Update=20transactions=5F?=
 =?UTF-8?q?test.py=20to=20support=20KIP-848=E2=80=99s=20group=20protocol?=
 =?UTF-8?q?=20config=20(#15567)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the methods involved.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../kafkatest/tests/core/transactions_test.py | 33 ++++++++++++-------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/tests/kafkatest/tests/core/transactions_test.py b/tests/kafkatest/tests/core/transactions_test.py
index b9b39f355e46a..ae649c3766727 100644
--- a/tests/kafkatest/tests/core/transactions_test.py
+++ b/tests/kafkatest/tests/core/transactions_test.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 
 from kafkatest.services.zookeeper import ZookeeperService
-from kafkatest.services.kafka import KafkaService, quorum
+from kafkatest.services.kafka import KafkaService, quorum, consumer_group
 from kafkatest.services.console_consumer import ConsoleConsumer
 from kafkatest.services.verifiable_producer import VerifiableProducer
 from kafkatest.services.transactional_message_copier import TransactionalMessageCopier
@@ -85,8 +85,8 @@ def seed_messages(self, topic, num_seed_messages):
                    (self.num_seed_messages, seed_timeout_sec))
         return seed_producer.acked
 
-    def get_messages_from_topic(self, topic, num_messages):
-        consumer = self.start_consumer(topic, group_id="verifying_consumer")
+    def get_messages_from_topic(self, topic, num_messages, group_protocol):
+        consumer = self.start_consumer(topic, group_id="verifying_consumer", group_protocol=group_protocol)
         return self.drain_consumer(consumer, num_messages)
 
     def bounce_brokers(self, clean_shutdown):
@@ -154,7 +154,7 @@ def create_and_start_copiers(self, input_topic, output_topic, num_copiers, use_g
             ))
         return copiers
 
-    def start_consumer(self, topic_to_read, group_id):
+    def start_consumer(self, topic_to_read, group_id, group_protocol):
         consumer = ConsoleConsumer(context=self.test_context,
                                    num_nodes=1,
                                    kafka=self.kafka,
@@ -162,7 +162,8 @@ def start_consumer(self, topic_to_read, group_id):
                                    group_id=group_id,
                                    message_validator=is_int,
                                    from_beginning=True,
-                                   isolation_level="read_committed")
+                                   isolation_level="read_committed",
+                                   consumer_properties=consumer_group.maybe_set_group_protocol(group_protocol))
         consumer.start()
         # ensure that the consumer is up.
         wait_until(lambda: (len(consumer.messages_consumed[1]) > 0) == True,
@@ -189,7 +190,7 @@ def drain_consumer(self, consumer, num_messages):
     def copy_messages_transactionally(self, failure_mode, bounce_target,
                                       input_topic, output_topic,
                                       num_copiers, num_messages_to_copy,
-                                      use_group_metadata):
+                                      use_group_metadata, group_protocol):
         """Copies messages transactionally from the seeded input topic to the
         output topic, either bouncing brokers or clients in a hard and soft
         way as it goes.
@@ -204,7 +205,8 @@ def copy_messages_transactionally(self, failure_mode, bounce_target,
                                                 num_copiers=num_copiers,
                                                 use_group_metadata=use_group_metadata)
         concurrent_consumer = self.start_consumer(output_topic,
-                                                  group_id="concurrent_consumer")
+                                                  group_id="concurrent_consumer",
+                                                  group_protocol=group_protocol)
         clean_shutdown = False
         if failure_mode == "clean_bounce":
             clean_shutdown = True
@@ -257,9 +259,18 @@ def setup_topics(self):
         check_order=[True, False],
         use_group_metadata=[True, False],
         metadata_quorum=quorum.all_kraft,
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[False]
+    )
+    @matrix(
+        failure_mode=["hard_bounce", "clean_bounce"],
+        bounce_target=["brokers", "clients"],
+        check_order=[True, False],
+        use_group_metadata=[True, False],
+        metadata_quorum=quorum.all_kraft,
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_transactions(self, failure_mode, bounce_target, check_order, use_group_metadata, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_transactions(self, failure_mode, bounce_target, check_order, use_group_metadata, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         security_protocol = 'PLAINTEXT'
         self.kafka.security_protocol = security_protocol
         self.kafka.interbroker_security_protocol = security_protocol
@@ -283,8 +294,8 @@ def test_transactions(self, failure_mode, bounce_target, check_order, use_group_
         concurrently_consumed_messages = self.copy_messages_transactionally(
             failure_mode, bounce_target, input_topic=self.input_topic,
             output_topic=self.output_topic, num_copiers=self.num_input_partitions,
-            num_messages_to_copy=self.num_seed_messages, use_group_metadata=use_group_metadata)
-        output_messages = self.get_messages_from_topic(self.output_topic, self.num_seed_messages)
+            num_messages_to_copy=self.num_seed_messages, use_group_metadata=use_group_metadata, group_protocol=group_protocol)
+        output_messages = self.get_messages_from_topic(self.output_topic, self.num_seed_messages, group_protocol)
 
         concurrently_consumed_message_set = set(concurrently_consumed_messages)
         output_message_set = set(output_messages)

From 2e8d69b78ca52196decd851c8520798aa856c073 Mon Sep 17 00:00:00 2001
From: Sanskar Jhajharia <122860866+sjhajharia@users.noreply.github.com>
Date: Fri, 22 Mar 2024 23:56:07 +0530
Subject: [PATCH 191/258] KAFKA-16314: Introducing the
 AbortableTransactionException (#15486)

As a part of KIP-890, we are introducing a new class of Exceptions which when encountered shall lead to Aborting the ongoing Transaction. The following PR introduces the same with client side handling and server side changes.

On client Side, the code attempts to handle the exception as an Abortable error and ensure that it doesn't take the producer to a fatal state. For each of the Transactional APIs, we have added the appropriate handling. For the produce request, we have verified that the exception transitions the state to Aborted.
On the server side, we have bumped the ProduceRequest, ProduceResponse, TxnOffestCommitRequest and TxnOffsetCommitResponse Version. The appropriate handling on the server side has been added to ensure that the new error case is sent back only for the new clients. The older clients will continue to get the old Invalid_txn_state exception to maintain backward compatibility.

Reviewers: Calvin Liu <caliu@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../internals/TransactionManager.java         |  30 +++--
 .../errors/AbortableTransactionException.java |  23 ++++
 .../apache/kafka/common/protocol/Errors.java  |   4 +-
 .../message/AddOffsetsToTxnRequest.json       |   4 +-
 .../message/AddOffsetsToTxnResponse.json      |   4 +-
 .../message/AddPartitionsToTxnRequest.json    |   4 +-
 .../message/AddPartitionsToTxnResponse.json   |   4 +-
 .../common/message/EndTxnRequest.json         |   4 +-
 .../common/message/EndTxnResponse.json        |   4 +-
 .../message/FindCoordinatorRequest.json       |   4 +-
 .../message/FindCoordinatorResponse.json      |   4 +-
 .../common/message/InitProducerIdRequest.json |   4 +-
 .../message/InitProducerIdResponse.json       |   4 +-
 .../common/message/ProduceRequest.json        |   4 +-
 .../common/message/ProduceResponse.json       |   4 +-
 .../message/TxnOffsetCommitRequest.json       |   4 +-
 .../message/TxnOffsetCommitResponse.json      |   4 +-
 .../producer/internals/SenderTest.java        |  40 ++++++
 .../internals/TransactionManagerTest.java     | 121 ++++++++++++++++++
 .../group/CoordinatorPartitionWriter.scala    |   9 +-
 .../coordinator/group/GroupCoordinator.scala  |   8 +-
 .../group/GroupCoordinatorAdapter.scala       |   3 +-
 .../transaction/TransactionCoordinator.scala  |   2 +-
 .../server/AddPartitionsToTxnManager.scala    |  28 +++-
 .../main/scala/kafka/server/KafkaApis.scala   |   5 +-
 .../scala/kafka/server/ReplicaManager.scala   |  33 +++--
 .../AbstractCoordinatorConcurrencyTest.scala  |   3 +-
 .../CoordinatorPartitionWriterTest.scala      |   8 +-
 .../group/GroupCoordinatorAdapterTest.scala   |   3 +-
 .../GroupCoordinatorConcurrencyTest.scala     |   6 +-
 .../group/GroupCoordinatorTest.scala          |   7 +-
 .../TransactionCoordinatorTest.scala          |   4 +-
 .../AddPartitionsToTxnManagerTest.scala       |  66 +++++++---
 .../AddPartitionsToTxnRequestServerTest.scala |   2 +-
 .../unit/kafka/server/KafkaApisTest.scala     |   5 +
 .../kafka/server/ReplicaManagerTest.scala     |  44 ++++---
 .../group/GroupCoordinatorService.java        |   3 +-
 .../group/runtime/CoordinatorRuntime.java     |   7 +-
 .../group/runtime/PartitionWriter.java        |   4 +-
 .../group/GroupCoordinatorServiceTest.java    |   2 +
 .../group/runtime/CoordinatorRuntimeTest.java |  26 +++-
 .../runtime/InMemoryPartitionWriter.java      |   3 +-
 42 files changed, 445 insertions(+), 110 deletions(-)
 create mode 100644 clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java

diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
index 14a519a7a0f4f..4ff54b1759875 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
@@ -1330,6 +1330,8 @@ public void handleResponse(AbstractResponse response) {
                 // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
                 // just treat it the same as PRODUCE_FENCED.
                 fatalError(Errors.PRODUCER_FENCED.exception());
+            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unexpected error in InitProducerIdResponse; " + error.message()));
             }
@@ -1386,10 +1388,8 @@ public void handleResponse(AbstractResponse response) {
                     // just treat it the same as PRODUCE_FENCED.
                     fatalError(Errors.PRODUCER_FENCED.exception());
                     return;
-                } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
-                    fatalError(error.exception());
-                    return;
-                } else if (error == Errors.INVALID_TXN_STATE) {
+                } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED ||
+                        error == Errors.INVALID_TXN_STATE) {
                     fatalError(error.exception());
                     return;
                 } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
@@ -1401,6 +1401,9 @@ public void handleResponse(AbstractResponse response) {
                 } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
                     abortableErrorIfPossible(error.exception());
                     return;
+                } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                    abortableError(error.exception());
+                    return;
                 } else {
                     log.error("Could not add partition {} due to unexpected error {}", topicPartition, error);
                     hasPartitionErrors = true;
@@ -1504,6 +1507,8 @@ public void handleResponse(AbstractResponse response) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                 abortableError(GroupAuthorizationException.forGroupId(key));
+            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                abortableError(error.exception());
             } else {
                 fatalError(new KafkaException(String.format("Could not find a coordinator with type %s with key %s due to " +
                         "unexpected error: %s", coordinatorType, key,
@@ -1552,12 +1557,13 @@ public void handleResponse(AbstractResponse response) {
                 // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
                 // just treat it the same as PRODUCE_FENCED.
                 fatalError(Errors.PRODUCER_FENCED.exception());
-            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
-                fatalError(error.exception());
-            } else if (error == Errors.INVALID_TXN_STATE) {
+            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED ||
+                    error == Errors.INVALID_TXN_STATE) {
                 fatalError(error.exception());
             } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
                 abortableErrorIfPossible(error.exception());
+            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unhandled error in EndTxnResponse: " + error.message()));
             }
@@ -1611,12 +1617,13 @@ public void handleResponse(AbstractResponse response) {
                 // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
                 // just treat it the same as PRODUCE_FENCED.
                 fatalError(Errors.PRODUCER_FENCED.exception());
-            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
-                fatalError(error.exception());
-            } else if (error == Errors.INVALID_TXN_STATE) {
+            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED ||
+                    error == Errors.INVALID_TXN_STATE) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                 abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
+            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unexpected error in AddOffsetsToTxnResponse: " + error.message()));
             }
@@ -1679,7 +1686,8 @@ public void handleResponse(AbstractResponse response) {
                 } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                     abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
                     break;
-                } else if (error == Errors.FENCED_INSTANCE_ID) {
+                } else if (error == Errors.FENCED_INSTANCE_ID ||
+                        error == Errors.ABORTABLE_TRANSACTION) {
                     abortableError(error.exception());
                     break;
                 } else if (error == Errors.UNKNOWN_MEMBER_ID
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java b/clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java
new file mode 100644
index 0000000000000..cfea649e750c6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.errors;
+
+public class AbortableTransactionException extends ApiException {
+    public AbortableTransactionException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
index 610b1b66e7633..30bc250c6dec2 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
@@ -138,6 +138,7 @@
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedSaslMechanismException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.errors.AbortableTransactionException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -392,7 +393,8 @@ public enum Errors {
     UNKNOWN_CONTROLLER_ID(116, "This controller ID is not known.", UnknownControllerIdException::new),
     UNKNOWN_SUBSCRIPTION_ID(117, "Client sent a push telemetry request with an invalid or outdated subscription ID.", UnknownSubscriptionIdException::new),
     TELEMETRY_TOO_LARGE(118, "Client sent a push telemetry request larger than the maximum size the broker will accept.", TelemetryTooLargeException::new),
-    INVALID_REGISTRATION(119, "The controller has considered the broker registration to be invalid.", InvalidRegistrationException::new);
+    INVALID_REGISTRATION(119, "The controller has considered the broker registration to be invalid.", InvalidRegistrationException::new),
+    ABORTABLE_TRANSACTION(120, "The server encountered an error with the transaction. The client can abort the transaction to continue using this transactional ID.", AbortableTransactionException::new);
 
     private static final Logger log = LoggerFactory.getLogger(Errors.class);
 
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
index ade3fc72c9a51..9d7b63c313332 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
@@ -23,7 +23,9 @@
   // Version 2 adds the support for new error code PRODUCER_FENCED.
   //
   // Version 3 enables flexible versions.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
index 71fa655227081..6b3b1c481d663 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
@@ -22,7 +22,9 @@
   // Version 2 adds the support for new error code PRODUCER_FENCED.
   //
   // Version 3 enables flexible versions.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
index 1b89c54d86403..139d1436a6520 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
@@ -25,9 +25,11 @@
   // Version 3 enables flexible versions.
   //
   // Version 4 adds VerifyOnly field to check if partitions are already in transaction and adds support to batch multiple transactions.
+  //
+  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
   // Versions 3 and below will be exclusively used by clients and versions 4 and above will be used by brokers.
   "latestVersionUnstable": false,
-  "validVersions": "0-4",
+  "validVersions": "0-5",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "Transactions", "type": "[]AddPartitionsToTxnTransaction", "versions":  "4+",
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
index 326b4acdb446a..a2af388dba55f 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
@@ -24,7 +24,9 @@
   // Version 3 enables flexible versions.
   //
   // Version 4 adds support to batch multiple transactions and a top level error code.
-  "validVersions": "0-4",
+  //
+  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-5",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/EndTxnRequest.json b/clients/src/main/resources/common/message/EndTxnRequest.json
index f16ef76246d35..5bf6c577343d6 100644
--- a/clients/src/main/resources/common/message/EndTxnRequest.json
+++ b/clients/src/main/resources/common/message/EndTxnRequest.json
@@ -23,7 +23,9 @@
   // Version 2 adds the support for new error code PRODUCER_FENCED.
   //
   // Version 3 enables flexible versions.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
diff --git a/clients/src/main/resources/common/message/EndTxnResponse.json b/clients/src/main/resources/common/message/EndTxnResponse.json
index 3071953185ffc..53b0250f808b9 100644
--- a/clients/src/main/resources/common/message/EndTxnResponse.json
+++ b/clients/src/main/resources/common/message/EndTxnResponse.json
@@ -22,7 +22,9 @@
   // Version 2 adds the support for new error code PRODUCER_FENCED.
   //
   // Version 3 enables flexible versions.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/FindCoordinatorRequest.json b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
index 49ae7e92f0bb2..e6786f5b10e3e 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorRequest.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
@@ -25,7 +25,9 @@
   // Version 3 is the first flexible version.
   //
   // Version 4 adds support for batching via CoordinatorKeys (KIP-699)
-  "validVersions": "0-4",
+  //
+  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-5",
   "deprecatedVersions": "0",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/FindCoordinatorResponse.json b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
index 9309c0177d67e..a744a1928d6c7 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorResponse.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
@@ -24,7 +24,9 @@
   // Version 3 is the first flexible version.
   //
   // Version 4 adds support for batching via Coordinators (KIP-699)
-  "validVersions": "0-4",
+  //
+  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-5",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "1+", "ignorable": true,
diff --git a/clients/src/main/resources/common/message/InitProducerIdRequest.json b/clients/src/main/resources/common/message/InitProducerIdRequest.json
index 4e75352db6f33..92e6a6b25376d 100644
--- a/clients/src/main/resources/common/message/InitProducerIdRequest.json
+++ b/clients/src/main/resources/common/message/InitProducerIdRequest.json
@@ -25,7 +25,9 @@
   // Version 3 adds ProducerId and ProducerEpoch, allowing producers to try to resume after an INVALID_PRODUCER_EPOCH error
   //
   // Version 4 adds the support for new error code PRODUCER_FENCED.
-  "validVersions": "0-4",
+  //
+  // Verison 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-5",
   "flexibleVersions": "2+",
   "fields": [
     { "name": "TransactionalId", "type": "string", "versions": "0+", "nullableVersions": "0+", "entityType": "transactionalId",
diff --git a/clients/src/main/resources/common/message/InitProducerIdResponse.json b/clients/src/main/resources/common/message/InitProducerIdResponse.json
index f56c2fe67e3da..c0f10b2e851a6 100644
--- a/clients/src/main/resources/common/message/InitProducerIdResponse.json
+++ b/clients/src/main/resources/common/message/InitProducerIdResponse.json
@@ -24,7 +24,9 @@
   // Version 3 is the same as version 2.
   //
   // Version 4 adds the support for new error code PRODUCER_FENCED.
-  "validVersions": "0-4",
+  //
+  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-5",
   "flexibleVersions": "2+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+", "ignorable": true,
diff --git a/clients/src/main/resources/common/message/ProduceRequest.json b/clients/src/main/resources/common/message/ProduceRequest.json
index 3451a8471df3e..d396d66070da9 100644
--- a/clients/src/main/resources/common/message/ProduceRequest.json
+++ b/clients/src/main/resources/common/message/ProduceRequest.json
@@ -35,7 +35,9 @@
   // Version 9 enables flexible versions.
   //
   // Version 10 is the same as version 9 (KIP-951).
-  "validVersions": "0-10",
+  //
+  // Version 11 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-11",
   "deprecatedVersions": "0-6",
   "flexibleVersions": "9+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/ProduceResponse.json b/clients/src/main/resources/common/message/ProduceResponse.json
index d294fb8aa2efe..1c097d3bc3b5d 100644
--- a/clients/src/main/resources/common/message/ProduceResponse.json
+++ b/clients/src/main/resources/common/message/ProduceResponse.json
@@ -34,7 +34,9 @@
   // Version 9 enables flexible versions.
   //
   // Version 10 adds 'CurrentLeader' and 'NodeEndpoints' as tagged fields (KIP-951)
-  "validVersions": "0-10",
+  //
+  // Version 11 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-11",
   "flexibleVersions": "9+",
   "fields": [
     { "name": "Responses", "type": "[]TopicProduceResponse", "versions": "0+",
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
index 0e7b187542040..1df18c64b52d8 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
@@ -23,7 +23,9 @@
   // Version 2 adds the committed leader epoch.
   //
   // Version 3 adds the member.id, group.instance.id and generation.id.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
index 96b03a078d877..0f6c1f272418a 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
@@ -22,7 +22,9 @@
   // Version 2 is the same as version 1.
   //
   // Version 3 adds illegal generation, fenced instance id, and unknown member id errors.
-  "validVersions": "0-3",
+  //
+  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
     { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+",
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
index 29f0b7b1a8faa..b9c734a2cb1bf 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
@@ -42,6 +42,7 @@
 import org.apache.kafka.common.errors.TransactionAbortedException;
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.errors.AbortableTransactionException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.AddPartitionsToTxnResponseData;
 import org.apache.kafka.common.message.ApiMessageType;
@@ -3146,6 +3147,45 @@ public void testInvalidTxnStateIsAnAbortableError() throws Exception {
 
         txnManager.beginTransaction();
     }
+
+    @Test
+    public void testAbortableTxnExceptionIsAnAbortableError() throws Exception {
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+        apiVersions.update("0", NodeApiVersions.create(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 3));
+        TransactionManager txnManager = new TransactionManager(logContext, "textAbortableTxnException", 60000, 100, apiVersions);
+
+        setupWithTransactionState(txnManager);
+        doInitTransactions(txnManager, producerIdAndEpoch);
+
+        txnManager.beginTransaction();
+        txnManager.maybeAddPartition(tp0);
+        client.prepareResponse(buildAddPartitionsToTxnResponseData(0, Collections.singletonMap(tp0, Errors.NONE)));
+        sender.runOnce();
+
+        Future<RecordMetadata> request = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+        sendIdempotentProducerResponse(0, tp0, Errors.ABORTABLE_TRANSACTION, -1);
+
+        // Return AbortableTransactionException error. It should be abortable.
+        sender.runOnce();
+        assertFutureFailure(request, AbortableTransactionException.class);
+        assertTrue(txnManager.hasAbortableError());
+        TransactionalRequestResult result = txnManager.beginAbort();
+        sender.runOnce();
+
+        // Once the transaction is aborted, we should be able to begin a new one.
+        respondToEndTxn(Errors.NONE);
+        sender.runOnce();
+        assertTrue(txnManager::isInitializing);
+        prepareInitProducerResponse(Errors.NONE, producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
+        sender.runOnce();
+        assertTrue(txnManager::isReady);
+
+        assertTrue(result.isSuccessful());
+        result.await();
+
+        txnManager.beginTransaction();
+    }
     @Test
     public void testProducerBatchRetriesWhenPartitionLeaderChanges() throws Exception {
         Metrics m = new Metrics();
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index bf9d77a395d2a..4d8c445be0ac5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -40,6 +40,7 @@
 import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.errors.AbortableTransactionException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
@@ -3514,6 +3515,126 @@ public void testForegroundInvalidStateTransitionIsRecoverable() {
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
+    @Test
+    public void testAbortableTransactionExceptionInInitProducerId() {
+        TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
+        assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
+
+        prepareInitPidResponse(Errors.ABORTABLE_TRANSACTION, false, producerId, RecordBatch.NO_PRODUCER_EPOCH);
+        runUntil(transactionManager::hasError);
+        assertTrue(initPidResult.isCompleted());
+        assertFalse(initPidResult.isSuccessful());
+        assertThrows(AbortableTransactionException.class, initPidResult::await);
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
+    @Test
+    public void testAbortableTransactionExceptionInAddPartitions() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.maybeAddPartition(tp);
+
+        prepareAddPartitionsToTxn(tp, Errors.ABORTABLE_TRANSACTION);
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
+    @Test
+    public void testAbortableTransactionExceptionInFindCoordinator() {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+                singletonMap(new TopicPartition("foo", 0), new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
+
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
+
+        prepareFindCoordinatorResponse(Errors.ABORTABLE_TRANSACTION, false, CoordinatorType.GROUP, consumerGroupId);
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+
+        runUntil(sendOffsetsResult::isCompleted);
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
+
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
+    @Test
+    public void testAbortableTransactionExceptionInEndTxn() throws InterruptedException {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.maybeAddPartition(tp0);
+        TransactionalRequestResult commitResult = transactionManager.beginCommit();
+
+        Future<RecordMetadata> responseFuture = appendToAccumulator(tp0);
+
+        assertFalse(responseFuture.isDone());
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        prepareEndTxnResponse(Errors.ABORTABLE_TRANSACTION, TransactionResult.COMMIT, producerId, epoch);
+
+        runUntil(commitResult::isCompleted);
+        runUntil(responseFuture::isDone);
+
+        assertThrows(KafkaException.class, commitResult::await);
+        assertFalse(commitResult.isSuccessful());
+        assertTrue(commitResult.isAcked());
+
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
+    @Test
+    public void testAbortableTransactionExceptionInAddOffsetsToTxn() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+                singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
+
+        prepareAddOffsetsToTxnResponse(Errors.ABORTABLE_TRANSACTION, consumerGroupId, producerId, epoch);
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(sendOffsetsResult.isCompleted());
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
+
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
+    @Test
+    public void testAbortableTransactionExceptionInTxnOffsetCommit() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+                singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
+
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.ABORTABLE_TRANSACTION));
+        runUntil(transactionManager::hasError);
+
+        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(sendOffsetsResult.isCompleted());
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
+        assertAbortableError(AbortableTransactionException.class);
+    }
+
     private FutureRecordMetadata appendToAccumulator(TopicPartition tp) throws InterruptedException {
         final long nowMs = time.milliseconds();
         return accumulator.append(tp.topic(), tp.partition(), nowMs, "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS,
diff --git a/core/src/main/scala/kafka/coordinator/group/CoordinatorPartitionWriter.scala b/core/src/main/scala/kafka/coordinator/group/CoordinatorPartitionWriter.scala
index 59af648c21b5b..b912881847602 100644
--- a/core/src/main/scala/kafka/coordinator/group/CoordinatorPartitionWriter.scala
+++ b/core/src/main/scala/kafka/coordinator/group/CoordinatorPartitionWriter.scala
@@ -17,7 +17,7 @@
 package kafka.coordinator.group
 
 import kafka.cluster.PartitionListener
-import kafka.server.{ActionQueue, ReplicaManager, RequestLocal}
+import kafka.server.{ActionQueue, ReplicaManager, RequestLocal, defaultError, genericError}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.RecordTooLargeException
 import org.apache.kafka.common.protocol.Errors
@@ -188,8 +188,10 @@ class CoordinatorPartitionWriter[T](
     tp: TopicPartition,
     transactionalId: String,
     producerId: Long,
-    producerEpoch: Short
+    producerEpoch: Short,
+    apiVersion: Short
   ): CompletableFuture[VerificationGuard] = {
+    val supportedOperation = if (apiVersion >= 4) genericError else defaultError
     val future = new CompletableFuture[VerificationGuard]()
     replicaManager.maybeStartTransactionVerificationForPartition(
       topicPartition = tp,
@@ -204,7 +206,8 @@ class CoordinatorPartitionWriter[T](
         } else {
           future.complete(verificationGuard)
         }
-      }
+      },
+      supportedOperation
     )
     future
   }
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
index 3136d2c8e64cc..6d234dccc8d99 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
@@ -905,7 +905,8 @@ private[group] class GroupCoordinator(
                              generationId: Int,
                              offsetMetadata: immutable.Map[TopicIdPartition, OffsetAndMetadata],
                              responseCallback: immutable.Map[TopicIdPartition, Errors] => Unit,
-                             requestLocal: RequestLocal = RequestLocal.NoCaching): Unit = {
+                             requestLocal: RequestLocal = RequestLocal.NoCaching,
+                             apiVersion: Short): Unit = {
     validateGroupStatus(groupId, ApiKeys.TXN_OFFSET_COMMIT) match {
       case Some(error) => responseCallback(offsetMetadata.map { case (k, _) => k -> error })
       case None =>
@@ -928,7 +929,7 @@ private[group] class GroupCoordinator(
               offsetTopicPartition, offsetMetadata, newRequestLocal, responseCallback, Some(verificationGuard))
           }
         }
-
+        val supportedOperation = if (apiVersion >= 4) genericError else defaultError
         groupManager.replicaManager.maybeStartTransactionVerificationForPartition(
           topicPartition = offsetTopicPartition,
           transactionalId,
@@ -941,7 +942,8 @@ private[group] class GroupCoordinator(
           KafkaRequestHandler.wrapAsyncCallback(
             postVerificationCallback,
             requestLocal
-          )
+          ),
+          supportedOperation
         )
     }
   }
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala
index 6b3cf072c807a..53d7c9b5e92e4 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala
@@ -476,7 +476,8 @@ private[group] class GroupCoordinatorAdapter(
       request.generationId,
       partitions.toMap,
       callback,
-      RequestLocal(bufferSupplier)
+      RequestLocal(bufferSupplier),
+      context.apiVersion()
     )
 
     future
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
index c7f5a16b0662e..4d9ba7fa55893 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
@@ -369,7 +369,7 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
                   if (txnMetadata.topicPartitions.contains(part))
                     (part, Errors.NONE)
                   else
-                    (part, Errors.INVALID_TXN_STATE)
+                    (part, Errors.ABORTABLE_TRANSACTION)
                 }.toMap)
               }
             }
diff --git a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
index 8604c97ecc303..ec909580a8400 100644
--- a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
+++ b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
@@ -42,6 +42,17 @@ object AddPartitionsToTxnManager {
   val VerificationTimeMsMetricName = "VerificationTimeMs"
 }
 
+/**
+ * This is an enum which handles the Partition Response based on the Request Version and the exact operation
+ *    defaultError:       This is the default workflow which maps to cases when the Produce Request Version or the Txn_offset_commit request was lower than the first version supporting the new Error Class
+ *    genericError:       This maps to the case when the clients are updated to handle the AbortableTxnException
+ *    addPartition:       This is a WIP. To be updated as a part of KIP-890 Part 2
+ */
+sealed trait SupportedOperation
+case object defaultError extends SupportedOperation
+case object genericError extends SupportedOperation
+case object addPartition extends SupportedOperation
+
 /*
  * Data structure to hold the transactional data to send to a node. Note -- at most one request per transactional ID
  * will exist at a time in the map. If a given transactional ID exists in the map, and a new request with the same ID
@@ -49,7 +60,8 @@ object AddPartitionsToTxnManager {
  */
 class TransactionDataAndCallbacks(val transactionData: AddPartitionsToTxnTransactionCollection,
                                   val callbacks: mutable.Map[String, AddPartitionsToTxnManager.AppendCallback],
-                                  val startTimeMs: mutable.Map[String, Long])
+                                  val startTimeMs: mutable.Map[String, Long],
+                                  val supportedOperation: SupportedOperation)
 
 class AddPartitionsToTxnManager(
   config: KafkaConfig,
@@ -79,7 +91,8 @@ class AddPartitionsToTxnManager(
     producerId: Long,
     producerEpoch: Short,
     topicPartitions: Seq[TopicPartition],
-    callback: AddPartitionsToTxnManager.AppendCallback
+    callback: AddPartitionsToTxnManager.AppendCallback,
+    supportedOperation: SupportedOperation
   ): Unit = {
     val coordinatorNode = getTransactionCoordinator(partitionFor(transactionalId))
     if (coordinatorNode.isEmpty) {
@@ -99,14 +112,16 @@ class AddPartitionsToTxnManager(
         .setVerifyOnly(true)
         .setTopics(topicCollection)
 
-      addTxnData(coordinatorNode.get, transactionData, callback)
+      addTxnData(coordinatorNode.get, transactionData, callback, supportedOperation)
+
     }
   }
 
   private def addTxnData(
     node: Node,
     transactionData: AddPartitionsToTxnTransaction,
-    callback: AddPartitionsToTxnManager.AppendCallback
+    callback: AddPartitionsToTxnManager.AppendCallback,
+    supportedOperation: SupportedOperation
   ): Unit = {
     nodesToTransactions.synchronized {
       val curTime = time.milliseconds()
@@ -115,7 +130,8 @@ class AddPartitionsToTxnManager(
         new TransactionDataAndCallbacks(
           new AddPartitionsToTxnTransactionCollection(1),
           mutable.Map[String, AddPartitionsToTxnManager.AppendCallback](),
-          mutable.Map[String, Long]()))
+          mutable.Map[String, Long](),
+          supportedOperation))
 
       val existingTransactionData = existingNodeAndTransactionData.transactionData.find(transactionData.transactionalId)
 
@@ -210,6 +226,8 @@ class AddPartitionsToTxnManager(
                   val code =
                     if (partitionResult.partitionErrorCode == Errors.PRODUCER_FENCED.code)
                       Errors.INVALID_PRODUCER_EPOCH.code
+                    else if (partitionResult.partitionErrorCode() == Errors.ABORTABLE_TRANSACTION.code && transactionDataAndCallbacks.supportedOperation != genericError) // For backward compatibility with clients.
+                      Errors.INVALID_TXN_STATE.code
                     else
                       partitionResult.partitionErrorCode
                   unverified.put(tp, Errors.forCode(code))
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index e376b053e44f8..3b76696f199f1 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -721,7 +721,7 @@ class KafkaApis(val requestChannel: RequestChannel,
       sendResponseCallback(Map.empty)
     else {
       val internalTopicsAllowed = request.header.clientId == AdminUtils.ADMIN_CLIENT_ID
-
+      val supportedOperation = if (request.header.apiVersion > 10) genericError else defaultError
       // call the replica manager to append messages to the replicas
       replicaManager.handleProduceAppend(
         timeout = produceRequest.timeout.toLong,
@@ -731,7 +731,8 @@ class KafkaApis(val requestChannel: RequestChannel,
         entriesPerPartition = authorizedRequestInfo,
         responseCallback = sendResponseCallback,
         recordValidationStatsCallback = processingStatsCallback,
-        requestLocal = requestLocal)
+        requestLocal = requestLocal,
+        supportedOperation = supportedOperation)
 
       // if the request is put into the purgatory, it will have a held reference and hence cannot be garbage collected;
       // hence we clear its data here in order to let GC reclaim its memory since it is already appended to log
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 2a1cbf2d2bbae..ab9003b0a4e18 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -795,6 +795,7 @@ class ReplicaManager(val config: KafkaConfig,
    * @param requestLocal                  container for the stateful instances scoped to this request -- this must correspond to the
    *                                      thread calling this method
    * @param actionQueue                   the action queue to use. ReplicaManager#defaultActionQueue is used by default.
+   * @param supportedOperation            determines the supported Operation based on the client's Request api version
    *
    * The responseCallback is wrapped so that it is scheduled on a request handler thread. There, it should be called with
    * that request handler thread's thread local and not the one supplied to this method.
@@ -807,7 +808,8 @@ class ReplicaManager(val config: KafkaConfig,
                           responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
                           recordValidationStatsCallback: Map[TopicPartition, RecordValidationStats] => Unit = _ => (),
                           requestLocal: RequestLocal = RequestLocal.NoCaching,
-                          actionQueue: ActionQueue = this.defaultActionQueue): Unit = {
+                          actionQueue: ActionQueue = this.defaultActionQueue,
+                          supportedOperation: SupportedOperation): Unit = {
 
     val transactionalProducerInfo = mutable.HashSet[(Long, Short)]()
     val topicPartitionBatchInfo = mutable.Map[TopicPartition, Int]()
@@ -884,7 +886,8 @@ class ReplicaManager(val config: KafkaConfig,
       KafkaRequestHandler.wrapAsyncCallback(
         postVerificationCallback,
         requestLocal
-      )
+      ),
+      supportedOperation
     )
   }
 
@@ -975,12 +978,13 @@ class ReplicaManager(val config: KafkaConfig,
 
   /**
    *
-   * @param topicPartition  the topic partition to maybe verify
-   * @param transactionalId the transactional id for the transaction
-   * @param producerId      the producer id for the producer writing to the transaction
-   * @param producerEpoch   the epoch of the producer writing to the transaction
-   * @param baseSequence    the base sequence of the first record in the batch we are trying to append
-   * @param callback        the method to execute once the verification is either completed or returns an error
+   * @param topicPartition        the topic partition to maybe verify
+   * @param transactionalId       the transactional id for the transaction
+   * @param producerId            the producer id for the producer writing to the transaction
+   * @param producerEpoch         the epoch of the producer writing to the transaction
+   * @param baseSequence          the base sequence of the first record in the batch we are trying to append
+   * @param callback              the method to execute once the verification is either completed or returns an error
+   * @param supportedOperation    determines the supported operation based on the client's Request API version
    *
    * When the verification returns, the callback will be supplied the error if it exists or Errors.NONE.
    * If the verification guard exists, it will also be supplied. Otherwise the SENTINEL verification guard will be returned.
@@ -992,7 +996,8 @@ class ReplicaManager(val config: KafkaConfig,
     producerId: Long,
     producerEpoch: Short,
     baseSequence: Int,
-    callback: ((Errors, VerificationGuard)) => Unit
+    callback: ((Errors, VerificationGuard)) => Unit,
+    supportedOperation: SupportedOperation
   ): Unit = {
     def generalizedCallback(results: (Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])): Unit = {
       val (preAppendErrors, verificationGuards) = results
@@ -1007,7 +1012,8 @@ class ReplicaManager(val config: KafkaConfig,
       transactionalId,
       producerId,
       producerEpoch,
-      generalizedCallback
+      generalizedCallback,
+      supportedOperation
     )
   }
 
@@ -1018,6 +1024,7 @@ class ReplicaManager(val config: KafkaConfig,
    * @param producerId               the producer id for the producer writing to the transaction
    * @param producerEpoch            the epoch of the producer writing to the transaction
    * @param callback                 the method to execute once the verification is either completed or returns an error
+   * @param supportedOperation       determines the supported operation based on the client's Request API version
    *
    * When the verification returns, the callback will be supplied the errors per topic partition if there were errors.
    * The callback will also be supplied the verification guards per partition if they exist. It is possible to have an
@@ -1029,7 +1036,8 @@ class ReplicaManager(val config: KafkaConfig,
     transactionalId: String,
     producerId: Long,
     producerEpoch: Short,
-    callback: ((Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])) => Unit
+    callback: ((Map[TopicPartition, Errors], Map[TopicPartition, VerificationGuard])) => Unit,
+    supportedOperation: SupportedOperation
   ): Unit = {
     // Skip verification if the request is not transactional or transaction verification is disabled.
     if (transactionalId == null ||
@@ -1075,7 +1083,8 @@ class ReplicaManager(val config: KafkaConfig,
       producerId = producerId,
       producerEpoch = producerEpoch,
       topicPartitions = verificationGuards.keys.toSeq,
-      callback = invokeCallback
+      callback = invokeCallback,
+      supportedOperation = supportedOperation
     ))
 
   }
diff --git a/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
index 533f67aa75c1f..956db24c010ec 100644
--- a/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
@@ -200,7 +200,8 @@ object AbstractCoordinatorConcurrencyTest {
       producerId: Long,
       producerEpoch: Short,
       baseSequence: Int,
-      callback: ((Errors, VerificationGuard)) => Unit
+      callback: ((Errors, VerificationGuard)) => Unit,
+      supportedOperation: SupportedOperation
     ): Unit = {
       // Skip verification
       callback((Errors.NONE, VerificationGuard.SENTINEL))
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/CoordinatorPartitionWriterTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/CoordinatorPartitionWriterTest.scala
index 9e9ba7426078a..7a3c5f194f9a5 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/CoordinatorPartitionWriterTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/CoordinatorPartitionWriterTest.scala
@@ -21,7 +21,7 @@ import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.config.TopicConfig
 import org.apache.kafka.common.errors.{NotLeaderOrFollowerException, RecordTooLargeException}
-import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.{CompressionType, ControlRecordType, MemoryRecords, RecordBatch}
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.TransactionResult
@@ -343,7 +343,8 @@ class CoordinatorPartitionWriterTest {
       ArgumentMatchers.eq(10L),
       ArgumentMatchers.eq(5.toShort),
       ArgumentMatchers.eq(RecordBatch.NO_SEQUENCE),
-      callbackCapture.capture()
+      callbackCapture.capture(),
+      ArgumentMatchers.any()
     )).thenAnswer(_ => {
       callbackCapture.getValue.apply((
         error,
@@ -355,7 +356,8 @@ class CoordinatorPartitionWriterTest {
       tp,
       "transactional-id",
       10L,
-      5.toShort
+      5.toShort,
+      ApiKeys.TXN_OFFSET_COMMIT.latestVersion()
     )
 
     if (error == Errors.NONE) {
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorAdapterTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorAdapterTest.scala
index 14b41b09f8b7b..647d6072fc48c 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorAdapterTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorAdapterTest.scala
@@ -761,7 +761,8 @@ class GroupCoordinatorAdapterTest {
         )
       )),
       capturedCallback.capture(),
-      ArgumentMatchers.eq(RequestLocal(bufferSupplier))
+      ArgumentMatchers.eq(RequestLocal(bufferSupplier)),
+      ArgumentMatchers.any()
     )
 
     capturedCallback.getValue.apply(Map(
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
index be0038c027e99..9bed515aed725 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
@@ -24,13 +24,13 @@ import kafka.common.OffsetAndMetadata
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest
 import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
 import kafka.coordinator.group.GroupCoordinatorConcurrencyTest._
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig, KafkaRequestHandler}
+import kafka.server.{DelayedOperationPurgatory, KafkaConfig, KafkaRequestHandler, RequestLocal}
 import kafka.utils.CoreUtils
 import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
-import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests.{JoinGroupRequest, OffsetFetchResponse}
 import org.apache.kafka.common.utils.Time
 import org.junit.jupiter.api.Assertions._
@@ -319,7 +319,7 @@ class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest
         // Since the replica manager is mocked we can use a dummy value for transactionalId.
         groupCoordinator.handleTxnCommitOffsets(member.group.groupId, "dummy-txn-id", producerId, producerEpoch,
           JoinGroupRequest.UNKNOWN_MEMBER_ID, Option.empty, JoinGroupRequest.UNKNOWN_GENERATION_ID,
-          offsets, callbackWithTxnCompletion)
+          offsets, callbackWithTxnCompletion, RequestLocal.NoCaching, ApiKeys.TXN_OFFSET_COMMIT.latestVersion())
         replicaManager.tryCompleteActions()
       } finally lock.foreach(_.unlock())
     }
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
index 7930a798c07be..c8417f678670b 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
@@ -22,7 +22,7 @@ import kafka.common.OffsetAndMetadata
 import kafka.server.{ActionQueue, DelayedOperationPurgatory, HostedPartition, KafkaConfig, KafkaRequestHandler, ReplicaManager, RequestLocal}
 import kafka.utils._
 import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid}
-import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.record.{MemoryRecords, RecordBatch}
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.{JoinGroupRequest, OffsetCommitRequest, OffsetFetchResponse, TransactionResult}
@@ -4173,7 +4173,8 @@ class GroupCoordinatorTest {
       ArgumentMatchers.eq(producerId),
       ArgumentMatchers.eq(producerEpoch),
       any(),
-      postVerificationCallback.capture()
+      postVerificationCallback.capture(),
+      any()
     )).thenAnswer(
       _ => postVerificationCallback.getValue()((verificationError, VerificationGuard.SENTINEL))
     )
@@ -4198,7 +4199,7 @@ class GroupCoordinatorTest {
     when(replicaManager.getMagic(any[TopicPartition])).thenReturn(Some(RecordBatch.MAGIC_VALUE_V2))
 
     groupCoordinator.handleTxnCommitOffsets(groupId, transactionalId, producerId, producerEpoch,
-      memberId, groupInstanceId, generationId, offsets, responseCallback)
+      memberId, groupInstanceId, generationId, offsets, responseCallback, RequestLocal.NoCaching, ApiKeys.TXN_OFFSET_COMMIT.latestVersion())
     val result = Await.result(responseFuture, Duration(40, TimeUnit.MILLISECONDS))
     result
   }
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
index 1efccc847c912..1da8d664801e0 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
@@ -283,7 +283,7 @@ class TransactionCoordinatorTest {
 
     coordinator.handleVerifyPartitionsInTransaction(transactionalId, 0L, 0, partitions, verifyPartitionsInTxnCallback)
     errors.foreach { case (_, error) =>
-      assertEquals(Errors.INVALID_TXN_STATE, error)
+      assertEquals(Errors.ABORTABLE_TRANSACTION, error)
     }
   }
 
@@ -399,7 +399,7 @@ class TransactionCoordinatorTest {
     val extraPartitions = partitions ++ Set(new TopicPartition("topic2", 0))
     
     coordinator.handleVerifyPartitionsInTransaction(transactionalId, 0L, 0, extraPartitions, verifyPartitionsInTxnCallback)
-    assertEquals(Errors.INVALID_TXN_STATE, errors(new TopicPartition("topic2", 0)))
+    assertEquals(Errors.ABORTABLE_TRANSACTION, errors(new TopicPartition("topic2", 0)))
     assertEquals(Errors.NONE, errors(new TopicPartition("topic1", 0)))
     verify(transactionManager).getTransactionState(ArgumentMatchers.eq(transactionalId))
   }
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
index e849af43bd474..5ee9cf8302698 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
@@ -71,6 +71,7 @@ class AddPartitionsToTxnManagerTest {
   private val authenticationErrorResponse = clientResponse(null, authException = new SaslAuthenticationException(""))
   private val versionMismatchResponse = clientResponse(null, mismatchException = new UnsupportedVersionException(""))
   private val disconnectedResponse = clientResponse(null, disconnected = true)
+  private val supportedOperation = genericError
 
   private val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:2181"))
 
@@ -106,9 +107,9 @@ class AddPartitionsToTxnManagerTest {
     val transaction2Errors = mutable.Map[TopicPartition, Errors]()
     val transaction3Errors = mutable.Map[TopicPartition, Errors]()
 
-    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1Errors))
-    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transaction2Errors))
-    addPartitionsToTxnManager.verifyTransaction(transactionalId3, producerId3, producerEpoch = 0, topicPartitions, setErrors(transaction3Errors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1Errors), supportedOperation)
+    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transaction2Errors), supportedOperation)
+    addPartitionsToTxnManager.verifyTransaction(transactionalId3, producerId3, producerEpoch = 0, topicPartitions, setErrors(transaction3Errors), supportedOperation)
 
     // We will try to add transaction1 3 more times (retries). One will have the same epoch, one will have a newer epoch, and one will have an older epoch than the new one we just added.
     val transaction1RetryWithSameEpochErrors = mutable.Map[TopicPartition, Errors]()
@@ -116,17 +117,17 @@ class AddPartitionsToTxnManagerTest {
     val transaction1RetryWithOldEpochErrors = mutable.Map[TopicPartition, Errors]()
 
     // Trying to add more transactional data for the same transactional ID, producer ID, and epoch should simply replace the old data and send a retriable response.
-    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1RetryWithSameEpochErrors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1RetryWithSameEpochErrors), supportedOperation)
     val expectedNetworkErrors = topicPartitions.map(_ -> Errors.NETWORK_EXCEPTION).toMap
     assertEquals(expectedNetworkErrors, transaction1Errors)
 
     // Trying to add more transactional data for the same transactional ID and producer ID, but new epoch should replace the old data and send an error response for it.
-    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 1, topicPartitions, setErrors(transaction1RetryWithNewerEpochErrors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 1, topicPartitions, setErrors(transaction1RetryWithNewerEpochErrors), supportedOperation)
     val expectedEpochErrors = topicPartitions.map(_ -> Errors.INVALID_PRODUCER_EPOCH).toMap
     assertEquals(expectedEpochErrors, transaction1RetryWithSameEpochErrors)
 
     // Trying to add more transactional data for the same transactional ID and producer ID, but an older epoch should immediately return with error and keep the old data queued to send.
-    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1RetryWithOldEpochErrors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1RetryWithOldEpochErrors), supportedOperation)
     assertEquals(expectedEpochErrors, transaction1RetryWithOldEpochErrors)
 
     val requestsAndHandlers = addPartitionsToTxnManager.generateRequests().asScala
@@ -159,8 +160,8 @@ class AddPartitionsToTxnManagerTest {
 
     val transactionErrors = mutable.Map[TopicPartition, Errors]()
 
-    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
-    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
+    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
 
     val requestsAndHandlers = addPartitionsToTxnManager.generateRequests().asScala
     assertEquals(2, requestsAndHandlers.size)
@@ -173,8 +174,8 @@ class AddPartitionsToTxnManagerTest {
       }
     }
 
-    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
-    addPartitionsToTxnManager.verifyTransaction(transactionalId3, producerId3, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
+    addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
+    addPartitionsToTxnManager.verifyTransaction(transactionalId3, producerId3, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
 
     // Test creationTimeMs increases too.
     time.sleep(10)
@@ -207,7 +208,8 @@ class AddPartitionsToTxnManagerTest {
         producerId1,
         producerEpoch = 0,
         topicPartitions,
-        setErrors(errors)
+        setErrors(errors),
+        supportedOperation
       )
 
       assertEquals(topicPartitions.map(tp => tp -> Errors.COORDINATOR_NOT_AVAILABLE).toMap, errors)
@@ -240,8 +242,16 @@ class AddPartitionsToTxnManagerTest {
       transaction1Errors.clear()
       transaction2Errors.clear()
 
-      addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1Errors))
-      addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transaction2Errors))
+      addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1Errors), supportedOperation)
+      addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transaction2Errors), supportedOperation)
+    }
+
+    def addTransactionsToVerifyRequestVersion(operationExpected: SupportedOperation): Unit = {
+      transaction1Errors.clear()
+      transaction2Errors.clear()
+
+      addPartitionsToTxnManager.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transaction1Errors), operationExpected)
+      addPartitionsToTxnManager.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transaction2Errors), operationExpected)
     }
 
     val expectedAuthErrors = topicPartitions.map(_ -> Errors.SASL_AUTHENTICATION_FAILED).toMap
@@ -288,6 +298,32 @@ class AddPartitionsToTxnManagerTest {
     receiveResponse(mixedErrorsResponse)
     assertEquals(expectedTransaction1Errors, transaction1Errors)
     assertEquals(expectedTransaction2Errors, transaction2Errors)
+
+    val preConvertedAbortableTransaction1Errors = topicPartitions.map(_ -> Errors.ABORTABLE_TRANSACTION).toMap
+    val preConvertedAbortableTransaction2Errors = Map(new TopicPartition("foo", 1) -> Errors.NONE,
+      new TopicPartition("foo", 2) -> Errors.ABORTABLE_TRANSACTION,
+      new TopicPartition("foo", 3) -> Errors.NONE)
+    val abortableTransaction1ErrorResponse = AddPartitionsToTxnResponse.resultForTransaction(transactionalId1, preConvertedAbortableTransaction1Errors.asJava)
+    val abortableTransaction2ErrorResponse = AddPartitionsToTxnResponse.resultForTransaction(transactionalId2, preConvertedAbortableTransaction2Errors.asJava)
+    val mixedErrorsAddPartitionsResponseAbortableError = new AddPartitionsToTxnResponse(new AddPartitionsToTxnResponseData()
+      .setResultsByTransaction(new AddPartitionsToTxnResultCollection(Seq(abortableTransaction1ErrorResponse, abortableTransaction2ErrorResponse).iterator.asJava)))
+    val mixedAbortableErrorsResponse = clientResponse(mixedErrorsAddPartitionsResponseAbortableError)
+
+    val expectedAbortableTransaction1ErrorsLowerVersion = topicPartitions.map(_ -> Errors.INVALID_TXN_STATE).toMap
+    val expectedAbortableTransaction2ErrorsLowerVersion = Map(new TopicPartition("foo", 2) -> Errors.INVALID_TXN_STATE)
+
+    val expectedAbortableTransaction1ErrorsHigherVersion = topicPartitions.map(_ -> Errors.ABORTABLE_TRANSACTION).toMap
+    val expectedAbortableTransaction2ErrorsHigherVersion = Map(new TopicPartition("foo", 2) -> Errors.ABORTABLE_TRANSACTION)
+
+    addTransactionsToVerifyRequestVersion(defaultError)
+    receiveResponse(mixedAbortableErrorsResponse)
+    assertEquals(expectedAbortableTransaction1ErrorsLowerVersion, transaction1Errors)
+    assertEquals(expectedAbortableTransaction2ErrorsLowerVersion, transaction2Errors)
+
+    addTransactionsToVerifyRequestVersion(genericError)
+    receiveResponse(mixedAbortableErrorsResponse)
+    assertEquals(expectedAbortableTransaction1ErrorsHigherVersion, transaction1Errors)
+    assertEquals(expectedAbortableTransaction2ErrorsHigherVersion, transaction2Errors)
   }
 
   @Test
@@ -325,8 +361,8 @@ class AddPartitionsToTxnManagerTest {
     )
 
     try {
-      addPartitionsManagerWithMockedMetrics.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
-      addPartitionsManagerWithMockedMetrics.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors))
+      addPartitionsManagerWithMockedMetrics.verifyTransaction(transactionalId1, producerId1, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
+      addPartitionsManagerWithMockedMetrics.verifyTransaction(transactionalId2, producerId2, producerEpoch = 0, topicPartitions, setErrors(transactionErrors), supportedOperation)
 
       time.sleep(100)
 
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
index 6e296c2892b03..bc3da1852abdc 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
@@ -163,7 +163,7 @@ class AddPartitionsToTxnRequestServerTest extends BaseRequestTest {
 
     val verifyErrors = verifyResponse.errors()
 
-    assertEquals(Collections.singletonMap(transactionalId, Collections.singletonMap(tp0, Errors.INVALID_TXN_STATE)), verifyErrors)
+    assertEquals(Collections.singletonMap(transactionalId, Collections.singletonMap(tp0, Errors.ABORTABLE_TRANSACTION)), verifyErrors)
   }
   
   private def setUpTransactions(transactionalId: String, verifyOnly: Boolean, partitions: Set[TopicPartition]): (Int, AddPartitionsToTxnTransaction) = {
diff --git a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
index 90ff37ce32957..154d7e82b7466 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
@@ -2487,6 +2487,7 @@ class KafkaApisTest extends Logging {
         responseCallback.capture(),
         any(),
         any(),
+        any(),
         any()
       )).thenAnswer(_ => responseCallback.getValue.apply(Map(tp -> new PartitionResponse(Errors.INVALID_PRODUCER_EPOCH))))
 
@@ -2548,6 +2549,7 @@ class KafkaApisTest extends Logging {
         responseCallback.capture(),
         any(),
         any(),
+        any(),
         any())
       ).thenAnswer(_ => responseCallback.getValue.apply(Map(tp -> new PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER))))
 
@@ -2612,6 +2614,7 @@ class KafkaApisTest extends Logging {
         responseCallback.capture(),
         any(),
         any(),
+        any(),
         any())
       ).thenAnswer(_ => responseCallback.getValue.apply(Map(tp -> new PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER))))
 
@@ -2675,6 +2678,7 @@ class KafkaApisTest extends Logging {
         responseCallback.capture(),
         any(),
         any(),
+        any(),
         any())
       ).thenAnswer(_ => responseCallback.getValue.apply(Map(tp -> new PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER))))
 
@@ -2742,6 +2746,7 @@ class KafkaApisTest extends Logging {
           any(),
           any(),
           any(),
+          any(),
           any())
       } finally {
         kafkaApis.close()
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index a33e86e470f76..bfb822983ec81 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -112,6 +112,7 @@ class ReplicaManagerTest {
   private var mockRemoteLogManager: RemoteLogManager = _
   private var addPartitionsToTxnManager: AddPartitionsToTxnManager = _
   private var brokerTopicStats: BrokerTopicStats = _
+  private val supportedOperation = genericError
 
   // Constants defined for readability
   private val zkVersion = 0
@@ -132,7 +133,7 @@ class ReplicaManagerTest {
     addPartitionsToTxnManager = mock(classOf[AddPartitionsToTxnManager])
 
     // Anytime we try to verify, just automatically run the callback as though the transaction was verified.
-    when(addPartitionsToTxnManager.verifyTransaction(any(), any(), any(), any(), any())).thenAnswer { invocationOnMock =>
+    when(addPartitionsToTxnManager.verifyTransaction(any(), any(), any(), any(), any(), any())).thenAnswer { invocationOnMock =>
       val callback = invocationOnMock.getArgument(4, classOf[AddPartitionsToTxnManager.AppendCallback])
       callback(Map.empty[TopicPartition, Errors].toMap)
     }
@@ -2182,7 +2183,7 @@ class ReplicaManagerTest {
       val idempotentRecords = MemoryRecords.withIdempotentRecords(CompressionType.NONE, producerId, producerEpoch, sequence,
         new SimpleRecord("message".getBytes))
       handleProduceAppend(replicaManager, tp0, idempotentRecords, transactionalId = null)
-      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any[AddPartitionsToTxnManager.AppendCallback]())
+      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any[AddPartitionsToTxnManager.AppendCallback](), any())
       assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp0, producerId))
 
       // If we supply a transactional ID and some transactional and some idempotent records, we should only verify the topic partition with transactional records.
@@ -2197,7 +2198,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        any[AddPartitionsToTxnManager.AppendCallback]()
+        any[AddPartitionsToTxnManager.AppendCallback](),
+        any()
       )
       assertNotEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp0, producerId))
       assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp1, producerId))
@@ -2233,7 +2235,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback.capture()
+        appendCallback.capture(),
+        any()
       )
       val verificationGuard = getVerificationGuard(replicaManager, tp0, producerId)
       assertEquals(verificationGuard, getVerificationGuard(replicaManager, tp0, producerId))
@@ -2252,7 +2255,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback2.capture()
+        appendCallback2.capture(),
+        any()
       )
       assertEquals(verificationGuard, getVerificationGuard(replicaManager, tp0, producerId))
 
@@ -2291,7 +2295,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback.capture()
+        appendCallback.capture(),
+        any()
       )
       val verificationGuard = getVerificationGuard(replicaManager, tp0, producerId)
       assertEquals(verificationGuard, getVerificationGuard(replicaManager, tp0, producerId))
@@ -2313,7 +2318,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback2.capture()
+        appendCallback2.capture(),
+        any()
       )
       assertEquals(verificationGuard, getVerificationGuard(replicaManager, tp0, producerId))
 
@@ -2391,7 +2397,7 @@ class ReplicaManagerTest {
       assertThrows(classOf[InvalidPidMappingException],
         () => handleProduceAppendToMultipleTopics(replicaManager, transactionalRecords, transactionalId = transactionalId))
       // We should not add these partitions to the manager to verify.
-      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any())
+      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any(), any())
     } finally {
       replicaManager.shutdown(checkpointHW = false)
     }
@@ -2415,7 +2421,7 @@ class ReplicaManagerTest {
       handleProduceAppend(replicaManager, tp0, transactionalRecords, transactionalId = transactionalId).onFire { response =>
         assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, response.error)
       }
-      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any())
+      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any(), any())
     } finally {
       replicaManager.shutdown(checkpointHW = false)
     }
@@ -2449,7 +2455,7 @@ class ReplicaManagerTest {
       assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp, producerId))
 
       // We should not add these partitions to the manager to verify.
-      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any())
+      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any(), any())
 
       // Dynamically enable verification.
       config.dynamicConfig.initialize(None, None)
@@ -2463,7 +2469,7 @@ class ReplicaManagerTest {
         new SimpleRecord("message".getBytes))
 
       handleProduceAppend(replicaManager, tp, moreTransactionalRecords, transactionalId = transactionalId)
-      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any())
+      verify(addPartitionsToTxnManager, times(0)).verifyTransaction(any(), any(), any(), any(), any(), any())
       assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp, producerId))
       assertTrue(replicaManager.localLog(tp).get.hasOngoingTransaction(producerId))
     } finally {
@@ -2497,7 +2503,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback.capture()
+        appendCallback.capture(),
+        any()
       )
       val verificationGuard = getVerificationGuard(replicaManager, tp0, producerId)
       assertEquals(verificationGuard, getVerificationGuard(replicaManager, tp0, producerId))
@@ -2517,7 +2524,7 @@ class ReplicaManagerTest {
 
       // This time we do not verify
       handleProduceAppend(replicaManager, tp0, transactionalRecords, transactionalId = transactionalId)
-      verify(addPartitionsToTxnManager, times(1)).verifyTransaction(any(), any(), any(), any(), any())
+      verify(addPartitionsToTxnManager, times(1)).verifyTransaction(any(), any(), any(), any(), any(), any())
       assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp0, producerId))
       assertTrue(replicaManager.localLog(tp0).get.hasOngoingTransaction(producerId))
     } finally {
@@ -2552,7 +2559,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback.capture()
+        appendCallback.capture(),
+        any()
       )
 
       // Confirm we did not write to the log and instead returned the converted error with the correct error message.
@@ -2582,7 +2590,8 @@ class ReplicaManagerTest {
         ArgumentMatchers.eq(producerId),
         ArgumentMatchers.eq(producerEpoch),
         ArgumentMatchers.eq(Seq(tp0)),
-        appendCallback.capture()
+        appendCallback.capture(),
+        any()
       )
       assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, result.assertFired.left.getOrElse(Errors.NONE))
     } finally {
@@ -3040,6 +3049,7 @@ class ReplicaManagerTest {
       transactionalId = transactionalId,
       entriesPerPartition = entriesToAppend,
       responseCallback = appendCallback,
+      supportedOperation = supportedOperation
     )
 
     result
@@ -3067,6 +3077,7 @@ class ReplicaManagerTest {
       transactionalId = transactionalId,
       entriesPerPartition = entriesPerPartition,
       responseCallback = appendCallback,
+      supportedOperation = supportedOperation
     )
 
     result
@@ -3091,7 +3102,8 @@ class ReplicaManagerTest {
       producerId,
       producerEpoch,
       baseSequence,
-      postVerificationCallback
+      postVerificationCallback,
+      supportedOperation
     )
     result
   }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
index 2fdc128c7b98c..d6193ce799523 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
@@ -858,7 +858,8 @@ public CompletableFuture<TxnOffsetCommitResponseData> commitTransactionalOffsets
             request.producerId(),
             request.producerEpoch(),
             Duration.ofMillis(config.offsetCommitTimeoutMs),
-            coordinator -> coordinator.commitTransactionalOffset(context, request)
+            coordinator -> coordinator.commitTransactionalOffset(context, request),
+            context.apiVersion()
         ).exceptionally(exception -> handleOperationException(
             "txn-commit-offset",
             request,
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
index 6b98a51dd47b9..fd348e01ecebd 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
@@ -1482,6 +1482,7 @@ public <T> List<CompletableFuture<T>> scheduleWriteAllOperation(
      * @param producerEpoch     The producer epoch.
      * @param timeout           The write operation timeout.
      * @param op                The write operation.
+     * @param apiVersion        The Version of the Txn_Offset_Commit request
      *
      * @return A future that will be completed with the result of the write operation
      * when the operation is completed or an exception if the write operation failed.
@@ -1495,7 +1496,8 @@ public <T> CompletableFuture<T> scheduleTransactionalWriteOperation(
         long producerId,
         short producerEpoch,
         Duration timeout,
-        CoordinatorWriteOperation<S, T, U> op
+        CoordinatorWriteOperation<S, T, U> op,
+        Short apiVersion
     ) {
         throwIfNotRunning();
         log.debug("Scheduled execution of transactional write operation {}.", name);
@@ -1503,7 +1505,8 @@ public <T> CompletableFuture<T> scheduleTransactionalWriteOperation(
             tp,
             transactionalId,
             producerId,
-            producerEpoch
+            producerEpoch,
+            apiVersion
         ).thenCompose(verificationGuard -> {
             CoordinatorWriteEvent<T> event = new CoordinatorWriteEvent<>(
                 name,
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/PartitionWriter.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/PartitionWriter.java
index d647fce149db9..6ca2d2733f7c8 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/PartitionWriter.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/PartitionWriter.java
@@ -128,6 +128,7 @@ long appendEndTransactionMarker(
      * @param transactionalId   The transactional id.
      * @param producerId        The producer id.
      * @param producerEpoch     The producer epoch.
+     * @param apiVersion        The version of the Request used.
      * @return A future failed with any error encountered; or the {@link VerificationGuard}
      *         if the transaction required verification and {@link VerificationGuard#SENTINEL}
      *         if it did not.
@@ -137,6 +138,7 @@ CompletableFuture<VerificationGuard> maybeStartTransactionVerification(
         TopicPartition tp,
         String transactionalId,
         long producerId,
-        short producerEpoch
+        short producerEpoch,
+        short apiVersion
     ) throws KafkaException;
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
index 0d93f450686f5..121e624f2b053 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
@@ -1887,6 +1887,7 @@ public void testCommitTransactionalOffsets() throws ExecutionException, Interrup
             ArgumentMatchers.eq(10L),
             ArgumentMatchers.eq((short) 5),
             ArgumentMatchers.eq(Duration.ofMillis(5000)),
+            ArgumentMatchers.any(),
             ArgumentMatchers.any()
         )).thenReturn(CompletableFuture.completedFuture(response));
 
@@ -1937,6 +1938,7 @@ public void testCommitTransactionalOffsetsWithWrappedError() throws ExecutionExc
             ArgumentMatchers.eq(10L),
             ArgumentMatchers.eq((short) 5),
             ArgumentMatchers.eq(Duration.ofMillis(5000)),
+            ArgumentMatchers.any(),
             ArgumentMatchers.any()
         )).thenReturn(FutureUtils.failedFuture(new CompletionException(Errors.NOT_ENOUGH_REPLICAS.exception())));
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
index 4e10978d35740..622d4335b9a6d 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.NotCoordinatorException;
 import org.apache.kafka.common.errors.NotEnoughReplicasException;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.requests.TransactionResult;
@@ -87,6 +88,8 @@ public class CoordinatorRuntimeTest {
     private static final TopicPartition TP = new TopicPartition("__consumer_offsets", 0);
     private static final Duration DEFAULT_WRITE_TIMEOUT = Duration.ofMillis(5);
 
+    private static final short TXN_OFFSET_COMMIT_LATEST_VERSION = ApiKeys.TXN_OFFSET_COMMIT.latestVersion();
+
     /**
      * A CoordinatorEventProcessor that directly executes the operations. This is
      * useful in unit tests where execution in threads is not required.
@@ -1246,7 +1249,8 @@ public CoordinatorShardBuilder<MockCoordinatorShard, String> get() {
             TP,
             "transactional-id",
             100L,
-            (short) 50
+            (short) 50,
+            TXN_OFFSET_COMMIT_LATEST_VERSION
         )).thenReturn(CompletableFuture.completedFuture(guard));
 
         // Schedule a transactional write.
@@ -1257,7 +1261,8 @@ public CoordinatorShardBuilder<MockCoordinatorShard, String> get() {
             100L,
             (short) 50,
             Duration.ofMillis(5000),
-            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response")
+            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response"),
+            TXN_OFFSET_COMMIT_LATEST_VERSION
         );
 
         // Verify that the writer got the records with the correct
@@ -1328,7 +1333,8 @@ public CoordinatorShardBuilder<MockCoordinatorShard, String> get() {
             TP,
             "transactional-id",
             100L,
-            (short) 50
+            (short) 50,
+            TXN_OFFSET_COMMIT_LATEST_VERSION
         )).thenReturn(FutureUtils.failedFuture(Errors.NOT_ENOUGH_REPLICAS.exception()));
 
         // Schedule a transactional write.
@@ -1339,7 +1345,8 @@ public CoordinatorShardBuilder<MockCoordinatorShard, String> get() {
             100L,
             (short) 50,
             Duration.ofMillis(5000),
-            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response")
+            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response"),
+            TXN_OFFSET_COMMIT_LATEST_VERSION
         );
 
         // Verify that the future is failed with the expected exception.
@@ -1391,7 +1398,8 @@ public void testScheduleTransactionCompletion(TransactionResult result) throws E
             100L,
             (short) 5,
             DEFAULT_WRITE_TIMEOUT,
-            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1")
+            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1"),
+            TXN_OFFSET_COMMIT_LATEST_VERSION
         );
 
         // Verify that the write is not committed yet.
@@ -1556,7 +1564,9 @@ public void testScheduleTransactionCompletionWhenWriteFails() {
             100L,
             (short) 5,
             DEFAULT_WRITE_TIMEOUT,
-            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1"));
+            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1"),
+            TXN_OFFSET_COMMIT_LATEST_VERSION
+        );
 
         // Verify that the state has been updated.
         assertEquals(2L, ctx.coordinator.lastWrittenOffset());
@@ -1638,7 +1648,9 @@ public void replayEndTransactionMarker(
             100L,
             (short) 5,
             DEFAULT_WRITE_TIMEOUT,
-            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1"));
+            state -> new CoordinatorResult<>(Arrays.asList("record1", "record2"), "response1"),
+            TXN_OFFSET_COMMIT_LATEST_VERSION
+        );
 
         // Verify that the state has been updated.
         assertEquals(2L, ctx.coordinator.lastWrittenOffset());
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/InMemoryPartitionWriter.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/InMemoryPartitionWriter.java
index 7dea0b83f7ee9..da6d662b6d5bc 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/InMemoryPartitionWriter.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/InMemoryPartitionWriter.java
@@ -274,7 +274,8 @@ public CompletableFuture<VerificationGuard> maybeStartTransactionVerification(
         TopicPartition tp,
         String transactionalId,
         long producerId,
-        short producerEpoch
+        short producerEpoch,
+        short apiVersion
     ) throws KafkaException {
         return CompletableFuture.completedFuture(new VerificationGuard());
     }

From 6eba9057b1d1200e270ab06351de6ece123ef407 Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Sat, 23 Mar 2024 11:37:55 +0800
Subject: [PATCH 192/258] KAFKA-16381 use volatile to guarantee
 KafkaMetric#config visibility across threads (#15550)

Reviewers: vamossagar12 <sagarmeansocean@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../main/java/org/apache/kafka/common/metrics/KafkaMetric.java  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
index 805ae2a52781b..c8e53ffc6cea0 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
@@ -26,7 +26,7 @@ public final class KafkaMetric implements Metric {
     private final Object lock;
     private final Time time;
     private final MetricValueProvider<?> metricValueProvider;
-    private MetricConfig config;
+    private volatile MetricConfig config;
 
     // public for testing
     /**

From 0f216b644874f2f89285f9d4736b50246cd5b0a2 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Sat, 23 Mar 2024 06:44:05 +0300
Subject: [PATCH 193/258] MINOR: Tuple2 replaced with Map.Entry (#15560)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../java/org/apache/kafka/tools/Tuple2.java   |  48 ------
 .../consumer/group/ConsumerGroupCommand.java  |  59 +++----
 .../reassign/ReassignPartitionsCommand.java   | 161 +++++++++---------
 .../apache/kafka/tools/ToolsTestUtils.java    |  14 +-
 .../group/ConsumerGroupServiceTest.java       |  24 +--
 .../group/DeleteConsumerGroupsTest.java       |   4 +-
 ...tsConsumerGroupCommandIntegrationTest.java |  12 +-
 .../group/DescribeConsumerGroupTest.java      | 156 ++++++++---------
 .../tools/other/ReplicationQuotasTestRig.java |   2 +-
 .../ReassignPartitionsIntegrationTest.java    |  11 +-
 .../reassign/ReassignPartitionsUnitTest.java  |  47 ++---
 11 files changed, 248 insertions(+), 290 deletions(-)
 delete mode 100644 tools/src/main/java/org/apache/kafka/tools/Tuple2.java

diff --git a/tools/src/main/java/org/apache/kafka/tools/Tuple2.java b/tools/src/main/java/org/apache/kafka/tools/Tuple2.java
deleted file mode 100644
index 02dd4bf3981a8..0000000000000
--- a/tools/src/main/java/org/apache/kafka/tools/Tuple2.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.tools;
-
-import java.util.Objects;
-
-public final class Tuple2<V1, V2> {
-    public final V1 v1;
-
-    public final V2 v2;
-
-    public Tuple2(V1 v1, V2 v2) {
-        this.v1 = v1;
-        this.v2 = v2;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        Tuple2<?, ?> tuple = (Tuple2<?, ?>) o;
-        return Objects.equals(v1, tuple.v1) && Objects.equals(v2, tuple.v2);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(v1, v2);
-    }
-
-    @Override
-    public String toString() {
-        return "Tuple2{v1=" + v1 + ", v2=" + v2 + '}';
-    }
-}
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
index 2b055de616191..35371ce75aa65 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
@@ -53,7 +53,6 @@
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.util.CommandLineUtils;
 import org.apache.kafka.tools.ToolsUtils;
-import org.apache.kafka.tools.Tuple2;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -61,6 +60,7 @@
 import java.text.ParseException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -71,6 +71,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Properties;
@@ -315,10 +316,10 @@ private Optional<Integer> size(Optional<? extends Collection<?>> colOpt) {
             return colOpt.map(Collection::size);
         }
 
-        private void printOffsets(Map<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets) {
+        private void printOffsets(Map<String, Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets) {
             offsets.forEach((groupId, tuple) -> {
-                Optional<String> state = tuple.v1;
-                Optional<Collection<PartitionAssignmentState>> assignments = tuple.v2;
+                Optional<String> state = tuple.getKey();
+                Optional<Collection<PartitionAssignmentState>> assignments = tuple.getValue();
 
                 if (shouldPrintMemberState(groupId, state, size(assignments))) {
                     String format = printOffsetFormat(assignments);
@@ -359,10 +360,10 @@ private static String printOffsetFormat(Optional<Collection<PartitionAssignmentS
             return format;
         }
 
-        private void printMembers(Map<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members, boolean verbose) {
+        private void printMembers(Map<String, Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members, boolean verbose) {
             members.forEach((groupId, tuple) -> {
-                Optional<String> state = tuple.v1;
-                Optional<Collection<MemberAssignmentState>> assignments = tuple.v2;
+                Optional<String> state = tuple.getKey();
+                Optional<Collection<MemberAssignmentState>> assignments = tuple.getValue();
                 int maxGroupLen = 15, maxConsumerIdLen = 15, maxGroupInstanceIdLen = 17, maxHostLen = 15, maxClientIdLen = 15;
                 boolean includeGroupInstanceId = false;
 
@@ -448,11 +449,11 @@ void describeGroups() throws Exception {
             long subActions = Stream.of(membersOptPresent, offsetsOptPresent, stateOptPresent).filter(x -> x).count();
 
             if (subActions == 0 || offsetsOptPresent) {
-                TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets
+                TreeMap<String, Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> offsets
                     = collectGroupsOffsets(groupIds);
                 printOffsets(offsets);
             } else if (membersOptPresent) {
-                TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members
+                TreeMap<String, Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>>> members
                     = collectGroupsMembers(groupIds, opts.options.has(opts.verboseOpt));
                 printMembers(members, opts.options.has(opts.verboseOpt));
             } else {
@@ -560,7 +561,7 @@ Map<String, Map<TopicPartition, OffsetAndMetadata>> resetOffsets() {
             return result;
         }
 
-        Tuple2<Errors, Map<TopicPartition, Throwable>> deleteOffsets(String groupId, List<String> topics) {
+        Entry<Errors, Map<TopicPartition, Throwable>> deleteOffsets(String groupId, List<String> topics) {
             Map<TopicPartition, Throwable> partitionLevelResult = new HashMap<>();
             Set<String> topicWithPartitions = new HashSet<>();
             Set<String> topicWithoutPartitions = new HashSet<>();
@@ -617,17 +618,17 @@ Tuple2<Errors, Map<TopicPartition, Throwable>> deleteOffsets(String groupId, Lis
                 }
             });
 
-            return new Tuple2<>(topLevelException, partitionLevelResult);
+            return new SimpleImmutableEntry<>(topLevelException, partitionLevelResult);
         }
 
         void deleteOffsets() {
             String groupId = opts.options.valueOf(opts.groupOpt);
             List<String> topics = opts.options.valuesOf(opts.topicOpt);
 
-            Tuple2<Errors, Map<TopicPartition, Throwable>> res = deleteOffsets(groupId, topics);
+            Entry<Errors, Map<TopicPartition, Throwable>> res = deleteOffsets(groupId, topics);
 
-            Errors topLevelResult = res.v1;
-            Map<TopicPartition, Throwable> partitionLevelResult = res.v2;
+            Errors topLevelResult = res.getKey();
+            Map<TopicPartition, Throwable> partitionLevelResult = res.getValue();
 
             switch (topLevelResult) {
                 case NONE:
@@ -677,7 +678,7 @@ Map<String, ConsumerGroupDescription> describeConsumerGroups(Collection<String>
                 withTimeoutMs(new DescribeConsumerGroupsOptions())
             ).describedGroups();
 
-            for (Map.Entry<String, KafkaFuture<ConsumerGroupDescription>> e : stringKafkaFutureMap.entrySet()) {
+            for (Entry<String, KafkaFuture<ConsumerGroupDescription>> e : stringKafkaFutureMap.entrySet()) {
                 res.put(e.getKey(), e.getValue().get());
             }
             return res;
@@ -686,16 +687,16 @@ Map<String, ConsumerGroupDescription> describeConsumerGroups(Collection<String>
         /**
          * Returns the state of the specified consumer group and partition assignment states
          */
-        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> collectGroupOffsets(String groupId) throws Exception {
-            return collectGroupsOffsets(Collections.singletonList(groupId)).getOrDefault(groupId, new Tuple2<>(Optional.empty(), Optional.empty()));
+        Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> collectGroupOffsets(String groupId) throws Exception {
+            return collectGroupsOffsets(Collections.singletonList(groupId)).getOrDefault(groupId, new SimpleImmutableEntry<>(Optional.empty(), Optional.empty()));
         }
 
         /**
          * Returns states of the specified consumer groups and partition assignment states
          */
-        TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> collectGroupsOffsets(Collection<String> groupIds) throws Exception {
+        TreeMap<String, Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> collectGroupsOffsets(Collection<String> groupIds) throws Exception {
             Map<String, ConsumerGroupDescription> consumerGroups = describeConsumerGroups(groupIds);
-            TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> groupOffsets = new TreeMap<>();
+            TreeMap<String, Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>>> groupOffsets = new TreeMap<>();
 
             consumerGroups.forEach((groupId, consumerGroup) -> {
                 ConsumerGroupState state = consumerGroup.state();
@@ -737,19 +738,19 @@ TreeMap<String, Tuple2<Optional<String>, Optional<Collection<PartitionAssignment
 
                 rowsWithConsumer.addAll(rowsWithoutConsumer);
 
-                groupOffsets.put(groupId, new Tuple2<>(Optional.of(state.toString()), Optional.of(rowsWithConsumer)));
+                groupOffsets.put(groupId, new SimpleImmutableEntry<>(Optional.of(state.toString()), Optional.of(rowsWithConsumer)));
             });
 
             return groupOffsets;
         }
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> collectGroupMembers(String groupId, boolean verbose) throws Exception {
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> collectGroupMembers(String groupId, boolean verbose) throws Exception {
             return collectGroupsMembers(Collections.singleton(groupId), verbose).get(groupId);
         }
 
-        TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> collectGroupsMembers(Collection<String> groupIds, boolean verbose) throws Exception {
+        TreeMap<String, Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>>> collectGroupsMembers(Collection<String> groupIds, boolean verbose) throws Exception {
             Map<String, ConsumerGroupDescription> consumerGroups = describeConsumerGroups(groupIds);
-            TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>>> res = new TreeMap<>();
+            TreeMap<String, Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>>> res = new TreeMap<>();
 
             consumerGroups.forEach((groupId, consumerGroup) -> {
                 String state = consumerGroup.state().toString();
@@ -763,7 +764,7 @@ TreeMap<String, Tuple2<Optional<String>, Optional<Collection<MemberAssignmentSta
                         consumer.assignment().topicPartitions().size(),
                         new ArrayList<>(verbose ? consumer.assignment().topicPartitions() : Collections.emptySet())
                 )).collect(Collectors.toList());
-                res.put(groupId, new Tuple2<>(Optional.of(state), Optional.of(memberAssignmentStates)));
+                res.put(groupId, new SimpleImmutableEntry<>(Optional.of(state), Optional.of(memberAssignmentStates)));
             });
             return res;
         }
@@ -836,7 +837,7 @@ private Map<TopicPartition, LogOffsetResult> getLogTimestampOffsets(Collection<T
                 });
 
                 Map<TopicPartition, LogOffsetResult> successfulLogTimestampOffsets = successfulOffsetsForTimes.entrySet().stream()
-                    .collect(Collectors.toMap(Map.Entry::getKey, e -> new LogOffset(e.getValue().offset())));
+                    .collect(Collectors.toMap(Entry::getKey, e -> new LogOffset(e.getValue().offset())));
 
                 unsuccessfulOffsetsForTimes.forEach((tp, offsetResultInfo) ->
                     System.out.println("\nWarn: Partition " + tp.partition() + " from topic " + tp.topic() +
@@ -991,7 +992,7 @@ private Map<TopicPartition, OffsetAndMetadata> prepareOffsetsToReset(String grou
             if (opts.options.has(opts.resetToOffsetOpt)) {
                 long offset = opts.options.valueOf(opts.resetToOffsetOpt);
                 return checkOffsetsRange(partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), tp -> offset)))
-                    .entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+                    .entrySet().stream().collect(Collectors.toMap(Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
             } else if (opts.options.has(opts.resetToEarliestOpt)) {
                 Map<TopicPartition, LogOffsetResult> logStartOffsets = getLogStartOffsets(partitionsToReset);
                 return partitionsToReset.stream().collect(Collectors.toMap(Function.identity(), topicPartition -> {
@@ -1029,7 +1030,7 @@ private Map<TopicPartition, OffsetAndMetadata> prepareOffsetsToReset(String grou
                     return currentOffset.offset() + shiftBy;
                 }));
                 return checkOffsetsRange(requestedOffsets).entrySet().stream()
-                    .collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+                    .collect(Collectors.toMap(Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
             } else if (opts.options.has(opts.resetToDatetimeOpt)) {
                 try {
                     long timestamp = Utils.getDateTime(opts.options.valueOf(opts.resetToDatetimeOpt));
@@ -1078,7 +1079,7 @@ private Map<TopicPartition, OffsetAndMetadata> prepareOffsetsToReset(String grou
                         topicPartition -> resetPlanForGroup.get(topicPartition).offset()));
 
                     return checkOffsetsRange(requestedOffsets).entrySet().stream()
-                        .collect(Collectors.toMap(Map.Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
+                        .collect(Collectors.toMap(Entry::getKey, e -> new OffsetAndMetadata(e.getValue())));
                 }).orElseGet(Collections::emptyMap);
             } else if (opts.options.has(opts.resetToCurrentOpt)) {
                 Map<TopicPartition, OffsetAndMetadata> currentCommittedOffsets = getCommittedOffsets(groupId);
@@ -1104,7 +1105,7 @@ private Map<TopicPartition, OffsetAndMetadata> prepareOffsetsToReset(String grou
                     }));
 
                 Map<TopicPartition, OffsetAndMetadata> preparedOffsetsForPartitionsWithoutCommittedOffset = getLogEndOffsets(partitionsToResetWithoutCommittedOffset)
-                    .entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> {
+                    .entrySet().stream().collect(Collectors.toMap(Entry::getKey, e -> {
                         if (!(e.getValue() instanceof LogOffset)) {
                             ToolsUtils.printUsageAndExit(opts.parser, "Error getting ending offset of topic partition: " + e.getKey());
                             return null;
diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
index 3623346980743..47c4a8234a177 100644
--- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
@@ -49,9 +49,9 @@
 import org.apache.kafka.server.util.json.JsonValue;
 import org.apache.kafka.tools.TerseException;
 import org.apache.kafka.tools.ToolsUtils;
-import org.apache.kafka.tools.Tuple2;
 
 import java.io.IOException;
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -62,6 +62,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
@@ -203,20 +204,20 @@ static VerifyAssignmentResult verifyAssignment(Admin adminClient,
                                                    String jsonString,
                                                    Boolean preserveThrottles
     ) throws ExecutionException, InterruptedException, JsonProcessingException {
-        Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(jsonString);
+        Entry<List<Entry<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(jsonString);
 
-        List<Tuple2<TopicPartition, List<Integer>>> targetParts = t0.v1;
-        Map<TopicPartitionReplica, String> targetLogDirs = t0.v2;
+        List<Entry<TopicPartition, List<Integer>>> targetParts = t0.getKey();
+        Map<TopicPartitionReplica, String> targetLogDirs = t0.getValue();
 
-        Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> t1 = verifyPartitionAssignments(adminClient, targetParts);
+        Entry<Map<TopicPartition, PartitionReassignmentState>, Boolean> t1 = verifyPartitionAssignments(adminClient, targetParts);
 
-        Map<TopicPartition, PartitionReassignmentState> partStates = t1.v1;
-        Boolean partsOngoing = t1.v2;
+        Map<TopicPartition, PartitionReassignmentState> partStates = t1.getKey();
+        Boolean partsOngoing = t1.getValue();
 
-        Tuple2<Map<TopicPartitionReplica, LogDirMoveState>, Boolean> t2 = verifyReplicaMoves(adminClient, targetLogDirs);
+        Entry<Map<TopicPartitionReplica, LogDirMoveState>, Boolean> t2 = verifyReplicaMoves(adminClient, targetLogDirs);
 
-        Map<TopicPartitionReplica, LogDirMoveState> moveStates = t2.v1;
-        Boolean movesOngoing = t2.v2;
+        Map<TopicPartitionReplica, LogDirMoveState> moveStates = t2.getKey();
+        Boolean movesOngoing = t2.getValue();
 
         if (!partsOngoing && !movesOngoing && !preserveThrottles) {
             // If the partition assignments and replica assignments are done, clear any throttles
@@ -240,11 +241,11 @@ static VerifyAssignmentResult verifyAssignment(Admin adminClient,
      *                              reassignments (including reassignments not described
      *                              in the JSON file.)
      */
-    private static Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> verifyPartitionAssignments(Admin adminClient,
-                                                                                                               List<Tuple2<TopicPartition, List<Integer>>> targets
+    private static Entry<Map<TopicPartition, PartitionReassignmentState>, Boolean> verifyPartitionAssignments(Admin adminClient,
+                                                                                                               List<Entry<TopicPartition, List<Integer>>> targets
     ) throws ExecutionException, InterruptedException {
-        Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> t0 = findPartitionReassignmentStates(adminClient, targets);
-        System.out.println(partitionReassignmentStatesToString(t0.v1));
+        Entry<Map<TopicPartition, PartitionReassignmentState>, Boolean> t0 = findPartitionReassignmentStates(adminClient, targets);
+        System.out.println(partitionReassignmentStatesToString(t0.getKey()));
         return t0;
     }
 
@@ -301,27 +302,27 @@ static String partitionReassignmentStatesToString(Map<TopicPartition, PartitionR
      * @return                     A tuple containing the reassignment states for each topic
      *                             partition, plus whether there are any ongoing reassignments.
      */
-    static Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> findPartitionReassignmentStates(Admin adminClient,
-                                                                                                            List<Tuple2<TopicPartition,
+    static Entry<Map<TopicPartition, PartitionReassignmentState>, Boolean> findPartitionReassignmentStates(Admin adminClient,
+                                                                                                            List<Entry<TopicPartition,
                                                                                                             List<Integer>>> targetReassignments
     ) throws ExecutionException, InterruptedException {
         Map<TopicPartition, PartitionReassignment> currentReassignments = adminClient.
             listPartitionReassignments().reassignments().get();
 
-        List<Tuple2<TopicPartition, List<Integer>>> foundReassignments = new ArrayList<>();
-        List<Tuple2<TopicPartition, List<Integer>>> notFoundReassignments = new ArrayList<>();
+        List<Entry<TopicPartition, List<Integer>>> foundReassignments = new ArrayList<>();
+        List<Entry<TopicPartition, List<Integer>>> notFoundReassignments = new ArrayList<>();
 
         targetReassignments.forEach(reassignment -> {
-            if (currentReassignments.containsKey(reassignment.v1))
+            if (currentReassignments.containsKey(reassignment.getKey()))
                 foundReassignments.add(reassignment);
             else
                 notFoundReassignments.add(reassignment);
         });
 
-        List<Tuple2<TopicPartition, PartitionReassignmentState>> foundResults = foundReassignments.stream().map(e -> {
-            TopicPartition part = e.v1;
-            List<Integer> targetReplicas = e.v2;
-            return new Tuple2<>(part,
+        List<Entry<TopicPartition, PartitionReassignmentState>> foundResults = foundReassignments.stream().map(e -> {
+            TopicPartition part = e.getKey();
+            List<Integer> targetReplicas = e.getValue();
+            return new SimpleImmutableEntry<>(part,
                 new PartitionReassignmentState(
                     currentReassignments.get(part).replicas(),
                     targetReplicas,
@@ -329,7 +330,7 @@ static Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> findPart
         }).collect(Collectors.toList());
 
         Set<String> topicNamesToLookUp = notFoundReassignments.stream()
-            .map(e -> e.v1)
+            .map(e -> e.getKey())
             .filter(part -> !currentReassignments.containsKey(part))
             .map(TopicPartition::topic)
             .collect(Collectors.toSet());
@@ -337,28 +338,28 @@ static Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> findPart
         Map<String, KafkaFuture<TopicDescription>> topicDescriptions = adminClient.
             describeTopics(topicNamesToLookUp).topicNameValues();
 
-        List<Tuple2<TopicPartition, PartitionReassignmentState>> notFoundResults = new ArrayList<>();
-        for (Tuple2<TopicPartition, List<Integer>> e : notFoundReassignments) {
-            TopicPartition part = e.v1;
-            List<Integer> targetReplicas = e.v2;
+        List<Entry<TopicPartition, PartitionReassignmentState>> notFoundResults = new ArrayList<>();
+        for (Entry<TopicPartition, List<Integer>> e : notFoundReassignments) {
+            TopicPartition part = e.getKey();
+            List<Integer> targetReplicas = e.getValue();
 
             if (currentReassignments.containsKey(part)) {
                 PartitionReassignment reassignment = currentReassignments.get(part);
-                notFoundResults.add(new Tuple2<>(part, new PartitionReassignmentState(
+                notFoundResults.add(new SimpleImmutableEntry<>(part, new PartitionReassignmentState(
                     reassignment.replicas(),
                     targetReplicas,
                     false)));
             } else {
-                notFoundResults.add(new Tuple2<>(part, topicDescriptionFutureToState(part.partition(),
+                notFoundResults.add(new SimpleImmutableEntry<>(part, topicDescriptionFutureToState(part.partition(),
                     topicDescriptions.get(part.topic()), targetReplicas)));
             }
         }
 
         Map<TopicPartition, PartitionReassignmentState> allResults = new HashMap<>();
-        foundResults.forEach(e -> allResults.put(e.v1, e.v2));
-        notFoundResults.forEach(e -> allResults.put(e.v1, e.v2));
+        foundResults.forEach(e -> allResults.put(e.getKey(), e.getValue()));
+        notFoundResults.forEach(e -> allResults.put(e.getKey(), e.getValue()));
 
-        return new Tuple2<>(allResults, !currentReassignments.isEmpty());
+        return new SimpleImmutableEntry<>(allResults, !currentReassignments.isEmpty());
     }
 
     private static PartitionReassignmentState topicDescriptionFutureToState(int partition,
@@ -396,12 +397,12 @@ private static PartitionReassignmentState topicDescriptionFutureToState(int part
      *                              reassignments. (We don't have an efficient API that
      *                              returns all ongoing replica reassignments.)
      */
-    private static Tuple2<Map<TopicPartitionReplica, LogDirMoveState>, Boolean> verifyReplicaMoves(Admin adminClient,
+    private static Entry<Map<TopicPartitionReplica, LogDirMoveState>, Boolean> verifyReplicaMoves(Admin adminClient,
                                                                                                    Map<TopicPartitionReplica, String> targetReassignments
     ) throws ExecutionException, InterruptedException {
         Map<TopicPartitionReplica, LogDirMoveState> moveStates = findLogDirMoveStates(adminClient, targetReassignments);
         System.out.println(replicaMoveStatesToString(moveStates));
-        return new Tuple2<>(moveStates, !moveStates.values().stream().allMatch(LogDirMoveState::done));
+        return new SimpleImmutableEntry<>(moveStates, !moveStates.values().stream().allMatch(LogDirMoveState::done));
     }
 
     /**
@@ -420,7 +421,7 @@ static Map<TopicPartitionReplica, LogDirMoveState> findLogDirMoveStates(Admin ad
         Map<TopicPartitionReplica, DescribeReplicaLogDirsResult.ReplicaLogDirInfo> replicaLogDirInfos = adminClient
             .describeReplicaLogDirs(targetMoves.keySet()).all().get();
 
-        return targetMoves.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> {
+        return targetMoves.entrySet().stream().collect(Collectors.toMap(Entry::getKey, e -> {
             TopicPartitionReplica replica = e.getKey();
             String targetLogDir = e.getValue();
 
@@ -492,16 +493,16 @@ static String replicaMoveStatesToString(Map<TopicPartitionReplica, LogDirMoveSta
      * @param targetParts     The target partitions loaded from the JSON file.
      */
     private static void clearAllThrottles(Admin adminClient,
-                                          List<Tuple2<TopicPartition, List<Integer>>> targetParts
+                                          List<Entry<TopicPartition, List<Integer>>> targetParts
     ) throws ExecutionException, InterruptedException {
         Set<Integer> brokers = adminClient.describeCluster().nodes().get().stream().map(Node::id).collect(Collectors.toSet());
-        targetParts.forEach(t -> brokers.addAll(t.v2));
+        targetParts.forEach(t -> brokers.addAll(t.getValue()));
 
         System.out.printf("Clearing broker-level throttles on broker%s %s%n",
             brokers.size() == 1 ? "" : "s", Utils.join(brokers, ","));
         clearBrokerLevelThrottles(adminClient, brokers);
 
-        Set<String> topics = targetParts.stream().map(t -> t.v1.topic()).collect(Collectors.toSet());
+        Set<String> topics = targetParts.stream().map(t -> t.getKey().topic()).collect(Collectors.toSet());
         System.out.printf("Clearing topic-level throttles on topic%s %s%n",
             topics.size() == 1 ? "" : "s", Utils.join(topics, ","));
         clearTopicLevelThrottles(adminClient, topics);
@@ -550,15 +551,15 @@ private static void clearTopicLevelThrottles(Admin adminClient, Set<String> topi
      * @return                      A tuple containing the proposed assignment and the
      *                              current assignment.
      */
-    public static Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>> generateAssignment(Admin adminClient,
+    public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>> generateAssignment(Admin adminClient,
                                                                                                              String reassignmentJson,
                                                                                                              String brokerListString,
                                                                                                              Boolean enableRackAwareness
     ) throws ExecutionException, InterruptedException, JsonProcessingException {
-        Tuple2<List<Integer>, List<String>> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString);
+        Entry<List<Integer>, List<String>> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString);
 
-        List<Integer> brokersToReassign = t0.v1;
-        List<String> topicsToReassign = t0.v2;
+        List<Integer> brokersToReassign = t0.getKey();
+        List<String> topicsToReassign = t0.getValue();
 
         Map<TopicPartition, List<Integer>> currentAssignments = getReplicaAssignmentForTopics(adminClient, topicsToReassign);
         List<BrokerMetadata> brokerMetadatas = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness);
@@ -567,7 +568,7 @@ public static Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartition, Lis
             formatAsReassignmentJson(currentAssignments, Collections.emptyMap()));
         System.out.printf("Proposed partition reassignment configuration%n%s%n",
             formatAsReassignmentJson(proposedAssignments, Collections.emptyMap()));
-        return new Tuple2<>(proposedAssignments, currentAssignments);
+        return new SimpleImmutableEntry<>(proposedAssignments, currentAssignments);
     }
 
     /**
@@ -580,8 +581,8 @@ public static Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartition, Lis
      */
     private static Map<TopicPartition, List<Integer>> calculateAssignment(Map<TopicPartition, List<Integer>> currentAssignment,
                                                                           List<BrokerMetadata> brokerMetadatas) {
-        Map<String, List<Map.Entry<TopicPartition, List<Integer>>>> groupedByTopic = new HashMap<>();
-        for (Map.Entry<TopicPartition, List<Integer>> e : currentAssignment.entrySet())
+        Map<String, List<Entry<TopicPartition, List<Integer>>>> groupedByTopic = new HashMap<>();
+        for (Entry<TopicPartition, List<Integer>> e : currentAssignment.entrySet())
             groupedByTopic.computeIfAbsent(e.getKey().topic(), k -> new ArrayList<>()).add(e);
         Map<TopicPartition, List<Integer>> proposedAssignments = new HashMap<>();
         groupedByTopic.forEach((topic, assignment) -> {
@@ -598,7 +599,7 @@ private static Map<String, TopicDescription> describeTopics(Admin adminClient,
                                                                 Set<String> topics) throws ExecutionException, InterruptedException {
         Map<String, KafkaFuture<TopicDescription>> futures = adminClient.describeTopics(topics).topicNameValues();
         Map<String, TopicDescription> res = new HashMap<>();
-        for (Map.Entry<String, KafkaFuture<TopicDescription>> e : futures.entrySet()) {
+        for (Entry<String, KafkaFuture<TopicDescription>> e : futures.entrySet()) {
             String topicName = e.getKey();
             KafkaFuture<TopicDescription> topicDescriptionFuture = e.getValue();
             try {
@@ -695,7 +696,7 @@ static List<BrokerMetadata> getBrokerMetadata(Admin adminClient, List<Integer> b
      *
      * @return                       A tuple of brokers to reassign, topics to reassign
      */
-    static Tuple2<List<Integer>, List<String>> parseGenerateAssignmentArgs(String reassignmentJson,
+    static Entry<List<Integer>, List<String>> parseGenerateAssignmentArgs(String reassignmentJson,
                                                                            String brokerList) throws JsonMappingException {
         List<Integer> brokerListToReassign = Arrays.stream(brokerList.split(",")).map(Integer::parseInt).collect(Collectors.toList());
         Set<Integer> duplicateReassignments = ToolsUtils.duplicates(brokerListToReassign);
@@ -706,7 +707,7 @@ static Tuple2<List<Integer>, List<String>> parseGenerateAssignmentArgs(String re
         if (!duplicateTopicsToReassign.isEmpty())
             throw new AdminCommandFailedException(String.format("List of topics to reassign contains duplicate entries: %s",
                 duplicateTopicsToReassign));
-        return new Tuple2<>(brokerListToReassign, topicsToReassign);
+        return new SimpleImmutableEntry<>(brokerListToReassign, topicsToReassign);
     }
 
     /**
@@ -731,10 +732,10 @@ public static void executeAssignment(Admin adminClient,
                                   Long timeoutMs,
                                   Time time
     ) throws ExecutionException, InterruptedException, JsonProcessingException, TerseException {
-        Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> t0 = parseExecuteAssignmentArgs(reassignmentJson);
+        Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> t0 = parseExecuteAssignmentArgs(reassignmentJson);
 
-        Map<TopicPartition, List<Integer>> proposedParts = t0.v1;
-        Map<TopicPartitionReplica, String> proposedReplicas = t0.v2;
+        Map<TopicPartition, List<Integer>> proposedParts = t0.getKey();
+        Map<TopicPartitionReplica, String> proposedReplicas = t0.getValue();
         Map<TopicPartition, PartitionReassignment> currentReassignments = adminClient.
             listPartitionReassignments().reassignments().get();
         // If there is an existing assignment, check for --additional before proceeding.
@@ -901,7 +902,7 @@ static String currentPartitionReplicaAssignmentToString(Map<TopicPartition, List
                                                             Map<TopicPartition, List<Integer>> currentParts) throws JsonProcessingException {
         Map<TopicPartition, List<Integer>> partitionsToBeReassigned = currentParts.entrySet().stream()
             .filter(e -> proposedParts.containsKey(e.getKey()))
-            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+            .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
 
         return String.format("Current partition replica assignment%n%n%s%n%nSave this to use as the %s",
             formatAsReassignmentJson(partitionsToBeReassigned, Collections.emptyMap()),
@@ -921,7 +922,7 @@ static Map<TopicPartition, Throwable> alterPartitionReassignments(Admin adminCli
         reassignments.forEach((part, replicas) -> args.put(part, Optional.of(new NewPartitionReassignment(replicas))));
         Map<TopicPartition, KafkaFuture<Void>> results = adminClient.alterPartitionReassignments(args).values();
         Map<TopicPartition, Throwable> errors = new HashMap<>();
-        for (Map.Entry<TopicPartition, KafkaFuture<Void>> e :  results.entrySet()) {
+        for (Entry<TopicPartition, KafkaFuture<Void>> e :  results.entrySet()) {
             try {
                 e.getValue().get();
             } catch (ExecutionException t) {
@@ -945,7 +946,7 @@ static Map<TopicPartition, Throwable> cancelPartitionReassignments(Admin adminCl
 
         Map<TopicPartition, KafkaFuture<Void>> results = adminClient.alterPartitionReassignments(args).values();
         Map<TopicPartition, Throwable> errors = new HashMap<>();
-        for (Map.Entry<TopicPartition, KafkaFuture<Void>> e :  results.entrySet()) {
+        for (Entry<TopicPartition, KafkaFuture<Void>> e :  results.entrySet()) {
             try {
                 e.getValue().get();
             } catch (ExecutionException t) {
@@ -993,7 +994,7 @@ static Map<String, Map<Integer, PartitionMove>> calculateProposedMoveMap(Map<Top
                                                                              Map<TopicPartition, List<Integer>> proposedParts,
                                                                              Map<TopicPartition, List<Integer>> currentParts) {
         Map<String, Map<Integer, PartitionMove>> moveMap = calculateCurrentMoveMap(currentReassignments);
-        for (Map.Entry<TopicPartition, List<Integer>> e : proposedParts.entrySet()) {
+        for (Entry<TopicPartition, List<Integer>> e : proposedParts.entrySet()) {
             TopicPartition part = e.getKey();
             List<Integer> replicas = e.getValue();
             Map<Integer, PartitionMove> partMoves = moveMap.computeIfAbsent(part.topic(), k -> new HashMap<>());
@@ -1178,39 +1179,39 @@ static void modifyLogDirThrottle(Admin admin,
      * @return                  A tuple of the partitions to be reassigned and the replicas
      *                          to be reassigned.
      */
-    static Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> parseExecuteAssignmentArgs(
+    static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> parseExecuteAssignmentArgs(
         String reassignmentJson
     ) throws JsonProcessingException {
-        Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(reassignmentJson);
+        Entry<List<Entry<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(reassignmentJson);
 
-        List<Tuple2<TopicPartition, List<Integer>>> partitionsToBeReassigned = t0.v1;
-        Map<TopicPartitionReplica, String> replicaAssignment = t0.v2;
+        List<Entry<TopicPartition, List<Integer>>> partitionsToBeReassigned = t0.getKey();
+        Map<TopicPartitionReplica, String> replicaAssignment = t0.getValue();
 
         if (partitionsToBeReassigned.isEmpty())
             throw new AdminCommandFailedException("Partition reassignment list cannot be empty");
-        if (partitionsToBeReassigned.stream().anyMatch(t -> t.v2.isEmpty())) {
+        if (partitionsToBeReassigned.stream().anyMatch(t -> t.getValue().isEmpty())) {
             throw new AdminCommandFailedException("Partition replica list cannot be empty");
         }
-        Set<TopicPartition> duplicateReassignedPartitions = ToolsUtils.duplicates(partitionsToBeReassigned.stream().map(t -> t.v1).collect(Collectors.toList()));
+        Set<TopicPartition> duplicateReassignedPartitions = ToolsUtils.duplicates(partitionsToBeReassigned.stream().map(t -> t.getKey()).collect(Collectors.toList()));
         if (!duplicateReassignedPartitions.isEmpty()) {
             throw new AdminCommandFailedException(String.format(
                 "Partition reassignment contains duplicate topic partitions: %s",
                 duplicateReassignedPartitions.stream().map(Object::toString).collect(Collectors.joining(",")))
             );
         }
-        List<Tuple2<TopicPartition, Set<Integer>>> duplicateEntries = partitionsToBeReassigned.stream()
-            .map(t -> new Tuple2<>(t.v1, ToolsUtils.duplicates(t.v2)))
-            .filter(t -> !t.v2.isEmpty())
+        List<Entry<TopicPartition, Set<Integer>>> duplicateEntries = partitionsToBeReassigned.stream()
+            .map(t -> new SimpleImmutableEntry<>(t.getKey(), ToolsUtils.duplicates(t.getValue())))
+            .filter(t -> !t.getValue().isEmpty())
             .collect(Collectors.toList());
         if (!duplicateEntries.isEmpty()) {
             String duplicatesMsg = duplicateEntries.stream().map(t ->
                 String.format("%s contains multiple entries for %s",
-                    t.v1,
-                    t.v2.stream().map(Object::toString).collect(Collectors.joining(",")))
+                    t.getKey(),
+                    t.getValue().stream().map(Object::toString).collect(Collectors.joining(",")))
             ).collect(Collectors.joining(". "));
             throw new AdminCommandFailedException(String.format("Partition replica lists may not contain duplicate entries: %s", duplicatesMsg));
         }
-        return new Tuple2<>(partitionsToBeReassigned.stream().collect(Collectors.toMap(t -> t.v1, t -> t.v2)), replicaAssignment);
+        return new SimpleImmutableEntry<>(partitionsToBeReassigned.stream().collect(Collectors.toMap(t -> t.getKey(), t -> t.getValue())), replicaAssignment);
     }
 
     /**
@@ -1226,17 +1227,17 @@ static Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, Str
      * @return                      A tuple of the partition reassignments that were cancelled,
      *                              and the replica movements that were cancelled.
      */
-    static Tuple2<Set<TopicPartition>, Set<TopicPartitionReplica>> cancelAssignment(Admin adminClient,
+    static Entry<Set<TopicPartition>, Set<TopicPartitionReplica>> cancelAssignment(Admin adminClient,
                                                                                     String jsonString,
                                                                                     Boolean preserveThrottles,
                                                                                     Long timeoutMs,
                                                                                     Time time
     ) throws ExecutionException, InterruptedException, JsonProcessingException, TerseException {
-        Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(jsonString);
+        Entry<List<Entry<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> t0 = parsePartitionReassignmentData(jsonString);
 
-        List<Tuple2<TopicPartition, List<Integer>>> targetParts = t0.v1;
-        Map<TopicPartitionReplica, String> targetReplicas = t0.v2;
-        Set<TopicPartition> targetPartsSet = targetParts.stream().map(t -> t.v1).collect(Collectors.toSet());
+        List<Entry<TopicPartition, List<Integer>>> targetParts = t0.getKey();
+        Map<TopicPartitionReplica, String> targetReplicas = t0.getValue();
+        Set<TopicPartition> targetPartsSet = targetParts.stream().map(t -> t.getKey()).collect(Collectors.toSet());
         Set<TopicPartition> curReassigningParts = new HashSet<>();
         adminClient.listPartitionReassignments(targetPartsSet).reassignments().get().forEach((part, reassignment) -> {
             if (reassignment.addingReplicas().isEmpty() || !reassignment.removingReplicas().isEmpty())
@@ -1273,7 +1274,7 @@ static Tuple2<Set<TopicPartition>, Set<TopicPartitionReplica>> cancelAssignment(
         if (!preserveThrottles) {
             clearAllThrottles(adminClient, targetParts);
         }
-        return new Tuple2<>(curReassigningParts, curMovingParts.keySet());
+        return new SimpleImmutableEntry<>(curReassigningParts, curMovingParts.keySet());
     }
 
     public static String formatAsReassignmentJson(Map<TopicPartition, List<Integer>> partitionsToBeReassigned,
@@ -1330,7 +1331,7 @@ private static List<String> parseTopicsData(int version, JsonValue js) throws Js
         }
     }
 
-    private static Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> parsePartitionReassignmentData(
+    private static Entry<List<Entry<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> parsePartitionReassignmentData(
         String jsonData
     ) throws JsonProcessingException {
         JsonValue js;
@@ -1344,12 +1345,12 @@ private static Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPart
     }
 
     // Parses without deduplicating keys so the data can be checked before allowing reassignment to proceed
-    private static Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> parsePartitionReassignmentData(
+    private static Entry<List<Entry<TopicPartition, List<Integer>>>, Map<TopicPartitionReplica, String>> parsePartitionReassignmentData(
         int version, JsonValue jsonData
     ) throws JsonMappingException {
         switch (version) {
             case 1:
-                List<Tuple2<TopicPartition, List<Integer>>> partitionAssignment = new ArrayList<>();
+                List<Entry<TopicPartition, List<Integer>>> partitionAssignment = new ArrayList<>();
                 Map<TopicPartitionReplica, String> replicaAssignment = new HashMap<>();
 
                 Optional<JsonValue> partitionsSeq = jsonData.asJsonObject().get("partitions");
@@ -1369,7 +1370,7 @@ private static Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPart
                         if (newReplicas.size() != newLogDirs.size())
                             throw new AdminCommandFailedException("Size of replicas list " + newReplicas + " is different from " +
                                 "size of log dirs list " + newLogDirs + " for partition " + new TopicPartition(topic, partition));
-                        partitionAssignment.add(new Tuple2<>(new TopicPartition(topic, partition), newReplicas));
+                        partitionAssignment.add(new SimpleImmutableEntry<>(new TopicPartition(topic, partition), newReplicas));
                         for (int i = 0; i < newLogDirs.size(); i++) {
                             Integer replica = newReplicas.get(i);
                             String logDir = newLogDirs.get(i);
@@ -1382,7 +1383,7 @@ private static Tuple2<List<Tuple2<TopicPartition, List<Integer>>>, Map<TopicPart
                     }
                 }
 
-                return new Tuple2<>(partitionAssignment, replicaAssignment);
+                return new SimpleImmutableEntry<>(partitionAssignment, replicaAssignment);
 
             default:
                 throw new AdminOperationException("Not supported version field value " + version);
@@ -1481,7 +1482,7 @@ static Set<TopicPartitionReplica> alterReplicaLogDirs(Admin adminClient,
         Set<TopicPartitionReplica> results = new HashSet<>();
         Map<TopicPartitionReplica, KafkaFuture<Void>> values = adminClient.alterReplicaLogDirs(assignment).values();
 
-        for (Map.Entry<TopicPartitionReplica, KafkaFuture<Void>> e : values.entrySet()) {
+        for (Entry<TopicPartitionReplica, KafkaFuture<Void>> e : values.entrySet()) {
             TopicPartitionReplica replica = e.getKey();
             KafkaFuture<Void> future = e.getValue();
             try {
diff --git a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
index 0bb2c9c9d1b71..d427828f2c7f1 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ToolsTestUtils.java
@@ -31,12 +31,14 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
@@ -155,7 +157,7 @@ public static void removeReplicationThrottleForPartitions(Admin admin, List<Inte
     }
 
     public static void assignThrottledPartitionReplicas(Admin adminClient, Map<TopicPartition, List<Integer>> allReplicasByPartition) throws InterruptedException, ExecutionException {
-        Map<ConfigResource, List<Map.Entry<TopicPartition, List<Integer>>>> configResourceToPartitionReplicas =
+        Map<ConfigResource, List<Entry<TopicPartition, List<Integer>>>> configResourceToPartitionReplicas =
             allReplicasByPartition.entrySet().stream()
             .collect(Collectors.groupingBy(
                 topicPartitionListEntry -> new ConfigResource(ConfigResource.Type.TOPIC, topicPartitionListEntry.getKey().topic()))
@@ -163,10 +165,10 @@ public static void assignThrottledPartitionReplicas(Admin adminClient, Map<Topic
 
         Map<ConfigResource, List<AlterConfigOp>> throttles = configResourceToPartitionReplicas.entrySet().stream()
             .collect(
-                Collectors.toMap(Map.Entry::getKey, entry -> {
+                Collectors.toMap(Entry::getKey, entry -> {
                     List<AlterConfigOp> alterConfigOps = new ArrayList<>();
                     Map<TopicPartition, List<Integer>> replicaThrottle =
-                        entry.getValue().stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+                        entry.getValue().stream().collect(Collectors.toMap(Entry::getKey, Entry::getValue));
                     alterConfigOps.add(new AlterConfigOp(
                         new ConfigEntry(LogConfig.LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG, formatReplicaThrottles(replicaThrottle)),
                         AlterConfigOp.OpType.SET));
@@ -204,7 +206,7 @@ public static String formatReplicaThrottles(Map<TopicPartition, List<Integer>> m
 
     public static File tempPropertiesFile(Map<String, String> properties) throws IOException {
         StringBuilder sb = new StringBuilder();
-        for (Map.Entry<String, String> entry : properties.entrySet()) {
+        for (Entry<String, String> entry : properties.entrySet()) {
             sb.append(entry.getKey() + "=" + entry.getValue() + System.lineSeparator());
         }
         return org.apache.kafka.test.TestUtils.tempFile(sb.toString());
@@ -249,7 +251,7 @@ public static String grabConsoleError(Runnable f) {
     /**
      * Capture both the console output and console error during the execution of the provided function.
      */
-    public static Tuple2<String, String> grabConsoleOutputAndError(Runnable f) {
+    public static Entry<String, String> grabConsoleOutputAndError(Runnable f) {
         ByteArrayOutputStream outBuf = new ByteArrayOutputStream();
         ByteArrayOutputStream errBuf = new ByteArrayOutputStream();
         PrintStream out = new PrintStream(outBuf);
@@ -267,7 +269,7 @@ public static Tuple2<String, String> grabConsoleOutputAndError(Runnable f) {
         }
         out.flush();
         err.flush();
-        return new Tuple2<>(outBuf.toString(), errBuf.toString());
+        return new SimpleImmutableEntry<>(outBuf.toString(), errBuf.toString());
     }
 
     public static class MockExitProcedure implements Exit.Procedure {
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
index b7583c1c44f55..4fd7e3b91970e 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupServiceTest.java
@@ -38,7 +38,6 @@
 import org.apache.kafka.common.TopicPartitionInfo;
 import org.apache.kafka.common.internals.KafkaFutureImpl;
 import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentMatcher;
 import org.mockito.ArgumentMatchers;
@@ -49,11 +48,12 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Objects;
-import java.util.Set;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -91,10 +91,10 @@ public void testAdminRequestsForDescribeOffsets() throws Exception {
         when(admin.listOffsets(offsetsArgMatcher(), any()))
                 .thenReturn(listOffsetsResult());
 
-        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
-        assertEquals(Optional.of("Stable"), statesAndAssignments.v1);
-        assertTrue(statesAndAssignments.v2.isPresent());
-        assertEquals(TOPIC_PARTITIONS.size(), statesAndAssignments.v2.get().size());
+        Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
+        assertEquals(Optional.of("Stable"), statesAndAssignments.getKey());
+        assertTrue(statesAndAssignments.getValue().isPresent());
+        assertEquals(TOPIC_PARTITIONS.size(), statesAndAssignments.getValue().get().size());
 
         verify(admin, times(1)).describeConsumerGroups(ArgumentMatchers.eq(Collections.singletonList(GROUP)), any());
         verify(admin, times(1)).listConsumerGroupOffsets(ArgumentMatchers.eq(listConsumerGroupOffsetsSpec()), any());
@@ -157,16 +157,16 @@ public void testAdminRequestsForDescribeNegativeOffsets() throws Exception {
                 ArgumentMatchers.argThat(offsetsArgMatcher.apply(assignedTopicPartitions)),
                 any()
         )).thenReturn(new ListOffsetsResult(endOffsets.entrySet().stream().filter(e -> assignedTopicPartitions.contains(e.getKey()))
-                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))));
+                .collect(Collectors.toMap(Entry::getKey, Entry::getValue))));
         when(admin.listOffsets(
                 ArgumentMatchers.argThat(offsetsArgMatcher.apply(unassignedTopicPartitions)),
                 any()
         )).thenReturn(new ListOffsetsResult(endOffsets.entrySet().stream().filter(e -> unassignedTopicPartitions.contains(e.getKey()))
-                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))));
+                .collect(Collectors.toMap(Entry::getKey, Entry::getValue))));
 
-        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
-        Optional<String> state = statesAndAssignments.v1;
-        Optional<Collection<PartitionAssignmentState>> assignments = statesAndAssignments.v2;
+        Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> statesAndAssignments = groupService.collectGroupOffsets(GROUP);
+        Optional<String> state = statesAndAssignments.getKey();
+        Optional<Collection<PartitionAssignmentState>> assignments = statesAndAssignments.getValue();
 
         Map<TopicPartition, Optional<Long>> returnedOffsets = assignments.map(results ->
             results.stream().collect(Collectors.toMap(
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
index 95fe212c53133..b8133433d6489 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteConsumerGroupsTest.java
@@ -92,7 +92,7 @@ public void testDeleteCmdNonEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.collectGroupMembers(GROUP, false).v2.get().size() == 1,
+            () -> service.collectGroupMembers(GROUP, false).getValue().get().size() == 1,
             "The group did not initialize as expected."
         );
 
@@ -112,7 +112,7 @@ public void testDeleteNonEmptyGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(
-            () -> service.collectGroupMembers(GROUP, false).v2.get().size() == 1,
+            () -> service.collectGroupMembers(GROUP, false).getValue().get().size() == 1,
             "The group did not initialize as expected."
         );
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
index 05fd00071e15e..292180b8af90b 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DeleteOffsetsConsumerGroupCommandIntegrationTest.java
@@ -29,12 +29,12 @@
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.config.Defaults;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.util.Collections;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Properties;
 import java.util.concurrent.ExecutionException;
 
@@ -60,8 +60,8 @@ public void testDeleteOffsetsNonExistingGroup() {
         String topic = "foo:1";
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(group, topic));
 
-        Tuple2<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(group, Collections.singletonList(topic));
-        assertEquals(Errors.GROUP_ID_NOT_FOUND, res.v1);
+        Entry<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(group, Collections.singletonList(topic));
+        assertEquals(Errors.GROUP_ID_NOT_FOUND, res.getKey());
     }
 
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
@@ -145,9 +145,9 @@ private void testWithConsumerGroup(java.util.function.Consumer<Runnable> withCon
         withConsumerGroup.accept(() -> {
             String topic = inputPartition >= 0 ? inputTopic + ":" + inputPartition : inputTopic;
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(getArgs(GROUP, topic));
-            Tuple2<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(GROUP, Collections.singletonList(topic));
-            Errors topLevelError = res.v1;
-            Map<TopicPartition, Throwable> partitions = res.v2;
+            Entry<Errors, Map<TopicPartition, Throwable>> res = service.deleteOffsets(GROUP, Collections.singletonList(topic));
+            Errors topLevelError = res.getKey();
+            Map<TopicPartition, Throwable> partitions = res.getValue();
             TopicPartition tp = new TopicPartition(inputTopic, expectedPartition);
             // Partition level error should propagate to top level, unless this is due to a missed partition attempt.
             if (inputPartition >= 0) {
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
index 85f9f5b836a16..f1b6bba9ac1af 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
@@ -24,7 +24,6 @@
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.test.TestUtils;
 import org.apache.kafka.tools.ToolsTestUtils;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -34,6 +33,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
@@ -138,8 +138,8 @@ public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupPro
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(group);
-        assertTrue(res.v1.map(s -> s.contains("Dead")).orElse(false) && res.v2.map(Collection::isEmpty).orElse(false),
+        Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(group);
+        assertTrue(res.getKey().map(s -> s.contains("Dead")).orElse(false) && res.getValue().map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
     }
 
@@ -155,12 +155,12 @@ public void testDescribeMembersOfNonExistingGroup(String quorum, String groupPro
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group};
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(group, false);
-        assertTrue(res.v1.map(s -> s.contains("Dead")).orElse(false) && res.v2.map(Collection::isEmpty).orElse(false),
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(group, false);
+        assertTrue(res.getKey().map(s -> s.contains("Dead")).orElse(false) && res.getValue().map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "'.");
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res2 = service.collectGroupMembers(group, true);
-        assertTrue(res2.v1.map(s -> s.contains("Dead")).orElse(false) && res2.v2.map(Collection::isEmpty).orElse(false),
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res2 = service.collectGroupMembers(group, true);
+        assertTrue(res2.getKey().map(s -> s.contains("Dead")).orElse(false) && res2.getValue().map(Collection::isEmpty).orElse(false),
             "Expected the state to be 'Dead', with no members in the group '" + group + "' (verbose option).");
     }
 
@@ -197,8 +197,8 @@ public void testDescribeExistingGroup(String quorum, String groupProtocol) throw
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
-                return res.v1.trim().split("\n").length == 2 && res.v2.isEmpty();
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res.getKey().trim().split("\n").length == 2 && res.getValue().isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -226,9 +226,9 @@ public void testDescribeExistingGroups(String quorum, String groupProtocol) thro
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
-                long numLines = Arrays.stream(res.v1.trim().split("\n")).filter(line -> !line.isEmpty()).count();
-                return (numLines == expectedNumLines) && res.v2.isEmpty();
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res.getKey().trim().split("\n")).filter(line -> !line.isEmpty()).count();
+                return (numLines == expectedNumLines) && res.getValue().isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -252,9 +252,9 @@ public void testDescribeAllExistingGroups(String quorum, String groupProtocol) t
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
-                long numLines = Arrays.stream(res.v1.trim().split("\n")).filter(s -> !s.isEmpty()).count();
-                return (numLines == expectedNumLines) && res.v2.isEmpty();
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                long numLines = Arrays.stream(res.getKey().trim().split("\n")).filter(s -> !s.isEmpty()).count();
+                return (numLines == expectedNumLines) && res.getValue().isEmpty();
             }, "Expected a data row and no error in describe results with describe type " + String.join(" ", describeType) + ".");
         }
     }
@@ -271,9 +271,9 @@ public void testDescribeOffsetsOfExistingGroup(String quorum, String groupProtoc
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
-            Optional<String> state = groupOffsets.v1;
-            Optional<Collection<PartitionAssignmentState>> assignments = groupOffsets.v2;
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+            Optional<String> state = groupOffsets.getKey();
+            Optional<Collection<PartitionAssignmentState>> assignments = groupOffsets.getValue();
 
             Predicate<PartitionAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
@@ -307,9 +307,9 @@ public void testDescribeMembersOfExistingGroup(String quorum, String groupProtoc
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> groupMembers = service.collectGroupMembers(GROUP, false);
-            Optional<String> state = groupMembers.v1;
-            Optional<Collection<MemberAssignmentState>> assignments = groupMembers.v2;
+            Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> groupMembers = service.collectGroupMembers(GROUP, false);
+            Optional<String> state = groupMembers.getKey();
+            Optional<Collection<MemberAssignmentState>> assignments = groupMembers.getValue();
 
             Predicate<MemberAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
@@ -331,10 +331,10 @@ public void testDescribeMembersOfExistingGroup(String quorum, String groupProtoc
                 !Objects.equals(assignmentState.host, ConsumerGroupCommand.MISSING_COLUMN_VALUE);
         }, "Expected a 'Stable' group status, rows and valid member information for group " + GROUP + ".");
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
 
-        if (res.v2.isPresent()) {
-            assertTrue(res.v2.get().size() == 1 && res.v2.get().iterator().next().assignment.size() == 1,
+        if (res.getValue().isPresent()) {
+            assertTrue(res.getValue().get().size() == 1 && res.getValue().get().iterator().next().assignment.size() == 1,
                 "Expected a topic partition assigned to the single group member for group " + GROUP);
         } else {
             fail("Expected partition assignments for members of group " + GROUP);
@@ -407,8 +407,8 @@ public void testDescribeExistingGroupWithNoMembers(String quorum, String groupPr
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
-                return res.v1.trim().split("\n").length == 2 && res.v2.isEmpty();
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                return res.getKey().trim().split("\n").length == 2 && res.getValue().isEmpty();
             }, "Expected describe group results with one data row for describe type '" + String.join(" ", describeType) + "'");
 
             // stop the consumer so the group has no active member anymore
@@ -431,18 +431,18 @@ public void testDescribeOffsetsOfExistingGroupWithNoMembers(String quorum, Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res.v1.map(s -> s.contains("Stable")).orElse(false)
-                && res.v2.map(c -> c.stream().anyMatch(assignment -> Objects.equals(assignment.group, GROUP) && assignment.offset.isPresent())).orElse(false);
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.getKey().map(s -> s.contains("Stable")).orElse(false)
+                && res.getValue().map(c -> c.stream().anyMatch(assignment -> Objects.equals(assignment.group, GROUP) && assignment.offset.isPresent())).orElse(false);
         }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
 
         // stop the consumer so the group has no active member anymore
         executor.shutdown();
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> offsets = service.collectGroupOffsets(GROUP);
-            Optional<String> state = offsets.v1;
-            Optional<Collection<PartitionAssignmentState>> assignments = offsets.v2;
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> offsets = service.collectGroupOffsets(GROUP);
+            Optional<String> state = offsets.getKey();
+            Optional<Collection<PartitionAssignmentState>> assignments = offsets.getValue();
             List<PartitionAssignmentState> testGroupAssignments = assignments.get().stream().filter(a -> Objects.equals(a.group, GROUP)).collect(Collectors.toList());
             PartitionAssignmentState assignment = testGroupAssignments.get(0);
             return state.map(s -> s.contains("Empty")).orElse(false) &&
@@ -465,17 +465,17 @@ public void testDescribeMembersOfExistingGroupWithNoMembers(String quorum, Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res.v1.map(s -> s.contains("Stable")).orElse(false)
-                && res.v2.map(c -> c.stream().anyMatch(m -> Objects.equals(m.group, GROUP))).orElse(false);
+            Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.getKey().map(s -> s.contains("Stable")).orElse(false)
+                && res.getValue().map(c -> c.stream().anyMatch(m -> Objects.equals(m.group, GROUP))).orElse(false);
         }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.");
 
         // stop the consumer so the group has no active member anymore
         executor.shutdown();
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res.v1.map(s -> s.contains("Empty")).orElse(false) && res.v2.isPresent() && res.v2.get().isEmpty();
+            Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.getKey().map(s -> s.contains("Empty")).orElse(false) && res.getValue().isPresent() && res.getValue().get().isEmpty();
         }, "Expected no member in describe group members results for group '" + GROUP + "'");
     }
 
@@ -521,9 +521,9 @@ public void testDescribeWithConsumersWithoutAssignedPartitions(String quorum, St
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
                 int expectedNumRows = DESCRIBE_TYPE_MEMBERS.contains(describeType) ? 3 : 2;
-                return res.v2.isEmpty() && res.v1.trim().split("\n").length == expectedNumRows;
+                return res.getValue().isEmpty() && res.getKey().trim().split("\n").length == expectedNumRows;
             }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
         }
     }
@@ -540,11 +540,11 @@ public void testDescribeOffsetsWithConsumersWithoutAssignedPartitions(String quo
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res.v1.map(s -> s.contains("Stable")).isPresent() &&
-                res.v2.isPresent() &&
-                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 1 &&
-                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 1;
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.getKey().map(s -> s.contains("Stable")).isPresent() &&
+                res.getValue().isPresent() &&
+                res.getValue().get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 1 &&
+                res.getValue().get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 1;
         }, "Expected rows for consumers with no assigned partitions in describe group results");
     }
 
@@ -560,18 +560,18 @@ public void testDescribeMembersWithConsumersWithoutAssignedPartitions(String quo
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
-                res.v2.isPresent() &&
-                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
-                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 1 &&
-                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0).count() == 1 &&
-                res.v2.get().stream().allMatch(s -> s.assignment.isEmpty());
+            Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.getKey().map(s -> s.contains("Stable")).orElse(false) &&
+                res.getValue().isPresent() &&
+                res.getValue().get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.getValue().get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 1 &&
+                res.getValue().get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0).count() == 1 &&
+                res.getValue().get().stream().allMatch(s -> s.assignment.isEmpty());
         }, "Expected rows for consumers with no assigned partitions in describe group results");
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
-        assertTrue(res.v1.map(s -> s.contains("Stable")).orElse(false)
-                && res.v2.map(c -> c.stream().anyMatch(s -> !s.assignment.isEmpty())).orElse(false),
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res.getKey().map(s -> s.contains("Stable")).orElse(false)
+                && res.getValue().map(c -> c.stream().anyMatch(s -> !s.assignment.isEmpty())).orElse(false),
             "Expected additional columns in verbose version of describe members");
     }
 
@@ -608,9 +608,9 @@ public void testDescribeWithMultiPartitionTopicAndMultipleConsumers(String quoru
             ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs.toArray(new String[0]));
 
             TestUtils.waitForCondition(() -> {
-                Tuple2<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
+                Entry<String, String> res = ToolsTestUtils.grabConsoleOutputAndError(describeGroups(service));
                 int expectedNumRows = DESCRIBE_TYPE_STATE.contains(describeType) ? 2 : 3;
-                return res.v2.isEmpty() && res.v1.trim().split("\n").length == expectedNumRows;
+                return res.getValue().isEmpty() && res.getKey().trim().split("\n").length == expectedNumRows;
             }, "Expected a single data row in describe group result with describe type '" + String.join(" ", describeType) + "'");
         }
     }
@@ -629,12 +629,12 @@ public void testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
-                res.v2.isPresent() &&
-                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
-                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 2 &&
-                res.v2.get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && !x.partition.isPresent());
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.getKey().map(s -> s.contains("Stable")).orElse(false) &&
+                res.getValue().isPresent() &&
+                res.getValue().get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.getValue().get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.partition.isPresent()).count() == 2 &&
+                res.getValue().get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && !x.partition.isPresent());
         }, "Expected two rows (one row per consumer) in describe group results.");
     }
 
@@ -652,16 +652,16 @@ public void testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(Strin
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
-            return res.v1.map(s -> s.contains("Stable")).orElse(false) &&
-                res.v2.isPresent() &&
-                res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
-                res.v2.get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 2 &&
-                res.v2.get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0);
+            Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, false);
+            return res.getKey().map(s -> s.contains("Stable")).orElse(false) &&
+                res.getValue().isPresent() &&
+                res.getValue().get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2 &&
+                res.getValue().get().stream().filter(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 1).count() == 2 &&
+                res.getValue().get().stream().noneMatch(x -> Objects.equals(x.group, GROUP) && x.numPartitions == 0);
         }, "Expected two rows (one row per consumer) in describe group members results.");
 
-        Tuple2<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
-        assertTrue(res.v1.map(s -> s.contains("Stable")).orElse(false) && res.v2.map(s -> s.stream().filter(x -> x.assignment.isEmpty()).count()).orElse(0L) == 0,
+        Entry<Optional<String>, Optional<Collection<MemberAssignmentState>>> res = service.collectGroupMembers(GROUP, true);
+        assertTrue(res.getKey().map(s -> s.contains("Stable")).orElse(false) && res.getValue().map(s -> s.stream().filter(x -> x.assignment.isEmpty()).count()).orElse(0L) == 0,
             "Expected additional columns in verbose version of describe members");
     }
 
@@ -698,9 +698,9 @@ public void testDescribeSimpleConsumerGroup(String quorum) throws Exception {
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
-            return res.v1.map(s -> s.contains("Empty")).orElse(false)
-                && res.v2.isPresent() && res.v2.get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2;
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> res = service.collectGroupOffsets(GROUP);
+            return res.getKey().map(s -> s.contains("Empty")).orElse(false)
+                && res.getValue().isPresent() && res.getValue().get().stream().filter(s -> Objects.equals(s.group, GROUP)).count() == 2;
         }, "Expected a stable group with two members in describe group state result.");
     }
 
@@ -798,18 +798,18 @@ public void testDescribeNonOffsetCommitGroup(String quorum, String groupProtocol
         ConsumerGroupCommand.ConsumerGroupService service = getConsumerGroupService(cgcArgs);
 
         TestUtils.waitForCondition(() -> {
-            Tuple2<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
+            Entry<Optional<String>, Optional<Collection<PartitionAssignmentState>>> groupOffsets = service.collectGroupOffsets(GROUP);
 
             Predicate<PartitionAssignmentState> isGrp = s -> Objects.equals(s.group, GROUP);
 
-            boolean res = groupOffsets.v1.map(s -> s.contains("Stable")).orElse(false) &&
-                groupOffsets.v2.isPresent() &&
-                groupOffsets.v2.get().stream().filter(isGrp).count() == 1;
+            boolean res = groupOffsets.getKey().map(s -> s.contains("Stable")).orElse(false) &&
+                groupOffsets.getValue().isPresent() &&
+                groupOffsets.getValue().get().stream().filter(isGrp).count() == 1;
 
             if (!res)
                 return false;
 
-            Optional<PartitionAssignmentState> maybeAssignmentState = groupOffsets.v2.get().stream().filter(isGrp).findFirst();
+            Optional<PartitionAssignmentState> maybeAssignmentState = groupOffsets.getValue().get().stream().filter(isGrp).findFirst();
             if (!maybeAssignmentState.isPresent())
                 return false;
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java
index 618cf972995b7..c7c4a4f35e65e 100644
--- a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java
+++ b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java
@@ -217,7 +217,7 @@ public void run(ExperimentDef config, Journal journal, boolean displayChartsOnSc
 
             System.out.println("Generating Reassignment");
             Map<TopicPartition, List<Integer>> newAssignment = ReassignPartitionsCommand.generateAssignment(adminClient,
-                json(TOPIC_NAME), brokers.stream().map(Object::toString).collect(Collectors.joining(",")), true).v1;
+                json(TOPIC_NAME), brokers.stream().map(Object::toString).collect(Collectors.joining(",")), true).getKey();
 
             System.out.println("Starting Reassignment");
             long start = System.currentTimeMillis();
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
index 83fc665e3e487..0348aa24ffe9f 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
@@ -43,7 +43,6 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.tools.TerseException;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -54,6 +53,7 @@
 import scala.collection.Seq;
 
 import java.io.Closeable;
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -62,6 +62,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
@@ -395,11 +396,11 @@ public void testCancellation(String quorum) throws Exception {
         waitForVerifyAssignment(cluster.adminClient, assignment, true,
             new VerifyAssignmentResult(partStates, true, Collections.emptyMap(), false));
         // Cancel the reassignment.
-        assertEquals(new Tuple2<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, true));
+        assertEquals(new SimpleImmutableEntry<>(new HashSet<>(asList(foo0, baz1)), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, true));
         // Broker throttles are still active because we passed --preserve-throttles
         waitForInterBrokerThrottle(asList(0, 1, 2, 3), interBrokerThrottle);
         // Cancelling the reassignment again should reveal nothing to cancel.
-        assertEquals(new Tuple2<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, false));
+        assertEquals(new SimpleImmutableEntry<>(Collections.emptySet(), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, false));
         // This time, the broker throttles were removed.
         waitForBrokerLevelThrottles(unthrottledBrokerConfigs);
         // Verify that there are no ongoing reassignments.
@@ -446,7 +447,7 @@ public void testCancellationWithAddingReplicaInIsr(String quorum) throws Excepti
         );
 
         // Now cancel the assignment and verify that the partition is removed from cancelled replicas
-        assertEquals(new Tuple2<>(Collections.singleton(foo0), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, true));
+        assertEquals(new SimpleImmutableEntry<>(Collections.singleton(foo0), Collections.emptySet()), runCancelAssignment(cluster.adminClient, assignment, true));
         verifyReplicaDeleted(foo0, 3);
         verifyReplicaDeleted(foo0, 4);
     }
@@ -708,7 +709,7 @@ private void runExecuteAssignment(Admin adminClient,
         }
     }
 
-    private Tuple2<Set<TopicPartition>, Set<TopicPartitionReplica>> runCancelAssignment(
+    private Entry<Set<TopicPartition>, Set<TopicPartitionReplica>> runCancelAssignment(
         Admin adminClient,
         String jsonString,
         Boolean preserveThrottles
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
index 699e048fb206e..c6f145d9a5598 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
@@ -31,18 +31,19 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.common.AdminCommandFailedException;
 import org.apache.kafka.server.common.AdminOperationException;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
@@ -174,14 +175,14 @@ public void testFindPartitionReassignmentStates() throws Exception {
             expStates.put(new TopicPartition("foo", 1),
                 new PartitionReassignmentState(asList(1, 2, 3), asList(1, 2, 3), true));
 
-            Tuple2<Map<TopicPartition, PartitionReassignmentState>, Boolean> actual =
+            Entry<Map<TopicPartition, PartitionReassignmentState>, Boolean> actual =
                 findPartitionReassignmentStates(adminClient, asList(
-                    new Tuple2<>(new TopicPartition("foo", 0), asList(0, 1, 3)),
-                    new Tuple2<>(new TopicPartition("foo", 1), asList(1, 2, 3))
+                    new SimpleImmutableEntry<>(new TopicPartition("foo", 0), asList(0, 1, 3)),
+                    new SimpleImmutableEntry<>(new TopicPartition("foo", 1), asList(1, 2, 3))
                 ));
 
-            assertEquals(expStates, actual.v1);
-            assertTrue(actual.v2);
+            assertEquals(expStates, actual.getKey());
+            assertTrue(actual.getValue());
 
             // Cancel the reassignment and test findPartitionReassignmentStates again.
             Map<TopicPartition, Throwable> cancelResult = cancelPartitionReassignments(adminClient,
@@ -198,12 +199,12 @@ public void testFindPartitionReassignmentStates() throws Exception {
                 new PartitionReassignmentState(asList(1, 2, 3), asList(1, 2, 3), true));
 
             actual = findPartitionReassignmentStates(adminClient, asList(
-                new Tuple2<>(new TopicPartition("foo", 0), asList(0, 1, 3)),
-                new Tuple2<>(new TopicPartition("foo", 1), asList(1, 2, 3))
+                new SimpleImmutableEntry<>(new TopicPartition("foo", 0), asList(0, 1, 3)),
+                new SimpleImmutableEntry<>(new TopicPartition("foo", 1), asList(1, 2, 3))
             ));
 
-            assertEquals(expStates, actual.v1);
-            assertFalse(actual.v2);
+            assertEquals(expStates, actual.getKey());
+            assertFalse(actual.getValue());
         }
     }
 
@@ -338,13 +339,13 @@ public void testParseGenerateAssignmentArgs() throws Exception {
             assertThrows(AdminCommandFailedException.class, () -> parseGenerateAssignmentArgs(
                 "{\"topics\": [{\"topic\": \"foo\"}], \"version\":1}", "5,2,3,4,5"),
                 "Expected to detect duplicate broker list entries").getMessage());
-        assertEquals(new Tuple2<>(asList(5, 2, 3, 4), asList("foo")),
+        assertEquals(new SimpleImmutableEntry<>(asList(5, 2, 3, 4), asList("foo")),
             parseGenerateAssignmentArgs("{\"topics\": [{\"topic\": \"foo\"}], \"version\":1}", "5,2,3,4"));
         assertStartsWith("List of topics to reassign contains duplicate entries",
             assertThrows(AdminCommandFailedException.class, () -> parseGenerateAssignmentArgs(
                 "{\"topics\": [{\"topic\": \"foo\"},{\"topic\": \"foo\"}], \"version\":1}", "5,2,3,4"),
                 "Expected to detect duplicate topic entries").getMessage());
-        assertEquals(new Tuple2<>(asList(5, 3, 4), asList("foo", "bar")),
+        assertEquals(new SimpleImmutableEntry<>(asList(5, 3, 4), asList("foo", "bar")),
             parseGenerateAssignmentArgs(
                 "{\"topics\": [{\"topic\": \"foo\"},{\"topic\": \"bar\"}], \"version\":1}", "5,3,4"));
     }
@@ -389,7 +390,7 @@ public void testGenerateAssignmentWithInconsistentRacks() throws Exception {
                     () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", true),
                     "Expected generateAssignment to fail").getMessage());
             // It should succeed when --disable-rack-aware is used.
-            Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>>
+            Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>>
                 proposedCurrent = generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", false);
 
             Map<TopicPartition, List<Integer>> expCurrent = new HashMap<>();
@@ -397,7 +398,7 @@ public void testGenerateAssignmentWithInconsistentRacks() throws Exception {
             expCurrent.put(new TopicPartition("foo", 0), asList(0, 1, 2));
             expCurrent.put(new TopicPartition("foo", 1), asList(1, 2, 3));
 
-            assertEquals(expCurrent, proposedCurrent.v2);
+            assertEquals(expCurrent, proposedCurrent.getValue());
         }
     }
 
@@ -407,7 +408,7 @@ public void testGenerateAssignmentWithFewerBrokers() throws Exception {
             addTopics(adminClient);
             List<Integer> goalBrokers = asList(0, 1, 3);
 
-            Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>>
+            Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>>
                 proposedCurrent = generateAssignment(adminClient,
                     "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"bar\"}]}",
                     goalBrokers.stream().map(Object::toString).collect(Collectors.joining(",")), false);
@@ -418,12 +419,12 @@ public void testGenerateAssignmentWithFewerBrokers() throws Exception {
             expCurrent.put(new TopicPartition("foo", 1), asList(1, 2, 3));
             expCurrent.put(new TopicPartition("bar", 0), asList(2, 3, 0));
 
-            assertEquals(expCurrent, proposedCurrent.v2);
+            assertEquals(expCurrent, proposedCurrent.getValue());
 
             // The proposed assignment should only span the provided brokers
-            proposedCurrent.v1.values().forEach(replicas ->
+            proposedCurrent.getKey().values().forEach(replicas ->
                 assertTrue(goalBrokers.containsAll(replicas),
-                    "Proposed assignment " + proposedCurrent.v1 + " puts replicas on brokers other than " + goalBrokers)
+                    "Proposed assignment " + proposedCurrent.getKey() + " puts replicas on brokers other than " + goalBrokers)
             );
         }
     }
@@ -567,14 +568,14 @@ public void testParseExecuteAssignmentArgs() throws Exception {
         partitionsToBeReassigned.put(new TopicPartition("foo", 0), asList(1, 2, 3));
         partitionsToBeReassigned.put(new TopicPartition("foo", 1), asList(3, 4, 5));
 
-        Tuple2<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> actual = parseExecuteAssignmentArgs(
+        Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartitionReplica, String>> actual = parseExecuteAssignmentArgs(
             "{\"version\":1,\"partitions\":" +
                 "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[1,2,3],\"log_dirs\":[\"any\",\"any\",\"any\"]}," +
                 "{\"topic\":\"foo\",\"partition\":1,\"replicas\":[3,4,5],\"log_dirs\":[\"any\",\"any\",\"any\"]}" +
                 "]}");
 
-        assertEquals(partitionsToBeReassigned, actual.v1);
-        assertTrue(actual.v2.isEmpty());
+        assertEquals(partitionsToBeReassigned, actual.getKey());
+        assertTrue(actual.getValue().isEmpty());
 
         Map<TopicPartitionReplica, String> replicaAssignment = new HashMap<>();
 
@@ -587,8 +588,8 @@ public void testParseExecuteAssignmentArgs() throws Exception {
                 "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[1,2,3],\"log_dirs\":[\"/tmp/a\",\"/tmp/b\",\"/tmp/c\"]}" +
                 "]}");
 
-        assertEquals(Collections.singletonMap(new TopicPartition("foo", 0), asList(1, 2, 3)), actual.v1);
-        assertEquals(replicaAssignment, actual.v2);
+        assertEquals(Collections.singletonMap(new TopicPartition("foo", 0), asList(1, 2, 3)), actual.getKey());
+        assertEquals(replicaAssignment, actual.getValue());
     }
 
     @Test

From bf9a27fefdb3d93c7a510f871433c4c9e07de71a Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Sun, 24 Mar 2024 13:09:21 +0800
Subject: [PATCH 194/258] KAFKA-16388 add production-ready test of 3.3 - 3.6
 release to MetadataVersionTest.testFromVersionString (#15563)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/server/common/MetadataVersionTest.java      | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/server-common/src/test/java/org/apache/kafka/server/common/MetadataVersionTest.java b/server-common/src/test/java/org/apache/kafka/server/common/MetadataVersionTest.java
index add6a88b49e02..bd9f2594b0080 100644
--- a/server-common/src/test/java/org/apache/kafka/server/common/MetadataVersionTest.java
+++ b/server-common/src/test/java/org/apache/kafka/server/common/MetadataVersionTest.java
@@ -132,6 +132,8 @@ public void testFromVersionString() {
         assertEquals(IBP_2_6_IV0, MetadataVersion.fromVersionString("2.6"));
         assertEquals(IBP_2_6_IV0, MetadataVersion.fromVersionString("2.6-IV0"));
 
+        // 2.7-IV2 is the latest production version in the 2.7 line
+        assertEquals(IBP_2_7_IV2, MetadataVersion.fromVersionString("2.7"));
         assertEquals(IBP_2_7_IV0, MetadataVersion.fromVersionString("2.7-IV0"));
         assertEquals(IBP_2_7_IV1, MetadataVersion.fromVersionString("2.7-IV1"));
         assertEquals(IBP_2_7_IV2, MetadataVersion.fromVersionString("2.7-IV2"));
@@ -150,24 +152,31 @@ public void testFromVersionString() {
         assertEquals(IBP_3_2_IV0, MetadataVersion.fromVersionString("3.2"));
         assertEquals(IBP_3_2_IV0, MetadataVersion.fromVersionString("3.2-IV0"));
 
+        // 3.3-IV3 is the latest production version in the 3.3 line
+        assertEquals(IBP_3_3_IV3, MetadataVersion.fromVersionString("3.3"));
         assertEquals(IBP_3_3_IV0, MetadataVersion.fromVersionString("3.3-IV0"));
         assertEquals(IBP_3_3_IV1, MetadataVersion.fromVersionString("3.3-IV1"));
         assertEquals(IBP_3_3_IV2, MetadataVersion.fromVersionString("3.3-IV2"));
         assertEquals(IBP_3_3_IV3, MetadataVersion.fromVersionString("3.3-IV3"));
 
+        // 3.4-IV0 is the latest production version in the 3.4 line
+        assertEquals(IBP_3_4_IV0, MetadataVersion.fromVersionString("3.4"));
         assertEquals(IBP_3_4_IV0, MetadataVersion.fromVersionString("3.4-IV0"));
 
+        // 3.5-IV2 is the latest production version in the 3.5 line
+        assertEquals(IBP_3_5_IV2, MetadataVersion.fromVersionString("3.5"));
         assertEquals(IBP_3_5_IV0, MetadataVersion.fromVersionString("3.5-IV0"));
         assertEquals(IBP_3_5_IV1, MetadataVersion.fromVersionString("3.5-IV1"));
         assertEquals(IBP_3_5_IV2, MetadataVersion.fromVersionString("3.5-IV2"));
 
+        // 3.6-IV2 is the latest production version in the 3.6 line
+        assertEquals(IBP_3_6_IV2, MetadataVersion.fromVersionString("3.6"));
         assertEquals(IBP_3_6_IV0, MetadataVersion.fromVersionString("3.6-IV0"));
         assertEquals(IBP_3_6_IV1, MetadataVersion.fromVersionString("3.6-IV1"));
         assertEquals(IBP_3_6_IV2, MetadataVersion.fromVersionString("3.6-IV2"));
 
         // 3.7-IV4 is the latest production version in the 3.7 line
         assertEquals(IBP_3_7_IV4, MetadataVersion.fromVersionString("3.7"));
-
         assertEquals(IBP_3_7_IV0, MetadataVersion.fromVersionString("3.7-IV0"));
         assertEquals(IBP_3_7_IV1, MetadataVersion.fromVersionString("3.7-IV1"));
         assertEquals(IBP_3_7_IV2, MetadataVersion.fromVersionString("3.7-IV2"));

From 0434c29e58ee3019d289745fff3fd1b248224ada Mon Sep 17 00:00:00 2001
From: Dmitry Werner <grimekillah@gmail.com>
Date: Mon, 25 Mar 2024 09:12:23 +0500
Subject: [PATCH 195/258] KAFKA-16408  kafka-get-offsets / GetOffsetShell
 doesn't handle --version or --help (#15583)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../java/org/apache/kafka/tools/GetOffsetShell.java |  5 ++++-
 .../org/apache/kafka/tools/GetOffsetShellTest.java  | 13 +++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/GetOffsetShell.java b/tools/src/main/java/org/apache/kafka/tools/GetOffsetShell.java
index 99551ca0545a4..0d90327bbe73f 100644
--- a/tools/src/main/java/org/apache/kafka/tools/GetOffsetShell.java
+++ b/tools/src/main/java/org/apache/kafka/tools/GetOffsetShell.java
@@ -59,6 +59,7 @@
 import java.util.stream.Collectors;
 
 public class GetOffsetShell {
+    static final String USAGE_TEXT = "An interactive shell for getting topic-partition offsets.";
     private static final Pattern TOPIC_PARTITION_PATTERN = Pattern.compile("([^:,]*)(?::(?:([0-9]*)|(?:([0-9]*)-([0-9]*))))?");
 
     public static void main(String... args) {
@@ -142,7 +143,7 @@ public GetOffsetShellOptions(String[] args) throws TerseException {
             excludeInternalTopicsOpt = parser.accepts("exclude-internal-topics", "By default, internal topics are included. If specified, internal topics are excluded.");
 
             if (args.length == 0) {
-                CommandLineUtils.printUsageAndExit(parser, "An interactive shell for getting topic-partition offsets.");
+                CommandLineUtils.printUsageAndExit(parser, USAGE_TEXT);
             }
 
             try {
@@ -157,6 +158,8 @@ public GetOffsetShellOptions(String[] args) throws TerseException {
                 effectiveBrokerListOpt = brokerListOpt;
             }
 
+            CommandLineUtils.maybePrintHelpOrVersion(this, USAGE_TEXT);
+
             CommandLineUtils.checkRequiredArgs(parser, options, effectiveBrokerListOpt);
 
             String brokerList = options.valueOf(effectiveBrokerListOpt);
diff --git a/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java b/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
index ce9c1718ffb67..dad41342b7be0 100644
--- a/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/GetOffsetShellTest.java
@@ -29,6 +29,7 @@
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.serialization.StringSerializer;
+import org.apache.kafka.common.utils.AppInfoParser;
 import org.apache.kafka.common.utils.Exit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
@@ -317,6 +318,18 @@ public void testTopicPartitionsFlagWithPartitionsFlagCauseExit() {
         assertExitCodeIsOne("--topic-partitions", "__consumer_offsets", "--partitions", "0");
     }
 
+    @ClusterTest
+    public void testPrintHelp() {
+        String out = ToolsTestUtils.captureStandardErr(() -> GetOffsetShell.mainNoExit("--help"));
+        assertTrue(out.startsWith(GetOffsetShell.USAGE_TEXT));
+    }
+
+    @ClusterTest
+    public void testPrintVersion() {
+        String out = ToolsTestUtils.captureStandardOut(() -> GetOffsetShell.mainNoExit("--version"));
+        assertEquals(AppInfoParser.getVersion(), out);
+    }
+
     private void assertExitCodeIsOne(String... args) {
         final int[] exitStatus = new int[1];
 

From 7b2fc469adc58329a96d238bf80eef10f4fb0b22 Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Mon, 25 Mar 2024 12:31:37 +0800
Subject: [PATCH 196/258] KAFKA-16410  kafka-leader-election /
 LeaderElectionCommand doesn't set exit code on error (#15591)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../kafka/tools/LeaderElectionCommand.java    |  9 ++++-
 .../tools/LeaderElectionCommandErrorTest.java | 33 ++++++++++++-------
 .../tools/LeaderElectionCommandTest.java      | 22 ++++++-------
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/LeaderElectionCommand.java b/tools/src/main/java/org/apache/kafka/tools/LeaderElectionCommand.java
index c704e418750c1..e1cdeed1e5885 100644
--- a/tools/src/main/java/org/apache/kafka/tools/LeaderElectionCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/LeaderElectionCommand.java
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.ElectionNotNeededException;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.common.AdminCommandFailedException;
 import org.apache.kafka.server.common.AdminOperationException;
@@ -62,11 +63,17 @@ public class LeaderElectionCommand {
     private static final DecodeJson.DecodeInteger INT = new DecodeJson.DecodeInteger();
 
     public static void main(String... args) {
+        Exit.exit(mainNoExit(args));
+    }
+
+    static int mainNoExit(String... args) {
         try {
             run(Duration.ofMillis(30000), args);
-        } catch (Exception e) {
+            return 0;
+        } catch (Throwable e) {
             System.err.println(e.getMessage());
             System.err.println(Utils.stackTrace(e));
+            return 1;
         }
     }
 
diff --git a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
index f293c465c4008..087621df7dea7 100644
--- a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandErrorTest.java
@@ -22,6 +22,7 @@
 
 import java.time.Duration;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -31,46 +32,54 @@
  * cluster creation and cleanup because the command is expected to fail immediately.
  */
 public class LeaderElectionCommandErrorTest {
+
     @Test
     public void testTopicWithoutPartition() {
-        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.main(
-                "--bootstrap-server", "nohost:9092",
-                "--election-type", "unclean",
-                "--topic", "some-topic"
-            ));
+        String[] args = {
+            "--bootstrap-server", "nohost:9092",
+            "--election-type", "unclean",
+            "--topic", "some-topic"
+        };
+        assertEquals(1, LeaderElectionCommand.mainNoExit(args));
+        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.mainNoExit(args));
         assertTrue(out.startsWith("Missing required option(s)"));
         assertTrue(out.contains(" partition"));
     }
 
     @Test
     public void testPartitionWithoutTopic() {
-        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.main(
+        String[] args = {
             "--bootstrap-server", "nohost:9092",
             "--election-type", "unclean",
             "--all-topic-partitions",
             "--partition", "0"
-        ));
-        String[] rows = out.split("\n");
+        };
+        assertEquals(1, LeaderElectionCommand.mainNoExit(args));
+        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.mainNoExit(args));
         assertTrue(out.startsWith("Option partition is only allowed if topic is used"));
     }
 
     @Test
     public void testMissingElectionType() {
-        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.main(
+        String[] args = {
             "--bootstrap-server", "nohost:9092",
             "--topic", "some-topic",
             "--partition", "0"
-        ));
+        };
+        assertEquals(1, LeaderElectionCommand.mainNoExit(args));
+        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.mainNoExit(args));
         assertTrue(out.startsWith("Missing required option(s)"));
         assertTrue(out.contains(" election-type"));
     }
 
     @Test
     public void testMissingTopicPartitionSelection() {
-        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.main(
+        String[] args = {
             "--bootstrap-server", "nohost:9092",
             "--election-type", "preferred"
-        ));
+        };
+        assertEquals(1, LeaderElectionCommand.mainNoExit(args));
+        String out = ToolsTestUtils.captureStandardErr(() -> LeaderElectionCommand.mainNoExit(args));
         assertTrue(out.startsWith("One and only one of the following options is required: "));
         assertTrue(out.contains(" all-topic-partitions"));
         assertTrue(out.contains(" topic"));
diff --git a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
index cc99fd02118f4..210741b8501a4 100644
--- a/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/LeaderElectionCommandTest.java
@@ -105,11 +105,11 @@ public void testAllTopicPartition() throws InterruptedException, ExecutionExcept
         cluster.startBroker(broker3);
         TestUtils.waitForOnlineBroker(client, broker3);
 
-        LeaderElectionCommand.main(
+        assertEquals(0, LeaderElectionCommand.mainNoExit(
             "--bootstrap-server", cluster.bootstrapServers(),
             "--election-type", "unclean",
             "--all-topic-partitions"
-        );
+        ));
 
         TestUtils.assertLeader(client, topicPartition, broker3);
     }
@@ -121,11 +121,11 @@ public void testAdminConfigCustomTimeouts() throws Exception {
         Path adminConfigPath = tempAdminConfig(defaultApiTimeoutMs, requestTimeoutMs);
 
         try (final MockedStatic<Admin> mockedAdmin = Mockito.mockStatic(Admin.class)) {
-            LeaderElectionCommand.main(
+            assertEquals(1, LeaderElectionCommand.mainNoExit(
                 "--bootstrap-server", cluster.bootstrapServers(),
                 "--election-type", "unclean", "--all-topic-partitions",
                 "--admin.config", adminConfigPath.toString()
-            );
+            ));
 
             ArgumentCaptor<Properties> argumentCaptor = ArgumentCaptor.forClass(Properties.class);
             mockedAdmin.verify(() -> Admin.create(argumentCaptor.capture()));
@@ -161,12 +161,12 @@ public void testTopicPartition() throws InterruptedException, ExecutionException
         cluster.startBroker(broker3);
         TestUtils.waitForOnlineBroker(client, broker3);
 
-        LeaderElectionCommand.main(
+        assertEquals(0, LeaderElectionCommand.mainNoExit(
             "--bootstrap-server", cluster.bootstrapServers(),
             "--election-type", "unclean",
             "--topic", topic,
             "--partition", Integer.toString(partition)
-        );
+        ));
 
         TestUtils.assertLeader(client, topicPartition, broker3);
     }
@@ -200,11 +200,11 @@ public void testPathToJsonFile() throws Exception {
 
         Path topicPartitionPath = tempTopicPartitionFile(Collections.singletonList(topicPartition));
 
-        LeaderElectionCommand.main(
+        assertEquals(0, LeaderElectionCommand.mainNoExit(
             "--bootstrap-server", cluster.bootstrapServers(),
             "--election-type", "unclean",
             "--path-to-json-file", topicPartitionPath.toString()
-        );
+        ));
 
         TestUtils.assertLeader(client, topicPartition, broker3);
     }
@@ -233,11 +233,11 @@ public void testPreferredReplicaElection() throws InterruptedException, Executio
             JavaConverters.asScalaBuffer(Collections.singletonList(broker2)).toSet()
         );
 
-        LeaderElectionCommand.main(
+        assertEquals(0, LeaderElectionCommand.mainNoExit(
             "--bootstrap-server", cluster.bootstrapServers(),
             "--election-type", "preferred",
             "--all-topic-partitions"
-        );
+        ));
 
         TestUtils.assertLeader(client, topicPartition, broker2);
     }
@@ -288,7 +288,7 @@ public void testElectionResultOutput() throws Exception {
 
         Path topicPartitionPath = tempTopicPartitionFile(Arrays.asList(topicPartition0, topicPartition1));
         String output = ToolsTestUtils.captureStandardOut(() ->
-            LeaderElectionCommand.main(
+            LeaderElectionCommand.mainNoExit(
                 "--bootstrap-server", cluster.bootstrapServers(),
                 "--election-type", "preferred",
                 "--path-to-json-file", topicPartitionPath.toString()

From be17df6fdac5cf269d8f80b49f209fb20de70eed Mon Sep 17 00:00:00 2001
From: David Jacot <djacot@confluent.io>
Date: Mon, 25 Mar 2024 09:20:10 +0100
Subject: [PATCH 197/258] KAFKA-16374; High watermark updates should have a
 higher priority (#15534)

When the group coordinator is under heavy load, the current mechanism to release pending events based on updated high watermark, which consist in pushing an event at the end of the queue, is bad because pending events pay the cost of the queue twice. A first time for the handling of the first event and a second time for the handling of the hwm update. This patch changes this logic to push the hwm update event to the front of the queue in order to release pending events as soon as as possible.

Reviewers: Jeff Kim <jeff.kim@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../runtime/CoordinatorEventProcessor.java    |  12 +-
 .../group/runtime/CoordinatorRuntime.java     | 120 ++++++++++++------
 .../group/runtime/EventAccumulator.java       |  42 ++++--
 .../runtime/MultiThreadedEventProcessor.java  |  17 ++-
 .../group/runtime/CoordinatorRuntimeTest.java |  99 +++++++++++++--
 .../group/runtime/EventAccumulatorTest.java   |  44 ++++++-
 .../MultiThreadedEventProcessorTest.java      |  18 +--
 7 files changed, 273 insertions(+), 79 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorEventProcessor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorEventProcessor.java
index 0195880c4a5c3..f2463a7cd758d 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorEventProcessor.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorEventProcessor.java
@@ -24,10 +24,18 @@
 public interface CoordinatorEventProcessor extends AutoCloseable {
 
     /**
-     * Enqueues a new {{@link CoordinatorEvent}}.
+     * Enqueues a new {{@link CoordinatorEvent}} at the end of the processor.
      *
      * @param event The event.
      * @throws RejectedExecutionException If the event processor is closed.
      */
-    void enqueue(CoordinatorEvent event) throws RejectedExecutionException;
+    void enqueueLast(CoordinatorEvent event) throws RejectedExecutionException;
+
+    /**
+     * Enqueues a new {{@link CoordinatorEvent}} at the front of the processor.
+     *
+     * @param event The event.
+     * @throws RejectedExecutionException If the event processor is closed.
+     */
+    void enqueueFirst(CoordinatorEvent event) throws RejectedExecutionException;
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
index fd348e01ecebd..2f52f1d11f9ae 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java
@@ -51,6 +51,7 @@
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -356,7 +357,7 @@ public void run() {
 
                     log.debug("Scheduling write event {} for timer {}.", event.name, key);
                     try {
-                        enqueue(event);
+                        enqueueLast(event);
                     } catch (NotCoordinatorException ex) {
                         log.info("Failed to enqueue write event {} for timer {} because the runtime is closed. Ignoring it.",
                             event.name, key);
@@ -438,6 +439,12 @@ class CoordinatorContext {
          */
         SnapshottableCoordinator<S, U> coordinator;
 
+        /**
+         * The high watermark listener registered to all the partitions
+         * backing the coordinators.
+         */
+        HighWatermarkListener highWatermarklistener;
+
         /**
          * Constructor.
          *
@@ -495,6 +502,7 @@ private void transitionTo(
 
                 case ACTIVE:
                     state = CoordinatorState.ACTIVE;
+                    highWatermarklistener = new HighWatermarkListener();
                     partitionWriter.registerListener(tp, highWatermarklistener);
                     coordinator.onLoaded(metadataImage);
                     break;
@@ -520,7 +528,10 @@ private void transitionTo(
          * Unloads the coordinator.
          */
         private void unload() {
-            partitionWriter.deregisterListener(tp, highWatermarklistener);
+            if (highWatermarklistener != null) {
+                partitionWriter.deregisterListener(tp, highWatermarklistener);
+                highWatermarklistener = null;
+            }
             timer.cancelAll();
             deferredEventQueue.failAll(Errors.NOT_COORDINATOR.exception());
             if (coordinator != null) {
@@ -1179,6 +1190,23 @@ public String toString() {
      * backing the coordinator are updated.
      */
     class HighWatermarkListener implements PartitionWriter.Listener {
+
+        private static final long NO_OFFSET = -1L;
+
+        /**
+         * The atomic long is used to store the last and unprocessed high watermark
+         * received from the partition. The atomic value is replaced by -1L when
+         * the high watermark is taken to update the context.
+         */
+        private final AtomicLong lastHighWatermark = new AtomicLong(NO_OFFSET);
+
+        /**
+         * @return The last high watermark received or NO_OFFSET is none is pending.
+         */
+        public long lastHighWatermark() {
+            return lastHighWatermark.get();
+        }
+
         /**
          * Updates the high watermark of the corresponding coordinator.
          *
@@ -1191,30 +1219,37 @@ public void onHighWatermarkUpdated(
             long offset
         ) {
             log.debug("High watermark of {} incremented to {}.", tp, offset);
-            scheduleInternalOperation("HighWatermarkUpdated(tp=" + tp + ", offset=" + offset + ")", tp, () -> {
-                CoordinatorContext context = coordinators.get(tp);
-                if (context != null) {
-                    context.lock.lock();
-                    try {
-                        if (context.state == CoordinatorState.ACTIVE) {
-                            // The updated high watermark can be applied to the coordinator only if the coordinator
-                            // exists and is in the active state.
-                            log.debug("Updating high watermark of {} to {}.", tp, offset);
-                            context.coordinator.updateLastCommittedOffset(offset);
-                            context.deferredEventQueue.completeUpTo(offset);
-                            coordinatorMetrics.onUpdateLastCommittedOffset(tp, offset);
-                        } else {
-                            log.debug("Ignored high watermark updated for {} to {} because the coordinator is not active.",
-                                tp, offset);
+            if (lastHighWatermark.getAndSet(offset) == NO_OFFSET) {
+                // An event to apply the new high watermark is pushed to the front of the
+                // queue only if the previous value was -1L. If it was not, it means that
+                // there is already an event waiting to process the last value.
+                enqueueFirst(new CoordinatorInternalEvent("HighWatermarkUpdate", tp, () -> {
+                    long newHighWatermark = lastHighWatermark.getAndSet(NO_OFFSET);
+
+                    CoordinatorContext context = coordinators.get(tp);
+                    if (context != null) {
+                        context.lock.lock();
+                        try {
+                            if (context.state == CoordinatorState.ACTIVE) {
+                                // The updated high watermark can be applied to the coordinator only if the coordinator
+                                // exists and is in the active state.
+                                log.debug("Updating high watermark of {} to {}.", tp, newHighWatermark);
+                                context.coordinator.updateLastCommittedOffset(newHighWatermark);
+                                context.deferredEventQueue.completeUpTo(newHighWatermark);
+                                coordinatorMetrics.onUpdateLastCommittedOffset(tp, newHighWatermark);
+                            } else {
+                                log.debug("Ignored high watermark updated for {} to {} because the coordinator is not active.",
+                                    tp, newHighWatermark);
+                            }
+                        } finally {
+                            context.lock.unlock();
                         }
-                    } finally {
-                        context.lock.unlock();
+                    } else {
+                        log.debug("Ignored high watermark updated for {} to {} because the coordinator does not exist.",
+                            tp, newHighWatermark);
                     }
-                } else {
-                    log.debug("Ignored high watermark updated for {} to {} because the coordinator does not exist.",
-                        tp, offset);
-                }
-            });
+                }));
+            }
         }
     }
 
@@ -1263,12 +1298,6 @@ public void onHighWatermarkUpdated(
      */
     private final PartitionWriter<U> partitionWriter;
 
-    /**
-     * The high watermark listener registered to all the partitions
-     * backing the coordinators.
-     */
-    private final HighWatermarkListener highWatermarklistener;
-
     /**
      * The coordinator loaded used by the runtime.
      */
@@ -1335,7 +1364,6 @@ private CoordinatorRuntime(
         this.coordinators = new ConcurrentHashMap<>();
         this.processor = processor;
         this.partitionWriter = partitionWriter;
-        this.highWatermarklistener = new HighWatermarkListener();
         this.loader = loader;
         this.coordinatorShardBuilderSupplier = coordinatorShardBuilderSupplier;
         this.runtimeMetrics = runtimeMetrics;
@@ -1353,14 +1381,28 @@ private void throwIfNotRunning() {
     }
 
     /**
-     * Enqueues a new event.
+     * Enqueues a new event at the end of the processing queue.
+     *
+     * @param event The event.
+     * @throws NotCoordinatorException If the event processor is closed.
+     */
+    private void enqueueLast(CoordinatorEvent event) {
+        try {
+            processor.enqueueLast(event);
+        } catch (RejectedExecutionException ex) {
+            throw new NotCoordinatorException("Can't accept an event because the processor is closed", ex);
+        }
+    }
+
+    /**
+     * Enqueues a new event at the front of the processing queue.
      *
      * @param event The event.
      * @throws NotCoordinatorException If the event processor is closed.
      */
-    private void enqueue(CoordinatorEvent event) {
+    private void enqueueFirst(CoordinatorEvent event) {
         try {
-            processor.enqueue(event);
+            processor.enqueueFirst(event);
         } catch (RejectedExecutionException ex) {
             throw new NotCoordinatorException("Can't accept an event because the processor is closed", ex);
         }
@@ -1442,7 +1484,7 @@ public <T> CompletableFuture<T> scheduleWriteOperation(
         throwIfNotRunning();
         log.debug("Scheduled execution of write operation {}.", name);
         CoordinatorWriteEvent<T> event = new CoordinatorWriteEvent<>(name, tp, timeout, op);
-        enqueue(event);
+        enqueueLast(event);
         return event.future;
     }
 
@@ -1518,7 +1560,7 @@ public <T> CompletableFuture<T> scheduleTransactionalWriteOperation(
                 timeout,
                 op
             );
-            enqueue(event);
+            enqueueLast(event);
             return event.future;
         });
     }
@@ -1557,7 +1599,7 @@ public CompletableFuture<Void> scheduleTransactionCompletion(
             result,
             timeout
         );
-        enqueue(event);
+        enqueueLast(event);
         return event.future;
     }
 
@@ -1581,7 +1623,7 @@ public <T> CompletableFuture<T> scheduleReadOperation(
         throwIfNotRunning();
         log.debug("Scheduled execution of read operation {}.", name);
         CoordinatorReadEvent<T> event = new CoordinatorReadEvent<>(name, tp, op);
-        enqueue(event);
+        enqueueLast(event);
         return event.future;
     }
 
@@ -1622,7 +1664,7 @@ private void scheduleInternalOperation(
         Runnable op
     ) {
         log.debug("Scheduled execution of internal operation {}.", name);
-        enqueue(new CoordinatorInternalEvent(name, tp, op));
+        enqueueLast(new CoordinatorInternalEvent(name, tp, op));
     }
 
     /**
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
index 16b61f8e99150..2c22232c47a1b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/EventAccumulator.java
@@ -18,12 +18,12 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.Queue;
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.RejectedExecutionException;
@@ -60,7 +60,7 @@ public interface Event<K> {
     /**
      * The map of queues keyed by K.
      */
-    private final Map<K, Queue<T>> queues;
+    private final Map<K, Deque<T>> queues;
 
     /**
      * The list of available keys. Keys in this list can
@@ -110,17 +110,17 @@ public EventAccumulator(
     }
 
     /**
-     * Adds an {{@link Event}} to the queue.
+     * Adds an {{@link Event}} at the end of the queue.
      *
      * @param event An {{@link Event}}.
      */
-    public void add(T event) throws RejectedExecutionException {
+    public void addLast(T event) throws RejectedExecutionException {
         lock.lock();
         try {
             if (closed) throw new RejectedExecutionException("Can't accept an event because the accumulator is closed.");
 
             K key = event.key();
-            Queue<T> queue = queues.get(key);
+            Deque<T> queue = queues.get(key);
             if (queue == null) {
                 queue = new LinkedList<>();
                 queues.put(key, queue);
@@ -128,7 +128,33 @@ public void add(T event) throws RejectedExecutionException {
                     addAvailableKey(key);
                 }
             }
-            queue.add(event);
+            queue.addLast(event);
+            size++;
+        } finally {
+            lock.unlock();
+        }
+    }
+
+    /**
+     * Adds an {{@link Event}} at the front of the queue.
+     *
+     * @param event An {{@link Event}}.
+     */
+    public void addFirst(T event) throws RejectedExecutionException {
+        lock.lock();
+        try {
+            if (closed) throw new RejectedExecutionException("Can't accept an event because the accumulator is closed.");
+
+            K key = event.key();
+            Deque<T> queue = queues.get(key);
+            if (queue == null) {
+                queue = new LinkedList<>();
+                queues.put(key, queue);
+                if (!inflightKeys.contains(key)) {
+                    addAvailableKey(key);
+                }
+            }
+            queue.addFirst(event);
             size++;
         } finally {
             lock.unlock();
@@ -147,7 +173,7 @@ public T poll() {
             K key = randomKey();
             if (key == null) return null;
 
-            Queue<T> queue = queues.get(key);
+            Deque<T> queue = queues.get(key);
             T event = queue.poll();
 
             if (queue.isEmpty()) queues.remove(key);
@@ -181,7 +207,7 @@ public T take() {
 
             if (key == null) return null;
 
-            Queue<T> queue = queues.get(key);
+            Deque<T> queue = queues.get(key);
             T event = queue.poll();
 
             if (queue.isEmpty()) queues.remove(key);
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
index 0e3d563861c86..0aca6a3b79a9e 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessor.java
@@ -198,14 +198,25 @@ private void recordPollEndTime(long pollEndMs) {
     }
 
     /**
-     * Enqueues a new {{@link CoordinatorEvent}}.
+     * Enqueues a new {{@link CoordinatorEvent}} at the end of the processor.
      *
      * @param event The event.
      * @throws RejectedExecutionException If the event processor is closed.
      */
     @Override
-    public void enqueue(CoordinatorEvent event) throws RejectedExecutionException {
-        accumulator.add(event);
+    public void enqueueLast(CoordinatorEvent event) throws RejectedExecutionException {
+        accumulator.addLast(event);
+    }
+
+    /**
+     * Enqueues a new {{@link CoordinatorEvent}} at the front of the processor.
+     *
+     * @param event The event.
+     * @throws RejectedExecutionException If the event processor is closed.
+     */
+    @Override
+    public void enqueueFirst(CoordinatorEvent event) throws RejectedExecutionException {
+        accumulator.addFirst(event);
     }
 
     /**
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
index 622d4335b9a6d..c8a1dc337cba2 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntimeTest.java
@@ -47,12 +47,12 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 import java.util.OptionalInt;
-import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -96,7 +96,16 @@ public class CoordinatorRuntimeTest {
      */
     private static class DirectEventProcessor implements CoordinatorEventProcessor {
         @Override
-        public void enqueue(CoordinatorEvent event) throws RejectedExecutionException {
+        public void enqueueLast(CoordinatorEvent event) throws RejectedExecutionException {
+            try {
+                event.run();
+            } catch (Throwable ex) {
+                event.complete(ex);
+            }
+        }
+
+        @Override
+        public void enqueueFirst(CoordinatorEvent event) throws RejectedExecutionException {
             try {
                 event.run();
             } catch (Throwable ex) {
@@ -113,11 +122,16 @@ public void close() throws Exception {}
      * when poll() is called.
      */
     private static class ManualEventProcessor implements CoordinatorEventProcessor {
-        private Queue<CoordinatorEvent> queue = new LinkedList<>();
+        private Deque<CoordinatorEvent> queue = new LinkedList<>();
+
+        @Override
+        public void enqueueLast(CoordinatorEvent event) throws RejectedExecutionException {
+            queue.addLast(event);
+        }
 
         @Override
-        public void enqueue(CoordinatorEvent event) throws RejectedExecutionException {
-            queue.add(event);
+        public void enqueueFirst(CoordinatorEvent event) throws RejectedExecutionException {
+            queue.addFirst(event);
         }
 
         public boolean poll() {
@@ -507,12 +521,6 @@ public void testScheduleLoadingWithFailure() {
 
         // Verify that onUnloaded is called.
         verify(coordinator, times(1)).onUnloaded();
-
-        // Verify that the listener is deregistered.
-        verify(writer, times(1)).deregisterListener(
-            eq(TP),
-            any(PartitionWriter.Listener.class)
-        );
     }
 
     @Test
@@ -2603,6 +2611,75 @@ public void testPartitionLoadGeneratesSnapshotAtHighWatermarkNoRecordsLoaded() {
         assertTrue(ctx.coordinator.snapshotRegistry().hasSnapshot(0L));
     }
 
+    @Test
+    public void testHighWatermarkUpdate() {
+        MockTimer timer = new MockTimer();
+        MockPartitionWriter writer = new MockPartitionWriter();
+        ManualEventProcessor processor = new ManualEventProcessor();
+
+        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
+            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
+                .withTime(timer.time())
+                .withTimer(timer)
+                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
+                .withLoader(new MockCoordinatorLoader())
+                .withEventProcessor(processor)
+                .withPartitionWriter(writer)
+                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
+                .withCoordinatorRuntimeMetrics(mock(GroupCoordinatorRuntimeMetrics.class))
+                .withCoordinatorMetrics(mock(GroupCoordinatorMetrics.class))
+                .build();
+
+        // Loads the coordinator. Poll once to execute the load operation and once
+        // to complete the load.
+        runtime.scheduleLoadOperation(TP, 10);
+        processor.poll();
+        processor.poll();
+
+        // Write #1.
+        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, DEFAULT_WRITE_TIMEOUT,
+            state -> new CoordinatorResult<>(Collections.singletonList("record1"), "response1")
+        );
+        processor.poll();
+
+        // Write #2.
+        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, DEFAULT_WRITE_TIMEOUT,
+            state -> new CoordinatorResult<>(Collections.singletonList("record2"), "response2")
+        );
+        processor.poll();
+
+        // Records have been written to the log.
+        assertEquals(Arrays.asList(
+            InMemoryPartitionWriter.LogEntry.value("record1"),
+            InMemoryPartitionWriter.LogEntry.value("record2")
+        ), writer.entries(TP));
+
+        // There is no pending high watermark.
+        assertEquals(-1, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());
+
+        // Commit the first record.
+        writer.commit(TP, 1);
+
+        // We should have one pending event and the pending high watermark should be set.
+        assertEquals(1, processor.size());
+        assertEquals(1, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());
+
+        // Commit the second record.
+        writer.commit(TP, 2);
+
+        // We should still have one pending event and the pending high watermark should be updated.
+        assertEquals(1, processor.size());
+        assertEquals(2, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());
+
+        // Poll once to process the high watermark update and complete the writes.
+        processor.poll();
+
+        assertEquals(-1, runtime.contextOrThrow(TP).highWatermarklistener.lastHighWatermark());
+        assertEquals(2, runtime.contextOrThrow(TP).coordinator.lastCommittedOffset());
+        assertTrue(write1.isDone());
+        assertTrue(write2.isDone());
+    }
+
     private static <S extends CoordinatorShard<U>, U> ArgumentMatcher<CoordinatorPlayback<U>> coordinatorMatcher(
         CoordinatorRuntime<S, U> runtime,
         TopicPartition tp
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
index e077fb5e0226c..faac0f46f6a48 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/EventAccumulatorTest.java
@@ -16,7 +16,9 @@
  */
 package org.apache.kafka.coordinator.group.runtime;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -92,7 +94,7 @@ public void testBasicOperations() {
             new MockEvent(3, 2)
         );
 
-        events.forEach(accumulator::add);
+        events.forEach(accumulator::addLast);
         assertEquals(9, accumulator.size());
 
         Set<MockEvent> polledEvents = new HashSet<>();
@@ -111,6 +113,34 @@ public void testBasicOperations() {
         accumulator.close();
     }
 
+    @Test
+    public void testAddFirst() {
+        EventAccumulator<Integer, MockEvent> accumulator = new EventAccumulator<>();
+
+        List<MockEvent> events = Arrays.asList(
+            new MockEvent(1, 0),
+            new MockEvent(1, 1),
+            new MockEvent(1, 2)
+        );
+
+        events.forEach(accumulator::addFirst);
+        assertEquals(3, accumulator.size());
+
+        List<MockEvent> polledEvents = new ArrayList<>(3);
+        for (int i = 0; i < events.size(); i++) {
+            MockEvent event = accumulator.poll();
+            assertNotNull(event);
+            polledEvents.add(event);
+            assertEquals(events.size() - 1 - i, accumulator.size());
+            accumulator.done(event);
+        }
+
+        Collections.reverse(events);
+        assertEquals(events, polledEvents);
+
+        accumulator.close();
+    }
+
     @Test
     public void testKeyConcurrentAndOrderingGuarantees() {
         EventAccumulator<Integer, MockEvent> accumulator = new EventAccumulator<>();
@@ -118,9 +148,9 @@ public void testKeyConcurrentAndOrderingGuarantees() {
         MockEvent event0 = new MockEvent(1, 0);
         MockEvent event1 = new MockEvent(1, 1);
         MockEvent event2 = new MockEvent(1, 2);
-        accumulator.add(event0);
-        accumulator.add(event1);
-        accumulator.add(event2);
+        accumulator.addLast(event0);
+        accumulator.addLast(event1);
+        accumulator.addLast(event2);
         assertEquals(3, accumulator.size());
 
         MockEvent event = null;
@@ -169,9 +199,9 @@ public void testDoneUnblockWaitingThreads() throws ExecutionException, Interrupt
         assertFalse(future1.isDone());
         assertFalse(future2.isDone());
 
-        accumulator.add(event0);
-        accumulator.add(event1);
-        accumulator.add(event2);
+        accumulator.addLast(event0);
+        accumulator.addLast(event1);
+        accumulator.addLast(event2);
 
         // One future should be completed with event0.
         assertEquals(event0, CompletableFuture
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
index 3708141827cc0..f01fc88310589 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/runtime/MultiThreadedEventProcessorTest.java
@@ -186,7 +186,7 @@ public void testEventsAreProcessed() throws Exception {
                 new FutureEvent<>(new TopicPartition("foo", 2), numEventsExecuted::incrementAndGet)
             );
 
-            events.forEach(eventProcessor::enqueue);
+            events.forEach(eventProcessor::enqueueLast);
 
             CompletableFuture.allOf(events
                 .stream()
@@ -223,7 +223,7 @@ public void testProcessingGuarantees() throws Exception {
                 new FutureEvent<>(new TopicPartition("foo", 1), numEventsExecuted::incrementAndGet, true)  // Event 5
             );
 
-            events.forEach(eventProcessor::enqueue);
+            events.forEach(eventProcessor::enqueueLast);
 
             // Events 0 and 1 are executed.
             assertTrue(events.get(0).awaitExecution(5, TimeUnit.SECONDS));
@@ -301,7 +301,7 @@ public void testEventsAreRejectedWhenClosed() throws Exception {
         eventProcessor.close();
 
         assertThrows(RejectedExecutionException.class,
-            () -> eventProcessor.enqueue(new FutureEvent<>(new TopicPartition("foo", 0), () -> 0)));
+            () -> eventProcessor.enqueueLast(new FutureEvent<>(new TopicPartition("foo", 0), () -> 0)));
     }
 
     @Test
@@ -332,14 +332,14 @@ public void testEventsAreDrainedWhenClosed() throws Exception {
             );
 
             // Enqueue the blocking event.
-            eventProcessor.enqueue(blockingEvent);
+            eventProcessor.enqueueLast(blockingEvent);
 
             // Ensure that the blocking event is executed.
             waitForCondition(() -> numEventsExecuted.get() > 0,
                 "Blocking event not executed.");
 
             // Enqueue the other events.
-            events.forEach(eventProcessor::enqueue);
+            events.forEach(eventProcessor::enqueueLast);
 
             // Events should not be completed.
             events.forEach(event -> assertFalse(event.future.isDone()));
@@ -349,7 +349,7 @@ public void testEventsAreDrainedWhenClosed() throws Exception {
 
             // Enqueuing a new event is rejected.
             assertThrows(RejectedExecutionException.class,
-                () -> eventProcessor.enqueue(blockingEvent));
+                () -> eventProcessor.enqueueLast(blockingEvent));
 
             // Release the blocking event to unblock the thread.
             blockingEvent.release();
@@ -398,7 +398,7 @@ public void testMetrics() throws Exception {
             new DelayEventAccumulator(mockTime, 500L)
         )) {
             // Enqueue the blocking event.
-            eventProcessor.enqueue(blockingEvent);
+            eventProcessor.enqueueLast(blockingEvent);
 
             // Ensure that the blocking event is executed.
             waitForCondition(() -> numEventsExecuted.get() > 0,
@@ -414,7 +414,7 @@ public void testMetrics() throws Exception {
                 mockTime.milliseconds()
             );
 
-            eventProcessor.enqueue(otherEvent);
+            eventProcessor.enqueueLast(otherEvent);
 
             // Pass the time.
             mockTime.sleep(3000L);
@@ -492,7 +492,7 @@ public void testRecordThreadIdleRatioTwoThreads() throws Exception {
                 new FutureEvent<>(new TopicPartition("foo", 2), numEventsExecuted::incrementAndGet)
             );
 
-            events.forEach(eventProcessor::enqueue);
+            events.forEach(eventProcessor::enqueueLast);
 
             CompletableFuture.allOf(events
                 .stream()

From e9e007aec8004576ac6d34de690f17a3f5fa8af6 Mon Sep 17 00:00:00 2001
From: Vedarth Sharma <142404391+VedarthConfluent@users.noreply.github.com>
Date: Mon, 25 Mar 2024 13:53:01 +0530
Subject: [PATCH 198/258] KAFKA-15882: Add nightly docker image scan job
 (#15013)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 .github/workflows/docker_scan.yml | 44 +++++++++++++++++++++++++++++++
 docker/README.md                  | 20 ++++++++++++++
 2 files changed, 64 insertions(+)
 create mode 100644 .github/workflows/docker_scan.yml

diff --git a/.github/workflows/docker_scan.yml b/.github/workflows/docker_scan.yml
new file mode 100644
index 0000000000000..7d9ecfe6192b5
--- /dev/null
+++ b/.github/workflows/docker_scan.yml
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Docker Image CVE Scanner
+on:
+  schedule:
+    # This job will run at 3:30 UTC daily
+    - cron: '30 3 * * *'
+  workflow_dispatch:
+jobs:
+  scan_jvm:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        # This is an array of supported tags. Make sure this array only contains the supported tags
+        supported_image_tag: ['latest', '3.7.0']
+    steps:
+      - name: Run CVE scan
+        uses: aquasecurity/trivy-action@master
+        if: always()
+        with:
+          image-ref: apache/kafka:${{ matrix.supported_image_tag }}
+          format: 'table'
+          severity: 'CRITICAL,HIGH'
+          output: scan_report_jvm_${{ matrix.supported_image_tag }}.txt
+          exit-code: '1'
+      - name: Upload CVE scan report
+        if: always()
+        uses: actions/upload-artifact@v3
+        with:
+          name: scan_report_jvm_${{ matrix.supported_image_tag }}.txt
+          path: scan_report_jvm_${{ matrix.supported_image_tag }}.txt
diff --git a/docker/README.md b/docker/README.md
index 0232604e61325..54d15c04feb17 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -63,6 +63,26 @@ rc_docker_image: apache/kafka:3.6.0-rc0
 promoted_docker_image: apache/kafka:3.6.0
 ```
 
+Cron job for checking CVEs in supported docker images
+-----------------------------------------------------
+
+- `Docker Image CVE Scanner` Github Action Workflow (present in `.github/workflows/docker_scan.yml`) will run nightly CVE scans and generate reports for docker image tags mentioned in the `supported_image_tag` array.
+- This workflow is branch independent. Only the workflow in trunk, i.e. the default branch will be considered.
+- In case a Critical or High CVE is detected, the workflow will fail.
+- It will generate the scan reports that can be checked by the community.
+- For every new release, this should be updated with the latest supported docker images.
+- For example:-
+```
+For supporting apache/kafka:3.6.0, apache/kafka:latest and apache/kafka:3.7.0-rc0, supported_image_tag array should be
+supported_image_tag: ['3.6.0', 'latest', '3.7.0-rc0']
+```
+- When RC for a version gets changed or when a bug fix release happens, this should be updated as well.
+- For example:-
+```
+For supporting apache/kafka:3.6.1, apache/kafka:latest and apache/kafka:3.7.0-rc1, tag array should be
+supported_image_tag: ['3.6.1', 'latest', '3.7.0-rc1']
+```
+
 Local Setup
 -----------
 

From 0104cd04316a74127be2502a01b5b356a2a06087 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Mon, 25 Mar 2024 11:30:08 +0100
Subject: [PATCH 199/258] KAFKA-16375: Fix for rejoin while reconciling
 (#15579)

This PR includes a fix to properly identify a reconciliation that should be interrupted and not applied because the member has rejoined. It does so simply based on a flag (not epochs, server or local). If the member has rejoined while reconciling, the reconciliation will be interrupted.

This also ensures that the check to abort the reconciliation is performed on all the 3 stages of the reconciliation that could be delayed: commit, onPartitionsRevoked, onPartitionsAssigned.

Reviewers: David Jacot <djacot@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../internals/MembershipManagerImpl.java      |  75 ++++---
 .../internals/MembershipManagerImplTest.java  | 185 +++++++++++++++---
 2 files changed, 199 insertions(+), 61 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 8dc033ea5fcf9..574db9ebcd0fe 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -228,11 +228,11 @@ public class MembershipManagerImpl implements MembershipManager {
     private boolean reconciliationInProgress;
 
     /**
-     * Epoch the member had when the reconciliation in progress started. This is used to identify if
-     * the member has rejoined while it was reconciling an assignment (in which case the result
-     * of the reconciliation is not applied.)
+     * True if a reconciliation is in progress and the member rejoined the group since the start
+     * of the reconciliation. Used to know that the reconciliation in progress should be
+     * interrupted and not be applied.
      */
-    private int memberEpochOnReconciliationStart;
+    private boolean rejoinedWhileReconciliationInProgress;
 
     /**
      * If the member is currently leaving the group after a call to {@link #leaveGroup()}}, this
@@ -641,6 +641,9 @@ public void transitionToJoining() {
                     "the member is in FATAL state");
             return;
         }
+        if (reconciliationInProgress) {
+            rejoinedWhileReconciliationInProgress = true;
+        }
         resetEpoch();
         transitionTo(MemberState.JOINING);
         clearPendingAssignmentsAndLocalNamesCache();
@@ -972,7 +975,10 @@ void maybeReconcile() {
                 log.debug("Auto-commit before reconciling new assignment completed successfully.");
             }
 
-            revokeAndAssign(resolvedAssignment, assignedTopicIdPartitions, revokedPartitions, addedPartitions);
+            if (!maybeAbortReconciliation()) {
+                revokeAndAssign(resolvedAssignment, assignedTopicIdPartitions, revokedPartitions, addedPartitions);
+            }
+
         }).exceptionally(error -> {
             if (error != null) {
                 log.error("Reconciliation failed.", error);
@@ -1010,33 +1016,23 @@ private void revokeAndAssign(LocalAssignment resolvedAssignment,
         // and assignment, executed sequentially).
         CompletableFuture<Void> reconciliationResult =
             revocationResult.thenCompose(__ -> {
-                boolean memberHasRejoined = memberEpochOnReconciliationStart != memberEpoch;
-                if (state == MemberState.RECONCILING && !memberHasRejoined) {
+                if (!maybeAbortReconciliation()) {
                     // Apply assignment
                     return assignPartitions(assignedTopicIdPartitions, addedPartitions);
-                } else {
-                    log.debug("Revocation callback completed but the member already " +
-                        "transitioned out of the reconciling state for epoch {} into " +
-                        "{} state with epoch {}. Interrupting reconciliation as it's " +
-                        "not relevant anymore,", memberEpochOnReconciliationStart, state, memberEpoch);
-                    String reason = interruptedReconciliationErrorMessage();
-                    CompletableFuture<Void> res = new CompletableFuture<>();
-                    res.completeExceptionally(new KafkaException("Interrupting reconciliation" +
-                        " after revocation. " + reason));
-                    return res;
                 }
+                return CompletableFuture.completedFuture(null);
             });
 
-        reconciliationResult.whenComplete((result, error) -> {
-            markReconciliationCompleted();
+        reconciliationResult.whenComplete((__, error) -> {
             if (error != null) {
                 // Leaving member in RECONCILING state after callbacks fail. The member
                 // won't send the ack, and the expectation is that the broker will kick the
                 // member out of the group after the rebalance timeout expires, leading to a
                 // RECONCILING -> FENCED transition.
                 log.error("Reconciliation failed.", error);
+                markReconciliationCompleted();
             } else {
-                if (state == MemberState.RECONCILING) {
+                if (reconciliationInProgress && !maybeAbortReconciliation()) {
                     currentAssignment = resolvedAssignment;
 
                     // Reschedule the auto commit starting from now that the member has a new assignment.
@@ -1044,15 +1040,30 @@ private void revokeAndAssign(LocalAssignment resolvedAssignment,
 
                     // Make assignment effective on the broker by transitioning to send acknowledge.
                     transitionTo(MemberState.ACKNOWLEDGING);
-                } else {
-                    String reason = interruptedReconciliationErrorMessage();
-                    log.error("Interrupting reconciliation after partitions assigned callback " +
-                        "completed. " + reason);
+                    markReconciliationCompleted();
                 }
             }
         });
     }
 
+    /**
+     * @return True if the reconciliation in progress should not continue. This could be because
+     * the member is not in RECONCILING state anymore (member failed or is leaving the group), or
+     * if it has rejoined the group (note that after rejoining the member could be RECONCILING
+     * again, so checking the state is not enough)
+     */
+    boolean maybeAbortReconciliation() {
+        boolean shouldAbort = state != MemberState.RECONCILING || rejoinedWhileReconciliationInProgress;
+        if (shouldAbort) {
+            String reason = rejoinedWhileReconciliationInProgress ?
+                "the member has re-joined the group" :
+                "the member already transitioned out of the reconciling state into " + state;
+            log.info("Interrupting reconciliation that is not relevant anymore because " + reason);
+            markReconciliationCompleted();
+        }
+        return shouldAbort;
+    }
+
     // Visible for testing.
     void updateAssignment(Map<Uuid, SortedSet<Integer>> partitions) {
         currentAssignment = new LocalAssignment(0, partitions);
@@ -1067,25 +1078,12 @@ private SortedSet<TopicPartition> toTopicPartitionSet(SortedSet<TopicIdPartition
         return result;
     }
 
-    /**
-     * @return Reason for interrupting a reconciliation progress when callbacks complete.
-     */
-    private String interruptedReconciliationErrorMessage() {
-        String reason;
-        if (state != MemberState.RECONCILING) {
-            reason = "The member already transitioned out of the reconciling state into " + state;
-        } else {
-            reason = "The member has re-joined the group.";
-        }
-        return reason;
-    }
-
     /**
      *  Visible for testing.
      */
     void markReconciliationInProgress() {
         reconciliationInProgress = true;
-        memberEpochOnReconciliationStart = memberEpoch;
+        rejoinedWhileReconciliationInProgress = false;
     }
 
     /**
@@ -1093,6 +1091,7 @@ void markReconciliationInProgress() {
      */
     void markReconciliationCompleted() {
         reconciliationInProgress = false;
+        rejoinedWhileReconciliationInProgress = false;
     }
 
     /**
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index 4f893a4efbc11..b2c75f495dd7e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -470,13 +470,13 @@ public void testDelayedReconciliationResultDiscardedIfMemberNotInReconcilingStat
     }
 
     /**
-     * This is the case where a member is stuck reconciling an assignment A (waiting on
-     * metadata, commit or callbacks), and it rejoins (due to fence or unsubscribe/subscribe). If
-     * the reconciliation of A completes it should not be applied (it should not update the
-     * assignment on the member or send ack).
+     * This is the case where a member is stuck reconciling an assignment A (waiting for commit
+     * to complete), and it rejoins (due to fence or unsubscribe/subscribe). If the
+     * reconciliation of A completes it should be interrupted, and it should not update the
+     * assignment on the member or send ack.
      */
     @Test
-    public void testDelayedReconciliationResultDiscardedIfMemberRejoins() {
+    public void testDelayedReconciliationResultDiscardedAfterCommitIfMemberRejoins() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         Uuid topicId1 = Uuid.randomUuid();
         String topic1 = "topic1";
@@ -496,25 +496,83 @@ public void testDelayedReconciliationResultDiscardedIfMemberRejoins() {
         testFencedMemberReleasesAssignmentAndTransitionsToJoining(membershipManager);
         clearInvocations(subscriptionState);
 
-        // Get new assignment A2 after rejoining. This should not trigger a reconciliation just
-        // yet because there is another on in progress, but should keep the new assignment ready
-        // to be reconciled next.
-        Uuid topicId3 = Uuid.randomUuid();
-        mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId3, "topic3", owned);
-        receiveAssignmentAfterRejoin(topicId3, Collections.singletonList(5), membershipManager);
-        verifyReconciliationNotTriggered(membershipManager);
-        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin = topicIdPartitionsMap(topicId3, 5);
-        assertEquals(assignmentAfterRejoin, membershipManager.topicPartitionsAwaitingReconciliation());
+        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin = receiveAssignmentAfterRejoin(
+            Collections.singletonList(5), membershipManager, owned);
 
         // Reconciliation completes when the member has already re-joined the group. Should not
-        // update the subscription state or send ack.
+        // proceed with the revocation, update the subscription state or send ack.
         commitResult.complete(null);
-        verify(subscriptionState, never()).assignFromSubscribed(anyCollection());
-        assertNotEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+        assertInitialReconciliationDiscardedAfterRejoin(membershipManager, assignmentAfterRejoin);
+    }
 
-        // Assignment received after rejoining should be ready to reconcile on the next
-        // reconciliation loop.
-        assertEquals(assignmentAfterRejoin, membershipManager.topicPartitionsAwaitingReconciliation());
+    /**
+     * This is the case where a member is stuck reconciling an assignment A (waiting for
+     * onPartitionsRevoked callback to complete), and it rejoins (due to fence or
+     * unsubscribe/subscribe). When the reconciliation of A completes it should be interrupted,
+     * and it should not update the assignment on the member or send ack.
+     */
+    @Test
+    public void testDelayedReconciliationResultDiscardedAfterPartitionsRevokedCallbackIfMemberRejoins() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        Uuid topicId1 = Uuid.randomUuid();
+        String topic1 = "topic1";
+        List<TopicIdPartition> owned = Collections.singletonList(new TopicIdPartition(topicId1,
+            new TopicPartition(topic1, 0)));
+        mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId1, topic1, owned);
+
+        // Reconciliation that does not complete stuck on onPartitionsRevoked callback
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+        ConsumerRebalanceListenerCallbackCompletedEvent callbackCompletedEvent =
+            mockNewAssignmentStuckOnPartitionsRevokedCallback(membershipManager, topicId1, topic1,
+            Arrays.asList(1, 2), owned.get(0).topicPartition(), invoker);
+        Map<Uuid, SortedSet<Integer>> assignment1 = topicIdPartitionsMap(topicId1,  1, 2);
+        assertEquals(assignment1, membershipManager.topicPartitionsAwaitingReconciliation());
+
+        // Get fenced and rejoin while still reconciling. Get new assignment to reconcile after rejoining.
+        testFencedMemberReleasesAssignmentAndTransitionsToJoining(membershipManager);
+        clearInvocations(subscriptionState);
+
+        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin = receiveAssignmentAfterRejoin(
+            Collections.singletonList(5), membershipManager, owned);
+
+        // onPartitionsRevoked callback completes when the member has already re-joined the group.
+        // Should not proceed with the assignment, update the subscription state or send ack.
+        completeCallback(callbackCompletedEvent, membershipManager);
+        assertInitialReconciliationDiscardedAfterRejoin(membershipManager, assignmentAfterRejoin);
+    }
+
+    /**
+     * This is the case where a member is stuck reconciling an assignment A (waiting for
+     * onPartitionsAssigned callback to complete), and it rejoins (due to fence or
+     * unsubscribe/subscribe). If the reconciliation of A completes it should be interrupted, and it
+     * should not update the assignment on the member or send ack.
+     */
+    @Test
+    public void testDelayedReconciliationResultDiscardedAfterPartitionsAssignedCallbackIfMemberRejoins() {
+        MembershipManagerImpl membershipManager = createMemberInStableState();
+        Uuid topicId1 = Uuid.randomUuid();
+        String topic1 = "topic1";
+
+        // Reconciliation that does not complete stuck on onPartitionsAssigned callback
+        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
+        int newPartition = 1;
+        ConsumerRebalanceListenerCallbackCompletedEvent callbackCompletedEvent =
+            mockNewAssignmentStuckOnPartitionsAssignedCallback(membershipManager, topicId1,
+                topic1, newPartition, invoker);
+        Map<Uuid, SortedSet<Integer>> assignment1 = topicIdPartitionsMap(topicId1,  newPartition);
+        assertEquals(assignment1, membershipManager.topicPartitionsAwaitingReconciliation());
+
+        // Get fenced and rejoin while still reconciling. Get new assignment to reconcile after rejoining.
+        testFencedMemberReleasesAssignmentAndTransitionsToJoining(membershipManager);
+        clearInvocations(subscriptionState);
+
+        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin = receiveAssignmentAfterRejoin(
+            Collections.singletonList(5), membershipManager, Collections.emptyList());
+
+        // onPartitionsAssigned callback completes when the member has already re-joined the group.
+        // Should not update the subscription state or send ack.
+        completeCallback(callbackCompletedEvent, membershipManager);
+        assertInitialReconciliationDiscardedAfterRejoin(membershipManager, assignmentAfterRejoin);
     }
 
     /**
@@ -2147,7 +2205,7 @@ public void testRebalanceMetricsOnSuccessfulRebalance() {
     }
 
     @Test
-    public void testRebalanceMetricsForMultipleReconcilations() {
+    public void testRebalanceMetricsForMultipleReconciliations() {
         MembershipManagerImpl membershipManager = createMemberInStableState();
         ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
 
@@ -2299,6 +2357,57 @@ private CompletableFuture<Void> mockNewAssignmentAndRevocationStuckOnCommit(
         return commitResult;
     }
 
+    private ConsumerRebalanceListenerCallbackCompletedEvent mockNewAssignmentStuckOnPartitionsRevokedCallback(
+        MembershipManagerImpl membershipManager, Uuid topicId, String topicName,
+        List<Integer> partitions, TopicPartition ownedPartition, ConsumerRebalanceListenerInvoker invoker) {
+        doNothing().when(subscriptionState).markPendingRevocation(anySet());
+        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
+        when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(ownedPartition));
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+        when(commitRequestManager.autoCommitEnabled()).thenReturn(false);
+
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
+        receiveAssignment(topicId, partitions, membershipManager);
+        membershipManager.poll(time.milliseconds());
+        verifyReconciliationTriggered(membershipManager);
+        clearInvocations(membershipManager);
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+
+        return performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED,
+            topicPartitions(ownedPartition.topic(), ownedPartition.partition()),
+            false
+        );
+    }
+
+    private ConsumerRebalanceListenerCallbackCompletedEvent mockNewAssignmentStuckOnPartitionsAssignedCallback(
+        MembershipManagerImpl membershipManager, Uuid topicId, String topicName, int newPartition,
+        ConsumerRebalanceListenerInvoker invoker) {
+        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
+        when(subscriptionState.assignedPartitions()).thenReturn(Collections.emptySet());
+        when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
+        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
+        when(commitRequestManager.autoCommitEnabled()).thenReturn(false);
+
+        when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
+        receiveAssignment(topicId, Collections.singletonList(newPartition), membershipManager);
+        membershipManager.poll(time.milliseconds());
+        verifyReconciliationTriggered(membershipManager);
+        clearInvocations(membershipManager);
+        assertEquals(MemberState.RECONCILING, membershipManager.state());
+
+        return performCallback(
+            membershipManager,
+            invoker,
+            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
+            topicPartitions(topicName, newPartition),
+            false
+        );
+    }
+
     private void verifyReconciliationTriggered(MembershipManagerImpl membershipManager) {
         verify(membershipManager).markReconciliationInProgress();
         assertEquals(MemberState.RECONCILING, membershipManager.state());
@@ -2480,7 +2589,15 @@ private void receiveAssignment(Uuid topicId, List<Integer> partitions, Membershi
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
     }
 
-    private void receiveAssignmentAfterRejoin(Uuid topicId, List<Integer> partitions, MembershipManager membershipManager) {
+    private Map<Uuid, SortedSet<Integer>> receiveAssignmentAfterRejoin(List<Integer> partitions,
+                                                                       MembershipManagerImpl membershipManager,
+                                                                       Collection<TopicIdPartition> owned) {
+        // Get new assignment after rejoining. This should not trigger a reconciliation just
+        // yet because there is another one in progress, but should keep the new assignment ready
+        // to be reconciled next.
+        Uuid topicId = Uuid.randomUuid();
+        mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId, "topic3", owned);
+
         ConsumerGroupHeartbeatResponseData.Assignment targetAssignment = new ConsumerGroupHeartbeatResponseData.Assignment()
                 .setTopicPartitions(Collections.singletonList(
                         new ConsumerGroupHeartbeatResponseData.TopicPartitions()
@@ -2489,6 +2606,29 @@ private void receiveAssignmentAfterRejoin(Uuid topicId, List<Integer> partitions
         ConsumerGroupHeartbeatResponse heartbeatResponse =
                 createConsumerGroupHeartbeatResponseWithBumpedEpoch(targetAssignment);
         membershipManager.onHeartbeatSuccess(heartbeatResponse.data());
+
+        verifyReconciliationNotTriggered(membershipManager);
+        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin = topicIdPartitionsMap(topicId, 5);
+        assertEquals(assignmentAfterRejoin, membershipManager.topicPartitionsAwaitingReconciliation());
+        return assignmentAfterRejoin;
+    }
+
+    private void assertInitialReconciliationDiscardedAfterRejoin(
+        MembershipManagerImpl membershipManager,
+        Map<Uuid, SortedSet<Integer>> assignmentAfterRejoin) {
+        verify(subscriptionState, never()).markPendingRevocation(any());
+        verify(subscriptionState, never()).assignFromSubscribed(anyCollection());
+        assertNotEquals(MemberState.ACKNOWLEDGING, membershipManager.state());
+
+        // Assignment received after rejoining should be ready to reconcile on the next
+        // reconciliation loop.
+        assertEquals(assignmentAfterRejoin, membershipManager.topicPartitionsAwaitingReconciliation());
+
+        // Stale reconciliation should have been aborted and a new one should be triggered on the next poll.
+        assertFalse(membershipManager.reconciliationInProgress());
+        clearInvocations(membershipManager);
+        membershipManager.poll(time.milliseconds());
+        verify(membershipManager).markReconciliationInProgress();
     }
 
     private void receiveEmptyAssignment(MembershipManager membershipManager) {
@@ -2578,7 +2718,6 @@ private ConsumerRebalanceListenerCallbackCompletedEvent mockFencedMemberStuckOnU
         when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(ownedPartition));
         when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
         when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));
-        // doNothing().when(subscriptionState).markPendingRevocation(anySet());
         when(commitRequestManager.autoCommitEnabled()).thenReturn(false);
         membershipManager.transitionToFenced();
         return performCallback(

From 51c9b0d0ad408754b1c5883a9c7fcc63a5f57eb8 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Mon, 25 Mar 2024 15:34:52 +0100
Subject: [PATCH 200/258] KAFKA-16406: Splitting consumer integration test
 (#15535)

Splitting consumer integration tests to allow for parallelization and reduce build times. This PR is only extracting tests from PlainTextConsumerTest into separate files, no changes in logic. Grouping tests by the feature they relate to so that they can be easily found

Reviewers: Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>
---
 .../kafka/api/AbstractConsumerTest.scala      |   74 +
 .../api/PlaintextConsumerAssignTest.scala     |  246 ++++
 .../api/PlaintextConsumerAssignorsTest.scala  |  403 +++++
 .../api/PlaintextConsumerFetchTest.scala      |  243 ++++
 .../kafka/api/PlaintextConsumerPollTest.scala |  284 ++++
 .../PlaintextConsumerSubscriptionTest.scala   |  237 +++
 .../kafka/api/PlaintextConsumerTest.scala     | 1290 +----------------
 7 files changed, 1489 insertions(+), 1288 deletions(-)
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignTest.scala
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerFetchTest.scala
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerPollTest.scala
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerSubscriptionTest.scala

diff --git a/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
index ee9d8426eeb96..78d51329b7322 100644
--- a/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
@@ -82,6 +82,80 @@ abstract class AbstractConsumerTest extends BaseRequestTest {
     createTopic(topic, 2, brokerCount, adminClientConfig = this.adminClientConfig)
   }
 
+  def awaitAssignment(consumer: Consumer[_, _], expectedAssignment: Set[TopicPartition])
+  : Unit = {
+    TestUtils.pollUntilTrue(consumer, () => consumer.assignment() == expectedAssignment.asJava,
+      s"Timed out while awaiting expected assignment $expectedAssignment. " +
+        s"The current assignment is ${consumer.assignment()}")
+  }
+
+  def awaitNonEmptyRecords[K, V](consumer: Consumer[K, V], partition: TopicPartition): ConsumerRecords[K, V] = {
+    TestUtils.pollRecordsUntilTrue(consumer, (polledRecords: ConsumerRecords[K, V]) => {
+      if (polledRecords.records(partition).asScala.nonEmpty)
+        return polledRecords
+      false
+    }, s"Consumer did not consume any messages for partition $partition before timeout.")
+    throw new IllegalStateException("Should have timed out before reaching here")
+  }
+
+  /**
+   * Creates 'consumerCount' consumers and consumer pollers, one per consumer; subscribes consumers to
+   * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
+   *
+   * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
+   *
+   * @param consumerCount     number of consumers to create
+   * @param topicsToSubscribe topics to which consumers will subscribe to
+   * @param subscriptions     set of all topic partitions
+   * @return collection of created consumers and collection of corresponding consumer pollers
+   */
+  def createConsumerGroupAndWaitForAssignment(consumerCount: Int,
+                                              topicsToSubscribe: List[String],
+                                              subscriptions: Set[TopicPartition]): (Buffer[Consumer[Array[Byte], Array[Byte]]], Buffer[ConsumerAssignmentPoller]) = {
+    assertTrue(consumerCount <= subscriptions.size)
+    val consumerGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
+    for (_ <- 0 until consumerCount)
+      consumerGroup += createConsumer()
+
+    // create consumer pollers, wait for assignment and validate it
+    val consumerPollers = subscribeConsumers(consumerGroup, topicsToSubscribe)
+    (consumerGroup, consumerPollers)
+  }
+
+  /**
+   * Creates consumer pollers corresponding to a given consumer group, one per consumer; subscribes consumers to
+   * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
+   *
+   * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
+   *
+   * @param consumerGroup     consumer group
+   * @param topicsToSubscribe topics to which consumers will subscribe to
+   * @return collection of consumer pollers
+   */
+  def subscribeConsumers(consumerGroup: mutable.Buffer[Consumer[Array[Byte], Array[Byte]]],
+                         topicsToSubscribe: List[String]): mutable.Buffer[ConsumerAssignmentPoller] = {
+    val consumerPollers = mutable.Buffer[ConsumerAssignmentPoller]()
+    for (consumer <- consumerGroup)
+      consumerPollers += subscribeConsumerAndStartPolling(consumer, topicsToSubscribe)
+    consumerPollers
+  }
+
+  def changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers: Buffer[ConsumerAssignmentPoller],
+                                                           topicsToSubscribe: List[String],
+                                                           subscriptions: Set[TopicPartition]): Unit = {
+    for (poller <- consumerPollers)
+      poller.subscribe(topicsToSubscribe)
+
+    // since subscribe call to poller does not actually call consumer subscribe right away, wait
+    // until subscribe is called on all consumers
+    TestUtils.waitUntilTrue(() => {
+      consumerPollers.forall { poller => poller.isSubscribeRequestProcessed }
+    }, s"Failed to call subscribe on all consumers in the group for subscription $subscriptions", 1000L)
+
+    validateGroupAssignment(consumerPollers, subscriptions,
+      Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after we changed subscription"))
+  }
+
   protected class TestConsumerReassignmentListener extends ConsumerRebalanceListener {
     var callsToAssigned = 0
     var callsToRevoked = 0
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignTest.scala
new file mode 100644
index 0000000000000..219ed9c2a1e6f
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignTest.scala
@@ -0,0 +1,246 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.{TestInfoUtils, TestUtils}
+import java.util.Properties
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.common.TopicPartition
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, MethodSource}
+
+import org.apache.kafka.common.PartitionInfo
+import java.util.stream.Stream
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable
+import org.junit.jupiter.params.provider.CsvSource
+
+/**
+ * Integration tests for the consumer that covers logic related to manual assignment.
+ */
+@Timeout(600)
+class PlaintextConsumerAssignTest extends AbstractConsumerTest {
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndCommitAsyncNotCommitted(quorum: String, groupProtocol: String): Unit = {
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props)
+    val producer = createProducer()
+    val numRecords = 10000
+    val startingTimestamp = System.currentTimeMillis()
+    val cb = new CountConsumerCommitCallback
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+    consumer.assign(List(tp).asJava)
+    consumer.commitAsync(cb)
+    TestUtils.pollUntilTrue(consumer, () => cb.successCount >= 1 || cb.lastError.isDefined,
+      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
+    val committedOffset = consumer.committed(Set(tp).asJava)
+    assertNotNull(committedOffset)
+    // No valid fetch position due to the absence of consumer.poll; and therefore no offset was committed to
+    // tp. The committed offset should be null. This is intentional.
+    assertNull(committedOffset.get(tp))
+    assertTrue(consumer.assignment.contains(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndCommitSyncNotCommitted(quorum: String, groupProtocol: String): Unit = {
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props)
+    val producer = createProducer()
+    val numRecords = 10000
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync()
+    val committedOffset = consumer.committed(Set(tp).asJava)
+    assertNotNull(committedOffset)
+    // No valid fetch position due to the absence of consumer.poll; and therefore no offset was committed to
+    // tp. The committed offset should be null. This is intentional.
+    assertNull(committedOffset.get(tp))
+    assertTrue(consumer.assignment.contains(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndCommitSyncAllConsumed(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10000
+
+    val producer = createProducer()
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props)
+    consumer.assign(List(tp).asJava)
+    consumer.seek(tp, 0)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
+
+    consumer.commitSync()
+    val committedOffset = consumer.committed(Set(tp).asJava)
+    assertNotNull(committedOffset)
+    assertNotNull(committedOffset.get(tp))
+    assertEquals(numRecords, committedOffset.get(tp).offset())
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndConsume(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10
+
+    val producer = createProducer()
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
+
+    assertEquals(numRecords, consumer.position(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndConsumeSkippingPosition(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10
+
+    val producer = createProducer()
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+    consumer.assign(List(tp).asJava)
+    val offset = 1
+    consumer.seek(tp, offset)
+    consumeAndVerifyRecords(consumer = consumer, numRecords - offset, startingOffset = offset,
+      startingKeyAndValueIndex = offset, startingTimestamp = startingTimestamp + offset)
+
+    assertEquals(numRecords, consumer.position(tp))
+  }
+
+  // partitionsFor not implemented in consumer group protocol and this test requires ZK also
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @CsvSource(Array(
+    "zk, classic"
+  ))
+  def testAssignAndConsumeWithLeaderChangeValidatingPositions(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10
+    val producer = createProducer()
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
+
+    // Force leader epoch change to trigger position validation
+    var parts: mutable.Buffer[PartitionInfo] = null
+    while (parts == null)
+      parts = consumer.partitionsFor(tp.topic()).asScala
+    val leader = parts.head.leader().id()
+    this.servers(leader).shutdown()
+    this.servers(leader).startup()
+
+    // Consume after leader change
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 10,
+      startingTimestamp = startingTimestamp)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndFetchCommittedOffsets(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 100
+    val startingTimestamp = System.currentTimeMillis()
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props)
+    consumer.assign(List(tp).asJava)
+    // First consumer consumes and commits offsets
+    consumer.seek(tp, 0)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0,
+      startingTimestamp = startingTimestamp)
+    consumer.commitSync()
+    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
+    // We should see the committed offsets from another consumer
+    val anotherConsumer = createConsumer(configOverrides = props)
+    anotherConsumer.assign(List(tp).asJava)
+    assertEquals(numRecords, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndConsumeFromCommittedOffsets(quorum: String, groupProtocol: String): Unit = {
+    val producer = createProducer()
+    val numRecords = 100
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords = numRecords, tp, startingTimestamp = startingTimestamp)
+
+    // Commit offset with first consumer
+    val props = new Properties()
+    props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "group1")
+    val consumer = createConsumer(configOverrides = props)
+    consumer.assign(List(tp).asJava)
+    val offset = 10
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(offset)))
+      .asJava)
+    assertEquals(offset, consumer.committed(Set(tp).asJava).get(tp).offset)
+    consumer.close()
+
+    // Consume from committed offsets with another consumer in same group
+    val anotherConsumer = createConsumer(configOverrides = props)
+    assertEquals(offset, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    anotherConsumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = anotherConsumer, numRecords - offset,
+      startingOffset = offset, startingKeyAndValueIndex = offset,
+      startingTimestamp = startingTimestamp + offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAssignAndRetrievingCommittedOffsetsMultipleTimes(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 100
+    val startingTimestamp = System.currentTimeMillis()
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+
+    val props = new Properties()
+    val consumer = createConsumer(configOverrides = props)
+    consumer.assign(List(tp).asJava)
+
+    // Consume and commit offsets
+    consumer.seek(tp, 0)
+    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0,
+      startingTimestamp = startingTimestamp)
+    consumer.commitSync()
+
+    // Check committed offsets twice with same consumer
+    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
+  }
+
+}
+
+object PlaintextConsumerAssignTest {
+  def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
new file mode 100644
index 0000000000000..e126774a13763
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
@@ -0,0 +1,403 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.{TestInfoUtils, TestUtils}
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.UnsupportedAssignorException
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, CsvSource, MethodSource}
+
+import java.util
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.locks.ReentrantLock
+import java.util.stream.Stream
+import scala.collection.mutable.Buffer
+import scala.jdk.CollectionConverters._
+
+/**
+ * Integration tests for the consumer that covers assignors logic (client and server side assignors)
+ */
+@Timeout(600)
+class PlaintextConsumerAssignorsTest extends AbstractConsumerTest {
+
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testRoundRobinAssignment(quorum: String, groupProtocol: String): Unit = {
+    // 1 consumer using round-robin assignment
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "roundrobin-group")
+    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[RoundRobinAssignor].getName)
+    val consumer = createConsumer()
+
+    // create two new topics, each having 2 partitions
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+    val producer = createProducer()
+    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100) ++
+      createTopicAndSendRecords(producer, topic2, 2, 100)
+
+    assertEquals(0, consumer.assignment().size)
+
+    // subscribe to two topics
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
+
+    // add one more topic with 2 partitions
+    val topic3 = "topic3"
+    createTopicAndSendRecords(producer, topic3, 2, 100)
+
+    val newExpectedAssignment = expectedAssignment ++ Set(new TopicPartition(topic3, 0), new TopicPartition(topic3, 1))
+    consumer.subscribe(List(topic1, topic2, topic3).asJava)
+    awaitAssignment(consumer, newExpectedAssignment)
+
+    // remove the topic we just added
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
+
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
+  }
+
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testMultiConsumerRoundRobinAssignor(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "roundrobin-group")
+    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[RoundRobinAssignor].getName)
+
+    // create two new topics, total number of partitions must be greater than number of consumers
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+    val producer = createProducer()
+    val subscriptions = createTopicAndSendRecords(producer, topic1, 5, 100) ++
+      createTopicAndSendRecords(producer, topic2, 8, 100)
+
+    // create a group of consumers, subscribe the consumers to all the topics and start polling
+    // for the topic partition assignment
+    val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(10, List(topic1, topic2), subscriptions)
+    try {
+      validateGroupAssignment(consumerPollers, subscriptions)
+
+      // add one more consumer and validate re-assignment
+      addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers,
+        List(topic1, topic2), subscriptions, "roundrobin-group")
+    } finally {
+      consumerPollers.foreach(_.shutdown())
+    }
+  }
+
+  /**
+   * This test runs the following scenario to verify sticky assignor behavior.
+   * Topics: single-topic, with random number of partitions, where #par is 10, 20, 30, 40, 50, 60, 70, 80, 90, or 100
+   * Consumers: 9 consumers subscribed to the single topic
+   * Expected initial assignment: partitions are assigned to consumers in a round robin fashion.
+   *  - (#par mod 9) consumers will get (#par / 9 + 1) partitions, and the rest get (#par / 9) partitions
+   *    Then consumer #10 is added to the list (subscribing to the same single topic)
+   *    Expected new assignment:
+   *  - (#par / 10) partition per consumer, where one partition from each of the early (#par mod 9) consumers
+   *    will move to consumer #10, leading to a total of (#par mod 9) partition movement
+   */
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testMultiConsumerStickyAssignor(quorum: String, groupProtocol: String): Unit = {
+
+    def reverse(m: Map[Long, Set[TopicPartition]]) =
+      m.values.toSet.flatten.map(v => (v, m.keys.filter(m(_).contains(v)).head)).toMap
+
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "sticky-group")
+    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[StickyAssignor].getName)
+
+    // create one new topic
+    val topic = "single-topic"
+    val rand = 1 + scala.util.Random.nextInt(10)
+    val producer = createProducer()
+    val partitions = createTopicAndSendRecords(producer, topic, rand * 10, 100)
+
+    // create a group of consumers, subscribe the consumers to the single topic and start polling
+    // for the topic partition assignment
+    val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(9, List(topic), partitions)
+    validateGroupAssignment(consumerPollers, partitions)
+    val prePartition2PollerId = reverse(consumerPollers.map(poller => (poller.getId, poller.consumerAssignment())).toMap)
+
+    // add one more consumer and validate re-assignment
+    addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers, List(topic), partitions, "sticky-group")
+
+    val postPartition2PollerId = reverse(consumerPollers.map(poller => (poller.getId, poller.consumerAssignment())).toMap)
+    val keys = prePartition2PollerId.keySet.union(postPartition2PollerId.keySet)
+    var changes = 0
+    keys.foreach { key =>
+      val preVal = prePartition2PollerId.get(key)
+      val postVal = postPartition2PollerId.get(key)
+      if (preVal.nonEmpty && postVal.nonEmpty) {
+        if (preVal.get != postVal.get)
+          changes += 1
+      } else
+        changes += 1
+    }
+
+    consumerPollers.foreach(_.shutdown())
+
+    assertEquals(rand, changes, "Expected only two topic partitions that have switched to other consumers.")
+  }
+
+  /**
+   * This test re-uses BaseConsumerTest's consumers.
+   * As a result, it is testing the default assignment strategy set by BaseConsumerTest
+   * It tests the assignment results is expected using default assignor (i.e. Range assignor)
+   */
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testMultiConsumerDefaultAssignorAndVerifyAssignment(quorum: String, groupProtocol: String): Unit = {
+    // create two new topics, each having 3 partitions
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+
+    createTopic(topic1, 3)
+    createTopic(topic2, 3)
+
+    val consumersInGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
+    consumersInGroup += createConsumer()
+    consumersInGroup += createConsumer()
+
+    val tp1_0 = new TopicPartition(topic1, 0)
+    val tp1_1 = new TopicPartition(topic1, 1)
+    val tp1_2 = new TopicPartition(topic1, 2)
+    val tp2_0 = new TopicPartition(topic2, 0)
+    val tp2_1 = new TopicPartition(topic2, 1)
+    val tp2_2 = new TopicPartition(topic2, 2)
+
+    val subscriptions = Set(tp1_0, tp1_1, tp1_2, tp2_0, tp2_1, tp2_2)
+    val consumerPollers = subscribeConsumers(consumersInGroup, List(topic1, topic2))
+
+    val expectedAssignment = Buffer(Set(tp1_0, tp1_1, tp2_0, tp2_1), Set(tp1_2, tp2_2))
+
+    try {
+      validateGroupAssignment(consumerPollers, subscriptions, expectedAssignment = expectedAssignment)
+    } finally {
+      consumerPollers.foreach(_.shutdown())
+    }
+  }
+
+  /**
+   * This test re-uses BaseConsumerTest's consumers.
+   * As a result, it is testing the default assignment strategy set by BaseConsumerTest
+   */
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testMultiConsumerDefaultAssignor(quorum: String, groupProtocol: String): Unit = {
+    // use consumers and topics defined in this class + one more topic
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 100, tp)
+    sendRecords(producer, numRecords = 100, tp2)
+    val topic1 = "topic1"
+    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 5, 100)
+
+    // subscribe all consumers to all topics and validate the assignment
+
+    val consumersInGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
+    consumersInGroup += createConsumer()
+    consumersInGroup += createConsumer()
+
+    val consumerPollers = subscribeConsumers(consumersInGroup, List(topic, topic1))
+    try {
+      validateGroupAssignment(consumerPollers, subscriptions)
+
+      // add 2 more consumers and validate re-assignment
+      addConsumersToGroupAndWaitForGroupAssignment(2, consumersInGroup, consumerPollers, List(topic, topic1), subscriptions)
+
+      // add one more topic and validate partition re-assignment
+      val topic2 = "topic2"
+      val expandedSubscriptions = subscriptions ++ createTopicAndSendRecords(producer, topic2, 3, 100)
+      changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers, List(topic, topic1, topic2), expandedSubscriptions)
+
+      // remove the topic we just added and validate re-assignment
+      changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers, List(topic, topic1), subscriptions)
+
+    } finally {
+      consumerPollers.foreach(_.shutdown())
+    }
+  }
+
+  // Remote assignors only supported with consumer group protocol
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @CsvSource(Array(
+    "kraft+kip848, consumer"
+  ))
+  def testRemoteAssignorInvalid(quorum: String, groupProtocol: String): Unit = {
+    // 1 consumer using invalid remote assignor
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "invalid-assignor-group")
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "invalid")
+    val consumer = createConsumer()
+
+    // create two new topics, each having 2 partitions
+    val topic1 = "topic1"
+    val producer = createProducer()
+    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100)
+
+    assertEquals(0, consumer.assignment().size)
+
+    // subscribe to two topics
+    consumer.subscribe(List(topic1).asJava)
+
+    val e: UnsupportedAssignorException = assertThrows(
+      classOf[UnsupportedAssignorException],
+      () => awaitAssignment(consumer, expectedAssignment)
+    )
+
+    assertTrue(e.getMessage.startsWith("ServerAssignor invalid is not supported. " +
+      "Supported assignors: "))
+  }
+
+  // Remote assignors only supported with consumer group protocol
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @CsvSource(Array(
+    "kraft+kip848, consumer"
+  ))
+  def testRemoteAssignorRange(quorum: String, groupProtocol: String): Unit = {
+    // 1 consumer using range assignment
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "range-group")
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "range")
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "30000")
+    val consumer = createConsumer()
+
+    // create two new topics, each having 2 partitions
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+    val producer = createProducer()
+    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100) ++
+      createTopicAndSendRecords(producer, topic2, 2, 100)
+
+    assertEquals(0, consumer.assignment().size)
+
+    // subscribe to two topics
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
+
+    // add one more topic with 2 partitions
+    val topic3 = "topic3"
+    val additionalAssignment = createTopicAndSendRecords(producer, topic3, 2, 100)
+
+    val newExpectedAssignment = expectedAssignment ++ additionalAssignment
+    consumer.subscribe(List(topic1, topic2, topic3).asJava)
+    awaitAssignment(consumer, newExpectedAssignment)
+
+    // remove the topic we just added
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
+
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
+  }
+
+  // Only the classic group protocol supports client-side assignors
+  @ParameterizedTest
+  @CsvSource(Array(
+    "org.apache.kafka.clients.consumer.CooperativeStickyAssignor,   zk",
+    "org.apache.kafka.clients.consumer.RangeAssignor,               zk",
+    "org.apache.kafka.clients.consumer.CooperativeStickyAssignor,   kraft",
+    "org.apache.kafka.clients.consumer.RangeAssignor,               kraft"
+  ))
+  def testRebalanceAndRejoin(assignmentStrategy: String, quorum: String): Unit = {
+    // create 2 consumers
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, "classic")
+    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "rebalance-and-rejoin-group")
+    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, assignmentStrategy)
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
+    val consumer1 = createConsumer()
+    val consumer2 = createConsumer()
+
+    // create a new topic, have 2 partitions
+    val topic = "topic1"
+    val producer = createProducer()
+    val expectedAssignment = createTopicAndSendRecords(producer, topic, 2, 100)
+
+    assertEquals(0, consumer1.assignment().size)
+    assertEquals(0, consumer2.assignment().size)
+
+    val lock = new ReentrantLock()
+    var generationId1 = -1
+    var memberId1 = ""
+    val customRebalanceListener = new ConsumerRebalanceListener {
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+      }
+
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+        if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
+          fail(s"Time out while awaiting for lock.")
+        }
+        try {
+          generationId1 = consumer1.groupMetadata().generationId()
+          memberId1 = consumer1.groupMetadata().memberId()
+        } finally {
+          lock.unlock()
+        }
+      }
+    }
+    val consumerPoller1 = new ConsumerAssignmentPoller(consumer1, List(topic), Set.empty, customRebalanceListener)
+    consumerPoller1.start()
+    TestUtils.waitUntilTrue(() => consumerPoller1.consumerAssignment() == expectedAssignment,
+      s"Timed out while awaiting expected assignment change to $expectedAssignment.")
+
+    // Since the consumer1 already completed the rebalance,
+    // the `onPartitionsAssigned` rebalance listener will be invoked to set the generationId and memberId
+    var stableGeneration = -1
+    var stableMemberId1 = ""
+    if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
+      fail(s"Time out while awaiting for lock.")
+    }
+    try {
+      stableGeneration = generationId1
+      stableMemberId1 = memberId1
+    } finally {
+      lock.unlock()
+    }
+
+    val consumerPoller2 = subscribeConsumerAndStartPolling(consumer2, List(topic))
+    TestUtils.waitUntilTrue(() => consumerPoller1.consumerAssignment().size == 1,
+      s"Timed out while awaiting expected assignment size change to 1.")
+    TestUtils.waitUntilTrue(() => consumerPoller2.consumerAssignment().size == 1,
+      s"Timed out while awaiting expected assignment size change to 1.")
+
+    if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
+      fail(s"Time out while awaiting for lock.")
+    }
+    try {
+      if (assignmentStrategy.equals(classOf[CooperativeStickyAssignor].getName)) {
+        // cooperative rebalance should rebalance twice before finally stable
+        assertEquals(stableGeneration + 2, generationId1)
+      } else {
+        // eager rebalance should rebalance once before finally stable
+        assertEquals(stableGeneration + 1, generationId1)
+      }
+      assertEquals(stableMemberId1, memberId1)
+    } finally {
+      lock.unlock()
+    }
+
+    consumerPoller1.shutdown()
+    consumerPoller2.shutdown()
+  }
+
+}
+
+object PlaintextConsumerAssignorsTest {
+  def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerFetchTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerFetchTest.scala
new file mode 100644
index 0000000000000..3012a6d25fbb7
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerFetchTest.scala
@@ -0,0 +1,243 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.TestInfoUtils
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, MethodSource}
+import org.apache.kafka.common.TopicPartition
+
+import java.time.Duration
+import java.util.stream.Stream
+import scala.jdk.CollectionConverters._
+
+/**
+ * Integration tests for the consumer that covers fetching logic
+ */
+@Timeout(600)
+class PlaintextConsumerFetchTest extends AbstractConsumerTest {
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchInvalidOffset(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none")
+    val consumer = createConsumer(configOverrides = this.consumerConfig)
+
+    // produce one record
+    val totalRecords = 2
+    val producer = createProducer()
+    sendRecords(producer, totalRecords, tp)
+    consumer.assign(List(tp).asJava)
+
+    // poll should fail because there is no offset reset strategy set.
+    // we fail only when resetting positions after coordinator is known, so using a long timeout.
+    assertThrows(classOf[NoOffsetForPartitionException], () => consumer.poll(Duration.ofMillis(15000)))
+
+    // seek to out of range position
+    val outOfRangePos = totalRecords + 1
+    consumer.seek(tp, outOfRangePos)
+    val e = assertThrows(classOf[OffsetOutOfRangeException], () => consumer.poll(Duration.ofMillis(20000)))
+    val outOfRangePartitions = e.offsetOutOfRangePartitions()
+    assertNotNull(outOfRangePartitions)
+    assertEquals(1, outOfRangePartitions.size)
+    assertEquals(outOfRangePos.toLong, outOfRangePartitions.get(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchOutOfRangeOffsetResetConfigEarliest(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+    // ensure no in-flight fetch request so that the offset can be reset immediately
+    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0")
+    val consumer = createConsumer(configOverrides = this.consumerConfig)
+    val totalRecords = 10L
+
+    val producer = createProducer()
+    val startingTimestamp = 0
+    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = startingTimestamp)
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = totalRecords.toInt, startingOffset = 0)
+    // seek to out of range position
+    val outOfRangePos = totalRecords + 1
+    consumer.seek(tp, outOfRangePos)
+    // assert that poll resets to the beginning position
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0)
+  }
+
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchOutOfRangeOffsetResetConfigLatest(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest")
+    // ensure no in-flight fetch request so that the offset can be reset immediately
+    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0")
+    val consumer = createConsumer(configOverrides = this.consumerConfig)
+    val totalRecords = 10L
+
+    val producer = createProducer()
+    val startingTimestamp = 0
+    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = startingTimestamp)
+    consumer.assign(List(tp).asJava)
+    consumer.seek(tp, 0)
+    // consume some, but not all of the records
+    consumeAndVerifyRecords(consumer = consumer, numRecords = totalRecords.toInt / 2, startingOffset = 0)
+    // seek to out of range position
+    val outOfRangePos = totalRecords + 17 // arbitrary, much higher offset
+    consumer.seek(tp, outOfRangePos)
+    // assert that poll resets to the ending position
+    assertTrue(consumer.poll(Duration.ofMillis(50)).isEmpty)
+    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = totalRecords)
+    val nextRecord = consumer.poll(Duration.ofMillis(50)).iterator().next()
+    // ensure the seek went to the last known record at the time of the previous poll
+    assertEquals(totalRecords, nextRecord.offset())
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchRecordLargerThanFetchMaxBytes(quorum: String, groupProtocol: String): Unit = {
+    val maxFetchBytes = 10 * 1024
+    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxFetchBytes.toString)
+    checkLargeRecord(maxFetchBytes + 1)
+  }
+
+  private def checkLargeRecord(producerRecordSize: Int): Unit = {
+    val consumer = createConsumer()
+
+    // produce a record that is larger than the configured fetch size
+    val record = new ProducerRecord(tp.topic(), tp.partition(), "key".getBytes,
+      new Array[Byte](producerRecordSize))
+    val producer = createProducer()
+    producer.send(record)
+
+    // consuming a record that is too large should succeed since KIP-74
+    consumer.assign(List(tp).asJava)
+    val records = consumer.poll(Duration.ofMillis(20000))
+    assertEquals(1, records.count)
+    val consumerRecord = records.iterator().next()
+    assertEquals(0L, consumerRecord.offset)
+    assertEquals(tp.topic(), consumerRecord.topic())
+    assertEquals(tp.partition(), consumerRecord.partition())
+    assertArrayEquals(record.key(), consumerRecord.key())
+    assertArrayEquals(record.value(), consumerRecord.value())
+  }
+
+  /** We should only return a large record if it's the first record in the first non-empty partition of the fetch request */
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchHonoursFetchSizeIfLargeRecordNotFirst(quorum: String, groupProtocol: String): Unit = {
+    val maxFetchBytes = 10 * 1024
+    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxFetchBytes.toString)
+    checkFetchHonoursSizeIfLargeRecordNotFirst(maxFetchBytes)
+  }
+
+  private def checkFetchHonoursSizeIfLargeRecordNotFirst(largeProducerRecordSize: Int): Unit = {
+    val consumer = createConsumer()
+
+    val smallRecord = new ProducerRecord(tp.topic(), tp.partition(), "small".getBytes,
+      "value".getBytes)
+    val largeRecord = new ProducerRecord(tp.topic(), tp.partition(), "large".getBytes,
+      new Array[Byte](largeProducerRecordSize))
+
+    val producer = createProducer()
+    producer.send(smallRecord).get
+    producer.send(largeRecord).get
+
+    // we should only get the small record in the first `poll`
+    consumer.assign(List(tp).asJava)
+    val records = consumer.poll(Duration.ofMillis(20000))
+    assertEquals(1, records.count)
+    val consumerRecord = records.iterator().next()
+    assertEquals(0L, consumerRecord.offset)
+    assertEquals(tp.topic(), consumerRecord.topic())
+    assertEquals(tp.partition(), consumerRecord.partition())
+    assertArrayEquals(smallRecord.key(), consumerRecord.key())
+    assertArrayEquals(smallRecord.value(), consumerRecord.value())
+  }
+
+  /** We should only return a large record if it's the first record in the first partition of the fetch request */
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchHonoursMaxPartitionFetchBytesIfLargeRecordNotFirst(quorum: String, groupProtocol: String): Unit = {
+    val maxPartitionFetchBytes = 10 * 1024
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, maxPartitionFetchBytes.toString)
+    checkFetchHonoursSizeIfLargeRecordNotFirst(maxPartitionFetchBytes)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testFetchRecordLargerThanMaxPartitionFetchBytes(quorum: String, groupProtocol: String): Unit = {
+    val maxPartitionFetchBytes = 10 * 1024
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, maxPartitionFetchBytes.toString)
+    checkLargeRecord(maxPartitionFetchBytes + 1)
+  }
+
+  /** Test that we consume all partitions if fetch max bytes and max.partition.fetch.bytes are low */
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testLowMaxFetchSizeForRequestAndPartition(quorum: String, groupProtocol: String): Unit = {
+    // one of the effects of this is that there will be some log reads where `0 > remaining limit bytes < message size`
+    // and we don't return the message because it's not the first message in the first non-empty partition of the fetch
+    // this behaves a little different than when remaining limit bytes is 0 and it's important to test it
+    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, "500")
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, "100")
+
+    // Avoid a rebalance while the records are being sent (the default is 6 seconds)
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 20000.toString)
+    val consumer = createConsumer()
+
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+    val topic3 = "topic3"
+    val partitionCount = 30
+    val topics = Seq(topic1, topic2, topic3)
+    topics.foreach { topicName =>
+      createTopic(topicName, partitionCount, brokerCount)
+    }
+
+    val partitions = topics.flatMap { topic =>
+      (0 until partitionCount).map(new TopicPartition(topic, _))
+    }
+
+    assertEquals(0, consumer.assignment().size)
+
+    consumer.subscribe(List(topic1, topic2, topic3).asJava)
+
+    awaitAssignment(consumer, partitions.toSet)
+
+    val producer = createProducer()
+
+    val producerRecords = partitions.flatMap(sendRecords(producer, numRecords = partitionCount, _))
+
+    val consumerRecords = consumeRecords(consumer, producerRecords.size)
+
+    val expected = producerRecords.map { record =>
+      (record.topic, record.partition, new String(record.key), new String(record.value), record.timestamp)
+    }.toSet
+
+    val actual = consumerRecords.map { record =>
+      (record.topic, record.partition, new String(record.key), new String(record.value), record.timestamp)
+    }.toSet
+
+    assertEquals(expected, actual)
+  }
+
+}
+
+object PlaintextConsumerFetchTest {
+  def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerPollTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerPollTest.scala
new file mode 100644
index 0000000000000..3bcc449162619
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerPollTest.scala
@@ -0,0 +1,284 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.TestInfoUtils
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.common.{MetricName, TopicPartition}
+import org.apache.kafka.common.utils.Utils
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, MethodSource}
+
+import java.time.Duration
+import java.util
+import java.util.stream.Stream
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.Buffer
+
+/**
+ * Integration tests for the consumer that covers the poll logic
+ */
+@Timeout(600)
+class PlaintextConsumerPollTest extends AbstractConsumerTest {
+
+  // Deprecated poll(timeout) not supported for consumer group protocol
+  @deprecated("poll(Duration) is the replacement", since = "2.0")
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testDeprecatedPollBlocksForAssignment(quorum: String, groupProtocol: String): Unit = {
+    val consumer = createConsumer()
+    consumer.subscribe(Set(topic).asJava)
+    consumer.poll(0)
+    assertEquals(Set(tp, tp2), consumer.assignment().asScala)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
+    val maxPollRecords = 2
+    val numRecords = 10000
+
+    val producer = createProducer()
+    val startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
+
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer, numRecords = numRecords, startingOffset = 0, maxPollRecords = maxPollRecords,
+      startingTimestamp = startingTimestamp)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollIntervalMs(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 2000.toString)
+
+    val consumer = createConsumer()
+
+    val listener = new TestConsumerReassignmentListener()
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
+    assertEquals(1, listener.callsToAssigned)
+    assertEquals(0, listener.callsToRevoked)
+
+    // after we extend longer than max.poll a rebalance should be triggered
+    // NOTE we need to have a relatively much larger value than max.poll to let heartbeat expired for sure
+    Thread.sleep(3000)
+
+    awaitRebalance(consumer, listener)
+    assertEquals(2, listener.callsToAssigned)
+    assertEquals(1, listener.callsToRevoked)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollIntervalMsDelayInRevocation(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 5000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 1000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
+
+    val consumer = createConsumer()
+    var commitCompleted = false
+    var committedPosition: Long = -1
+
+    val listener = new TestConsumerReassignmentListener {
+      override def onPartitionsLost(partitions: util.Collection[TopicPartition]): Unit = {}
+
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+        if (!partitions.isEmpty && partitions.contains(tp)) {
+          // on the second rebalance (after we have joined the group initially), sleep longer
+          // than session timeout and then try a commit. We should still be in the group,
+          // so the commit should succeed
+          Utils.sleep(1500)
+          committedPosition = consumer.position(tp)
+          consumer.commitSync(Map(tp -> new OffsetAndMetadata(committedPosition)).asJava)
+          commitCompleted = true
+        }
+        super.onPartitionsRevoked(partitions)
+      }
+    }
+
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
+
+    // force a rebalance to trigger an invocation of the revocation callback while in the group
+    consumer.subscribe(List("otherTopic").asJava, listener)
+    awaitRebalance(consumer, listener)
+
+    assertEquals(0, committedPosition)
+    assertTrue(commitCompleted)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollIntervalMsDelayInAssignment(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 5000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 1000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
+
+    val consumer = createConsumer()
+    val listener = new TestConsumerReassignmentListener {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+        // sleep longer than the session timeout, we should still be in the group after invocation
+        Utils.sleep(1500)
+        super.onPartitionsAssigned(partitions)
+      }
+    }
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
+
+    // We should still be in the group after this invocation
+    ensureNoRebalance(consumer, listener)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMaxPollIntervalMsShorterThanPollTimeout(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
+
+    val consumer = createConsumer()
+    val listener = new TestConsumerReassignmentListener
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
+
+    val callsToAssignedAfterFirstRebalance = listener.callsToAssigned
+
+    consumer.poll(Duration.ofMillis(2000))
+
+    // If the poll poll above times out, it would trigger a rebalance.
+    // Leave some time for the rebalance to happen and check for the rebalance event.
+    consumer.poll(Duration.ofMillis(500))
+    consumer.poll(Duration.ofMillis(500))
+
+    assertEquals(callsToAssignedAfterFirstRebalance, listener.callsToAssigned)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testPerPartitionLeadWithMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
+    val numMessages = 1000
+    val maxPollRecords = 10
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    awaitNonEmptyRecords(consumer, tp)
+
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLeadWithMaxPollRecords")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val lead = consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags))
+    assertEquals(maxPollRecords, lead.metricValue().asInstanceOf[Double], s"The lead should be $maxPollRecords")
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testPerPartitionLagWithMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
+    val numMessages = 1000
+    val maxPollRecords = 10
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val records = awaitNonEmptyRecords(consumer, tp)
+
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLagWithMaxPollRecords")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val lag = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags))
+
+    assertEquals(numMessages - records.count, lag.metricValue.asInstanceOf[Double], epsilon, s"The lag should be ${numMessages - records.count}")
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMultiConsumerSessionTimeoutOnStopPolling(quorum: String, groupProtocol: String): Unit = {
+    runMultiConsumerSessionTimeoutTest(false)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testMultiConsumerSessionTimeoutOnClose(quorum: String, groupProtocol: String): Unit = {
+    runMultiConsumerSessionTimeoutTest(true)
+  }
+
+  def runMultiConsumerSessionTimeoutTest(closeConsumer: Boolean): Unit = {
+    // use consumers defined in this class plus one additional consumer
+    // Use topic defined in this class + one additional topic
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 100, tp)
+    sendRecords(producer, numRecords = 100, tp2)
+    val topic1 = "topic1"
+    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 6, 100)
+
+    // first subscribe consumers that are defined in this class
+    val consumerPollers = Buffer[ConsumerAssignmentPoller]()
+    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
+    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
+
+    // create one more consumer and add it to the group; we will timeout this consumer
+    val timeoutConsumer = createConsumer()
+    val timeoutPoller = subscribeConsumerAndStartPolling(timeoutConsumer, List(topic, topic1))
+    consumerPollers += timeoutPoller
+
+    // validate the initial assignment
+    validateGroupAssignment(consumerPollers, subscriptions)
+
+    // stop polling and close one of the consumers, should trigger partition re-assignment among alive consumers
+    timeoutPoller.shutdown()
+    consumerPollers -= timeoutPoller
+    if (closeConsumer)
+      timeoutConsumer.close()
+
+    validateGroupAssignment(consumerPollers, subscriptions,
+      Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after one consumer left"), 3 * groupMaxSessionTimeoutMs)
+
+    // done with pollers and consumers
+    for (poller <- consumerPollers)
+      poller.shutdown()
+  }
+}
+
+object PlaintextConsumerPollTest {
+  def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly()
+
+  def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerSubscriptionTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerSubscriptionTest.scala
new file mode 100644
index 0000000000000..7ac0501bb5425
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerSubscriptionTest.scala
@@ -0,0 +1,237 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.TestInfoUtils
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.common.TopicPartition
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, MethodSource}
+
+import java.util.regex.Pattern
+import java.util.stream.Stream
+import scala.jdk.CollectionConverters._
+
+/**
+ * Integration tests for the consumer that covers the subscribe and unsubscribe logic.
+ */
+@Timeout(600)
+class PlaintextConsumerSubscriptionTest extends AbstractConsumerTest {
+
+  /**
+   * Verifies that pattern subscription performs as expected.
+   * The pattern matches the topics 'topic' and 'tblablac', but not 'tblablak' or 'tblab1'.
+   * It is expected that the consumer is subscribed to all partitions of 'topic' and
+   * 'tblablac' after the subscription when metadata is refreshed.
+   * When a new topic 'tsomec' is added afterwards, it is expected that upon the next
+   * metadata refresh the consumer becomes subscribed to this new topic and all partitions
+   * of that topic are assigned to it.
+   */
+  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testPatternSubscription(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    val topic1 = "tblablac" // matches subscribed pattern
+    createTopic(topic1, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
+
+    val topic2 = "tblablak" // does not match subscribed pattern
+    createTopic(topic2, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic2, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic2, 1))
+
+    val topic3 = "tblab1" // does not match subscribed pattern
+    createTopic(topic3, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 1))
+
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment().size)
+
+    val pattern = Pattern.compile("t.*c")
+    consumer.subscribe(pattern, new TestConsumerReassignmentListener)
+
+    var assignment = Set(
+      new TopicPartition(topic, 0),
+      new TopicPartition(topic, 1),
+      new TopicPartition(topic1, 0),
+      new TopicPartition(topic1, 1))
+    awaitAssignment(consumer, assignment)
+
+    val topic4 = "tsomec" // matches subscribed pattern
+    createTopic(topic4, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 1))
+
+    assignment ++= Set(
+      new TopicPartition(topic4, 0),
+      new TopicPartition(topic4, 1))
+    awaitAssignment(consumer, assignment)
+
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
+  }
+
+  /**
+   * Verifies that a second call to pattern subscription succeeds and performs as expected.
+   * The initial subscription is to a pattern that matches two topics 'topic' and 'foo'.
+   * The second subscription is to a pattern that matches 'foo' and a new topic 'bar'.
+   * It is expected that the consumer is subscribed to all partitions of 'topic' and 'foo' after
+   * the first subscription, and to all partitions of 'foo' and 'bar' after the second.
+   * The metadata refresh interval is intentionally increased to a large enough value to guarantee
+   * that it is the subscription call that triggers a metadata refresh, and not the timeout.
+   */
+  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testSubsequentPatternSubscription(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "30000")
+    val consumer = createConsumer()
+
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords = numRecords, tp)
+
+    // the first topic ('topic')  matches first subscription pattern only
+
+    val fooTopic = "foo" // matches both subscription patterns
+    createTopic(fooTopic, 1, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(fooTopic, 0))
+
+    assertEquals(0, consumer.assignment().size)
+
+    val pattern1 = Pattern.compile(".*o.*") // only 'topic' and 'foo' match this
+    consumer.subscribe(pattern1, new TestConsumerReassignmentListener)
+
+    var assignment = Set(
+      new TopicPartition(topic, 0),
+      new TopicPartition(topic, 1),
+      new TopicPartition(fooTopic, 0))
+    awaitAssignment(consumer, assignment)
+
+    val barTopic = "bar" // matches the next subscription pattern
+    createTopic(barTopic, 1, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(barTopic, 0))
+
+    val pattern2 = Pattern.compile("...") // only 'foo' and 'bar' match this
+    consumer.subscribe(pattern2, new TestConsumerReassignmentListener)
+    assignment --= Set(
+      new TopicPartition(topic, 0),
+      new TopicPartition(topic, 1))
+    assignment ++= Set(
+      new TopicPartition(barTopic, 0))
+    awaitAssignment(consumer, assignment)
+
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
+  }
+
+  /**
+   * Verifies that pattern unsubscription performs as expected.
+   * The pattern matches the topics 'topic' and 'tblablac'.
+   * It is expected that the consumer is subscribed to all partitions of 'topic' and
+   * 'tblablac' after the subscription when metadata is refreshed.
+   * When consumer unsubscribes from all its subscriptions, it is expected that its
+   * assignments are cleared right away.
+   */
+  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
+  def testPatternUnsubscription(quorum: String, groupProtocol: String): Unit = {
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    val topic1 = "tblablac" // matches the subscription pattern
+    createTopic(topic1, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
+
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment().size)
+
+    consumer.subscribe(Pattern.compile("t.*c"), new TestConsumerReassignmentListener)
+    val assignment = Set(
+      new TopicPartition(topic, 0),
+      new TopicPartition(topic, 1),
+      new TopicPartition(topic1, 0),
+      new TopicPartition(topic1, 1))
+    awaitAssignment(consumer, assignment)
+
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testExpandingTopicSubscriptions(quorum: String, groupProtocol: String): Unit = {
+    val otherTopic = "other"
+    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, initialAssignment)
+
+    createTopic(otherTopic, 2, brokerCount)
+    val expandedAssignment = initialAssignment ++ Set(new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
+    consumer.subscribe(List(topic, otherTopic).asJava)
+    awaitAssignment(consumer, expandedAssignment)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testShrinkingTopicSubscriptions(quorum: String, groupProtocol: String): Unit = {
+    val otherTopic = "other"
+    createTopic(otherTopic, 2, brokerCount)
+    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic, otherTopic).asJava)
+    awaitAssignment(consumer, initialAssignment)
+
+    val shrunkenAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, shrunkenAssignment)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testUnsubscribeTopic(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "100") // timeout quickly to avoid slow test
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "30")
+    val consumer = createConsumer()
+
+    val listener = new TestConsumerReassignmentListener()
+    consumer.subscribe(List(topic).asJava, listener)
+
+    // the initial subscription should cause a callback execution
+    awaitRebalance(consumer, listener)
+
+    consumer.subscribe(List[String]().asJava)
+    assertEquals(0, consumer.assignment.size())
+  }
+
+}
+
+object PlaintextConsumerSubscriptionTest {
+  def getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly()
+
+  def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 01dbb206a07cb..8784e5b43fc23 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -15,30 +15,24 @@ package kafka.api
 import java.time.Duration
 import java.util
 import java.util.Arrays.asList
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.locks.ReentrantLock
-import java.util.regex.Pattern
 import java.util.{Locale, Optional, Properties}
 import kafka.server.{KafkaBroker, QuotaType}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin.{NewPartitions, NewTopic}
 import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
-import org.apache.kafka.common.{KafkaException, MetricName, PartitionInfo, TopicPartition}
+import org.apache.kafka.common.{KafkaException, MetricName, TopicPartition}
 import org.apache.kafka.common.config.TopicConfig
-import org.apache.kafka.common.errors.{InvalidGroupIdException, InvalidTopicException, UnsupportedAssignorException}
+import org.apache.kafka.common.errors.{InvalidGroupIdException, InvalidTopicException}
 import org.apache.kafka.common.header.Headers
 import org.apache.kafka.common.record.{CompressionType, TimestampType}
 import org.apache.kafka.common.serialization._
-import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.test.{MockConsumerInterceptor, MockProducerInterceptor}
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Timeout
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.{CsvSource, MethodSource}
 
-import scala.collection.mutable
-import scala.collection.mutable.Buffer
 import scala.jdk.CollectionConverters._
 
 @Timeout(600)
@@ -131,17 +125,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertEquals(numRecords, records.size)
   }
 
-  // Deprecated poll(timeout) not supported for consumer group protocol
-  @deprecated("poll(Duration) is the replacement", since = "2.0")
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testDeprecatedPollBlocksForAssignment(quorum: String, groupProtocol: String): Unit = {
-    val consumer = createConsumer()
-    consumer.subscribe(Set(topic).asJava)
-    consumer.poll(0)
-    assertEquals(Set(tp, tp2), consumer.assignment().asScala)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testHeadersSerializerDeserializer(quorum: String, groupProtocol: String): Unit = {
@@ -152,141 +135,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     testHeadersSerializeDeserialize(extendedSerializer, extendedDeserializer)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
-    val maxPollRecords = 2
-    val numRecords = 10000
-
-    val producer = createProducer()
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer = createConsumer()
-    consumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer, numRecords = numRecords, startingOffset = 0, maxPollRecords = maxPollRecords,
-      startingTimestamp = startingTimestamp)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMaxPollIntervalMs(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 2000.toString)
-
-    val consumer = createConsumer()
-
-    val listener = new TestConsumerReassignmentListener()
-    consumer.subscribe(List(topic).asJava, listener)
-
-    // rebalance to get the initial assignment
-    awaitRebalance(consumer, listener)
-    assertEquals(1, listener.callsToAssigned)
-    assertEquals(0, listener.callsToRevoked)
-
-    // after we extend longer than max.poll a rebalance should be triggered
-    // NOTE we need to have a relatively much larger value than max.poll to let heartbeat expired for sure
-    Thread.sleep(3000)
-
-    awaitRebalance(consumer, listener)
-    assertEquals(2, listener.callsToAssigned)
-    assertEquals(1, listener.callsToRevoked)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMaxPollIntervalMsDelayInRevocation(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 5000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 1000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
-
-    val consumer = createConsumer()
-    var commitCompleted = false
-    var committedPosition: Long = -1
-
-    val listener = new TestConsumerReassignmentListener {
-      override def onPartitionsLost(partitions: util.Collection[TopicPartition]): Unit = {}
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
-        if (!partitions.isEmpty && partitions.contains(tp)) {
-          // on the second rebalance (after we have joined the group initially), sleep longer
-          // than session timeout and then try a commit. We should still be in the group,
-          // so the commit should succeed
-          Utils.sleep(1500)
-          committedPosition = consumer.position(tp)
-          consumer.commitSync(Map(tp -> new OffsetAndMetadata(committedPosition)).asJava)
-          commitCompleted = true
-        }
-        super.onPartitionsRevoked(partitions)
-      }
-    }
-
-    consumer.subscribe(List(topic).asJava, listener)
-
-    // rebalance to get the initial assignment
-    awaitRebalance(consumer, listener)
-
-    // force a rebalance to trigger an invocation of the revocation callback while in the group
-    consumer.subscribe(List("otherTopic").asJava, listener)
-    awaitRebalance(consumer, listener)
-
-    assertEquals(0, committedPosition)
-    assertTrue(commitCompleted)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMaxPollIntervalMsDelayInAssignment(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 5000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 1000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
-
-    val consumer = createConsumer()
-    val listener = new TestConsumerReassignmentListener {
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
-        // sleep longer than the session timeout, we should still be in the group after invocation
-        Utils.sleep(1500)
-        super.onPartitionsAssigned(partitions)
-      }
-    }
-    consumer.subscribe(List(topic).asJava, listener)
-
-    // rebalance to get the initial assignment
-    awaitRebalance(consumer, listener)
-
-    // We should still be in the group after this invocation
-    ensureNoRebalance(consumer, listener)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMaxPollIntervalMsShorterThanPollTimeout(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
-
-    val consumer = createConsumer()
-    val listener = new TestConsumerReassignmentListener
-    consumer.subscribe(List(topic).asJava, listener)
-
-    // rebalance to get the initial assignment
-    awaitRebalance(consumer, listener)
-
-    val callsToAssignedAfterFirstRebalance = listener.callsToAssigned
-
-    consumer.poll(Duration.ofMillis(2000))
-
-    // If the poll poll above times out, it would trigger a rebalance.
-    // Leave some time for the rebalance to happen and check for the rebalance event.
-    consumer.poll(Duration.ofMillis(500))
-    consumer.poll(Duration.ofMillis(500))
-
-    assertEquals(callsToAssignedAfterFirstRebalance, listener.callsToAssigned)
-  }
-
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testAutoCommitOnClose(quorum: String, groupProtocol: String): Unit = {
@@ -363,155 +211,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0, startingTimestamp = startingTimestamp)
   }
 
-  /**
-   * Verifies that pattern subscription performs as expected.
-   * The pattern matches the topics 'topic' and 'tblablac', but not 'tblablak' or 'tblab1'.
-   * It is expected that the consumer is subscribed to all partitions of 'topic' and
-   * 'tblablac' after the subscription when metadata is refreshed.
-   * When a new topic 'tsomec' is added afterwards, it is expected that upon the next
-   * metadata refresh the consumer becomes subscribed to this new topic and all partitions
-   * of that topic are assigned to it.
-   */
-  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testPatternSubscription(quorum: String, groupProtocol: String): Unit = {
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp)
-
-    val topic1 = "tblablac" // matches subscribed pattern
-    createTopic(topic1, 2, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
-
-    val topic2 = "tblablak" // does not match subscribed pattern
-    createTopic(topic2, 2, brokerCount)
-    sendRecords(producer,numRecords = 1000, new TopicPartition(topic2, 0))
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic2, 1))
-
-    val topic3 = "tblab1" // does not match subscribed pattern
-    createTopic(topic3, 2, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 0))
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 1))
-
-    val consumer = createConsumer()
-    assertEquals(0, consumer.assignment().size)
-
-    val pattern = Pattern.compile("t.*c")
-    consumer.subscribe(pattern, new TestConsumerReassignmentListener)
-
-    var assignment = Set(
-      new TopicPartition(topic, 0),
-      new TopicPartition(topic, 1),
-      new TopicPartition(topic1, 0),
-      new TopicPartition(topic1, 1))
-    awaitAssignment(consumer, assignment)
-
-    val topic4 = "tsomec" // matches subscribed pattern
-    createTopic(topic4, 2, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 0))
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 1))
-
-    assignment ++= Set(
-      new TopicPartition(topic4, 0),
-      new TopicPartition(topic4, 1))
-    awaitAssignment(consumer, assignment)
-
-    consumer.unsubscribe()
-    assertEquals(0, consumer.assignment().size)
-  }
-
-  /**
-   * Verifies that a second call to pattern subscription succeeds and performs as expected.
-   * The initial subscription is to a pattern that matches two topics 'topic' and 'foo'.
-   * The second subscription is to a pattern that matches 'foo' and a new topic 'bar'.
-   * It is expected that the consumer is subscribed to all partitions of 'topic' and 'foo' after
-   * the first subscription, and to all partitions of 'foo' and 'bar' after the second.
-   * The metadata refresh interval is intentionally increased to a large enough value to guarantee
-   * that it is the subscription call that triggers a metadata refresh, and not the timeout.
-   */
-  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testSubsequentPatternSubscription(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.METADATA_MAX_AGE_CONFIG, "30000")
-    val consumer = createConsumer()
-
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords = numRecords, tp)
-
-    // the first topic ('topic')  matches first subscription pattern only
-
-    val fooTopic = "foo" // matches both subscription patterns
-    createTopic(fooTopic, 1, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(fooTopic, 0))
-
-    assertEquals(0, consumer.assignment().size)
-
-    val pattern1 = Pattern.compile(".*o.*") // only 'topic' and 'foo' match this
-    consumer.subscribe(pattern1, new TestConsumerReassignmentListener)
-
-    var assignment = Set(
-      new TopicPartition(topic, 0),
-      new TopicPartition(topic, 1),
-      new TopicPartition(fooTopic, 0))
-    awaitAssignment(consumer, assignment)
-
-    val barTopic = "bar" // matches the next subscription pattern
-    createTopic(barTopic, 1, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(barTopic, 0))
-
-    val pattern2 = Pattern.compile("...") // only 'foo' and 'bar' match this
-    consumer.subscribe(pattern2, new TestConsumerReassignmentListener)
-    assignment --= Set(
-      new TopicPartition(topic, 0),
-      new TopicPartition(topic, 1))
-    assignment ++= Set(
-      new TopicPartition(barTopic, 0))
-    awaitAssignment(consumer, assignment)
-
-    consumer.unsubscribe()
-    assertEquals(0, consumer.assignment().size)
-  }
-
-  /**
-   * Verifies that pattern unsubscription performs as expected.
-   * The pattern matches the topics 'topic' and 'tblablac'.
-   * It is expected that the consumer is subscribed to all partitions of 'topic' and
-   * 'tblablac' after the subscription when metadata is refreshed.
-   * When consumer unsubscribes from all its subscriptions, it is expected that its
-   * assignments are cleared right away.
-   */
-  // TODO: enable this test for the consumer group protocol when support for pattern subscriptions is implemented.
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testPatternUnsubscription(quorum: String, groupProtocol: String): Unit = {
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp)
-
-    val topic1 = "tblablac" // matches the subscription pattern
-    createTopic(topic1, 2, brokerCount)
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
-    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
-
-    val consumer = createConsumer()
-    assertEquals(0, consumer.assignment().size)
-
-    consumer.subscribe(Pattern.compile("t.*c"), new TestConsumerReassignmentListener)
-    val assignment = Set(
-      new TopicPartition(topic, 0),
-      new TopicPartition(topic, 1),
-      new TopicPartition(topic1, 0),
-      new TopicPartition(topic1, 1))
-    awaitAssignment(consumer, assignment)
-
-    consumer.unsubscribe()
-    assertEquals(0, consumer.assignment().size)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testCommitMetadata(quorum: String, groupProtocol: String): Unit = {
@@ -554,36 +253,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertEquals(new OffsetAndMetadata(count), consumer.committed(Set(tp).asJava).get(tp))
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testExpandingTopicSubscriptions(quorum: String, groupProtocol: String): Unit = {
-    val otherTopic = "other"
-    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
-    val consumer = createConsumer()
-    consumer.subscribe(List(topic).asJava)
-    awaitAssignment(consumer, initialAssignment)
-
-    createTopic(otherTopic, 2, brokerCount)
-    val expandedAssignment = initialAssignment ++ Set(new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
-    consumer.subscribe(List(topic, otherTopic).asJava)
-    awaitAssignment(consumer, expandedAssignment)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testShrinkingTopicSubscriptions(quorum: String, groupProtocol: String): Unit = {
-    val otherTopic = "other"
-    createTopic(otherTopic, 2, brokerCount)
-    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
-    val consumer = createConsumer()
-    consumer.subscribe(List(topic, otherTopic).asJava)
-    awaitAssignment(consumer, initialAssignment)
-
-    val shrunkenAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
-    consumer.subscribe(List(topic).asJava)
-    awaitAssignment(consumer, shrunkenAssignment)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testPartitionsFor(quorum: String, groupProtocol: String): Unit = {
@@ -721,583 +390,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 5, startingTimestamp = startingTimestamp)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchInvalidOffset(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none")
-    val consumer = createConsumer(configOverrides = this.consumerConfig)
-
-    // produce one record
-    val totalRecords = 2
-    val producer = createProducer()
-    sendRecords(producer, totalRecords, tp)
-    consumer.assign(List(tp).asJava)
-
-    // poll should fail because there is no offset reset strategy set.
-    // we fail only when resetting positions after coordinator is known, so using a long timeout.
-    assertThrows(classOf[NoOffsetForPartitionException], () => consumer.poll(Duration.ofMillis(15000)))
-
-    // seek to out of range position
-    val outOfRangePos = totalRecords + 1
-    consumer.seek(tp, outOfRangePos)
-    val e = assertThrows(classOf[OffsetOutOfRangeException], () => consumer.poll(Duration.ofMillis(20000)))
-    val outOfRangePartitions = e.offsetOutOfRangePartitions()
-    assertNotNull(outOfRangePartitions)
-    assertEquals(1, outOfRangePartitions.size)
-    assertEquals(outOfRangePos.toLong, outOfRangePartitions.get(tp))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchOutOfRangeOffsetResetConfigEarliest(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
-    // ensure no in-flight fetch request so that the offset can be reset immediately
-    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0")
-    val consumer = createConsumer(configOverrides = this.consumerConfig)
-    val totalRecords = 10L
-
-    val producer = createProducer()
-    val startingTimestamp = 0
-    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = startingTimestamp)
-    consumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = consumer, numRecords = totalRecords.toInt, startingOffset =0)
-    // seek to out of range position
-    val outOfRangePos = totalRecords + 1
-    consumer.seek(tp, outOfRangePos)
-    // assert that poll resets to the beginning position
-    consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0)
-  }
-
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchOutOfRangeOffsetResetConfigLatest(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest")
-    // ensure no in-flight fetch request so that the offset can be reset immediately
-    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0")
-    val consumer = createConsumer(configOverrides = this.consumerConfig)
-    val totalRecords = 10L
-
-    val producer = createProducer()
-    val startingTimestamp = 0
-    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = startingTimestamp)
-    consumer.assign(List(tp).asJava)
-    consumer.seek(tp, 0)
-    // consume some, but not all of the records
-    consumeAndVerifyRecords(consumer = consumer, numRecords = totalRecords.toInt/2, startingOffset = 0)
-    // seek to out of range position
-    val outOfRangePos = totalRecords + 17 // arbitrary, much higher offset
-    consumer.seek(tp, outOfRangePos)
-    // assert that poll resets to the ending position
-    assertTrue(consumer.poll(Duration.ofMillis(50)).isEmpty)
-    sendRecords(producer, totalRecords.toInt, tp, startingTimestamp = totalRecords)
-    val nextRecord = consumer.poll(Duration.ofMillis(50)).iterator().next()
-    // ensure the seek went to the last known record at the time of the previous poll
-    assertEquals(totalRecords, nextRecord.offset())
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchRecordLargerThanFetchMaxBytes(quorum: String, groupProtocol: String): Unit = {
-    val maxFetchBytes = 10 * 1024
-    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxFetchBytes.toString)
-    checkLargeRecord(maxFetchBytes + 1)
-  }
-
-  private def checkLargeRecord(producerRecordSize: Int): Unit = {
-    val consumer = createConsumer()
-
-    // produce a record that is larger than the configured fetch size
-    val record = new ProducerRecord(tp.topic(), tp.partition(), "key".getBytes,
-      new Array[Byte](producerRecordSize))
-    val producer = createProducer()
-    producer.send(record)
-
-    // consuming a record that is too large should succeed since KIP-74
-    consumer.assign(List(tp).asJava)
-    val records = consumer.poll(Duration.ofMillis(20000))
-    assertEquals(1, records.count)
-    val consumerRecord = records.iterator().next()
-    assertEquals(0L, consumerRecord.offset)
-    assertEquals(tp.topic(), consumerRecord.topic())
-    assertEquals(tp.partition(), consumerRecord.partition())
-    assertArrayEquals(record.key(), consumerRecord.key())
-    assertArrayEquals(record.value(), consumerRecord.value())
-  }
-
-  /** We should only return a large record if it's the first record in the first non-empty partition of the fetch request */
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchHonoursFetchSizeIfLargeRecordNotFirst(quorum: String, groupProtocol: String): Unit = {
-    val maxFetchBytes = 10 * 1024
-    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxFetchBytes.toString)
-    checkFetchHonoursSizeIfLargeRecordNotFirst(maxFetchBytes)
-  }
-
-  private def checkFetchHonoursSizeIfLargeRecordNotFirst(largeProducerRecordSize: Int): Unit = {
-    val consumer = createConsumer()
-
-    val smallRecord = new ProducerRecord(tp.topic(), tp.partition(), "small".getBytes,
-      "value".getBytes)
-    val largeRecord = new ProducerRecord(tp.topic(), tp.partition(), "large".getBytes,
-      new Array[Byte](largeProducerRecordSize))
-
-    val producer = createProducer()
-    producer.send(smallRecord).get
-    producer.send(largeRecord).get
-
-    // we should only get the small record in the first `poll`
-    consumer.assign(List(tp).asJava)
-    val records = consumer.poll(Duration.ofMillis(20000))
-    assertEquals(1, records.count)
-    val consumerRecord = records.iterator().next()
-    assertEquals(0L, consumerRecord.offset)
-    assertEquals(tp.topic(), consumerRecord.topic())
-    assertEquals(tp.partition(), consumerRecord.partition())
-    assertArrayEquals(smallRecord.key(), consumerRecord.key())
-    assertArrayEquals(smallRecord.value(), consumerRecord.value())
-  }
-
-  /** We should only return a large record if it's the first record in the first partition of the fetch request */
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchHonoursMaxPartitionFetchBytesIfLargeRecordNotFirst(quorum: String, groupProtocol: String): Unit = {
-    val maxPartitionFetchBytes = 10 * 1024
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, maxPartitionFetchBytes.toString)
-    checkFetchHonoursSizeIfLargeRecordNotFirst(maxPartitionFetchBytes)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testFetchRecordLargerThanMaxPartitionFetchBytes(quorum: String, groupProtocol: String): Unit = {
-    val maxPartitionFetchBytes = 10 * 1024
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, maxPartitionFetchBytes.toString)
-    checkLargeRecord(maxPartitionFetchBytes + 1)
-  }
-
-  /** Test that we consume all partitions if fetch max bytes and max.partition.fetch.bytes are low */
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testLowMaxFetchSizeForRequestAndPartition(quorum: String, groupProtocol: String): Unit = {
-    // one of the effects of this is that there will be some log reads where `0 > remaining limit bytes < message size`
-    // and we don't return the message because it's not the first message in the first non-empty partition of the fetch
-    // this behaves a little different than when remaining limit bytes is 0 and it's important to test it
-    this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, "500")
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, "100")
-
-    // Avoid a rebalance while the records are being sent (the default is 6 seconds)
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 20000.toString)
-    val consumer = createConsumer()
-
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-    val topic3 = "topic3"
-    val partitionCount = 30
-    val topics = Seq(topic1, topic2, topic3)
-    topics.foreach { topicName =>
-      createTopic(topicName, partitionCount, brokerCount)
-    }
-
-    val partitions = topics.flatMap { topic =>
-      (0 until partitionCount).map(new TopicPartition(topic, _))
-    }
-
-    assertEquals(0, consumer.assignment().size)
-
-    consumer.subscribe(List(topic1, topic2, topic3).asJava)
-
-    awaitAssignment(consumer, partitions.toSet)
-
-    val producer = createProducer()
-
-    val producerRecords = partitions.flatMap(sendRecords(producer, numRecords = partitionCount, _))
-
-    val consumerRecords = consumeRecords(consumer, producerRecords.size)
-
-    val expected = producerRecords.map { record =>
-      (record.topic, record.partition, new String(record.key), new String(record.value), record.timestamp)
-    }.toSet
-
-    val actual = consumerRecords.map { record =>
-      (record.topic, record.partition, new String(record.key), new String(record.value), record.timestamp)
-    }.toSet
-
-    assertEquals(expected, actual)
-  }
-
-
-  // Remote assignors only supported with consumer group protocol
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @CsvSource(Array(
-    "kraft+kip848, consumer"
-  ))
-  def testRemoteAssignorInvalid(quorum: String, groupProtocol: String): Unit = {
-    // 1 consumer using invalid remote assignor
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "invalid-assignor-group")
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "invalid")
-    val consumer = createConsumer()
-
-    // create two new topics, each having 2 partitions
-    val topic1 = "topic1"
-    val producer = createProducer()
-    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100)
-
-    assertEquals(0, consumer.assignment().size)
-
-    // subscribe to two topics
-    consumer.subscribe(List(topic1).asJava)
-
-    val e:UnsupportedAssignorException = assertThrows(
-      classOf[UnsupportedAssignorException],
-      () => awaitAssignment(consumer, expectedAssignment)
-    )
-
-    assertTrue(e.getMessage.startsWith("ServerAssignor invalid is not supported. " +
-      "Supported assignors: "))
-  }
-
-  // Remote assignors only supported with consumer group protocol
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @CsvSource(Array(
-    "kraft+kip848, consumer"
-  ))
-  def testRemoteAssignorRange(quorum: String, groupProtocol: String): Unit = {
-    // 1 consumer using range assignment
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "range-group")
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "range")
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "30000")
-    val consumer = createConsumer()
-
-    // create two new topics, each having 2 partitions
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-    val producer = createProducer()
-    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100) ++
-      createTopicAndSendRecords(producer, topic2, 2, 100)
-
-    assertEquals(0, consumer.assignment().size)
-
-    // subscribe to two topics
-    consumer.subscribe(List(topic1, topic2).asJava)
-    awaitAssignment(consumer, expectedAssignment)
-
-    // add one more topic with 2 partitions
-    val topic3 = "topic3"
-    val additionalAssignment = createTopicAndSendRecords(producer, topic3, 2, 100)
-
-    val newExpectedAssignment = expectedAssignment ++ additionalAssignment
-    consumer.subscribe(List(topic1, topic2, topic3).asJava)
-    awaitAssignment(consumer, newExpectedAssignment)
-
-    // remove the topic we just added
-    consumer.subscribe(List(topic1, topic2).asJava)
-    awaitAssignment(consumer, expectedAssignment)
-
-    consumer.unsubscribe()
-    assertEquals(0, consumer.assignment().size)
-  }
-
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testRoundRobinAssignment(quorum: String, groupProtocol: String): Unit = {
-    // 1 consumer using round-robin assignment
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "roundrobin-group")
-    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[RoundRobinAssignor].getName)
-    val consumer = createConsumer()
-
-    // create two new topics, each having 2 partitions
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-    val producer = createProducer()
-    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100) ++
-      createTopicAndSendRecords(producer, topic2, 2, 100)
-
-    assertEquals(0, consumer.assignment().size)
-
-    // subscribe to two topics
-    consumer.subscribe(List(topic1, topic2).asJava)
-    awaitAssignment(consumer, expectedAssignment)
-
-    // add one more topic with 2 partitions
-    val topic3 = "topic3"
-    createTopicAndSendRecords(producer, topic3, 2, 100)
-
-    val newExpectedAssignment = expectedAssignment ++ Set(new TopicPartition(topic3, 0), new TopicPartition(topic3, 1))
-    consumer.subscribe(List(topic1, topic2, topic3).asJava)
-    awaitAssignment(consumer, newExpectedAssignment)
-
-    // remove the topic we just added
-    consumer.subscribe(List(topic1, topic2).asJava)
-    awaitAssignment(consumer, expectedAssignment)
-
-    consumer.unsubscribe()
-    assertEquals(0, consumer.assignment().size)
-  }
-
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testMultiConsumerRoundRobinAssignor(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "roundrobin-group")
-    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[RoundRobinAssignor].getName)
-
-    // create two new topics, total number of partitions must be greater than number of consumers
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-    val producer = createProducer()
-    val subscriptions = createTopicAndSendRecords(producer, topic1, 5, 100) ++
-      createTopicAndSendRecords(producer, topic2, 8, 100)
-
-    // create a group of consumers, subscribe the consumers to all the topics and start polling
-    // for the topic partition assignment
-    val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(10, List(topic1, topic2), subscriptions)
-    try {
-      validateGroupAssignment(consumerPollers, subscriptions)
-
-      // add one more consumer and validate re-assignment
-      addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers,
-        List(topic1, topic2), subscriptions, "roundrobin-group")
-    } finally {
-      consumerPollers.foreach(_.shutdown())
-    }
-  }
-
-  /**
-   * This test runs the following scenario to verify sticky assignor behavior.
-   * Topics: single-topic, with random number of partitions, where #par is 10, 20, 30, 40, 50, 60, 70, 80, 90, or 100
-   * Consumers: 9 consumers subscribed to the single topic
-   * Expected initial assignment: partitions are assigned to consumers in a round robin fashion.
-   *  - (#par mod 9) consumers will get (#par / 9 + 1) partitions, and the rest get (#par / 9) partitions
-   * Then consumer #10 is added to the list (subscribing to the same single topic)
-   * Expected new assignment:
-   *  - (#par / 10) partition per consumer, where one partition from each of the early (#par mod 9) consumers
-   *    will move to consumer #10, leading to a total of (#par mod 9) partition movement
-   */
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testMultiConsumerStickyAssignor(quorum: String, groupProtocol: String): Unit = {
-
-    def reverse(m: Map[Long, Set[TopicPartition]]) =
-      m.values.toSet.flatten.map(v => (v, m.keys.filter(m(_).contains(v)).head)).toMap
-
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "sticky-group")
-    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[StickyAssignor].getName)
-
-    // create one new topic
-    val topic = "single-topic"
-    val rand = 1 + scala.util.Random.nextInt(10)
-    val producer = createProducer()
-    val partitions = createTopicAndSendRecords(producer, topic, rand * 10, 100)
-
-    // create a group of consumers, subscribe the consumers to the single topic and start polling
-    // for the topic partition assignment
-    val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(9, List(topic), partitions)
-    validateGroupAssignment(consumerPollers, partitions)
-    val prePartition2PollerId = reverse(consumerPollers.map(poller => (poller.getId, poller.consumerAssignment())).toMap)
-
-    // add one more consumer and validate re-assignment
-    addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers, List(topic), partitions, "sticky-group")
-
-    val postPartition2PollerId = reverse(consumerPollers.map(poller => (poller.getId, poller.consumerAssignment())).toMap)
-    val keys = prePartition2PollerId.keySet.union(postPartition2PollerId.keySet)
-    var changes = 0
-    keys.foreach { key =>
-      val preVal = prePartition2PollerId.get(key)
-      val postVal = postPartition2PollerId.get(key)
-      if (preVal.nonEmpty && postVal.nonEmpty) {
-        if (preVal.get != postVal.get)
-          changes += 1
-      } else
-        changes += 1
-    }
-
-    consumerPollers.foreach(_.shutdown())
-
-    assertEquals(rand, changes, "Expected only two topic partitions that have switched to other consumers.")
-  }
-
-  /**
-   * This test re-uses BaseConsumerTest's consumers.
-   * As a result, it is testing the default assignment strategy set by BaseConsumerTest
-   */
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testMultiConsumerDefaultAssignor(quorum: String, groupProtocol: String): Unit = {
-    // use consumers and topics defined in this class + one more topic
-    val producer = createProducer()
-    sendRecords(producer, numRecords = 100, tp)
-    sendRecords(producer, numRecords = 100, tp2)
-    val topic1 = "topic1"
-    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 5, 100)
-
-    // subscribe all consumers to all topics and validate the assignment
-
-    val consumersInGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
-    consumersInGroup += createConsumer()
-    consumersInGroup += createConsumer()
-
-    val consumerPollers = subscribeConsumers(consumersInGroup, List(topic, topic1))
-    try {
-      validateGroupAssignment(consumerPollers, subscriptions)
-
-      // add 2 more consumers and validate re-assignment
-      addConsumersToGroupAndWaitForGroupAssignment(2, consumersInGroup, consumerPollers, List(topic, topic1), subscriptions)
-
-      // add one more topic and validate partition re-assignment
-      val topic2 = "topic2"
-      val expandedSubscriptions = subscriptions ++ createTopicAndSendRecords(producer, topic2, 3, 100)
-      changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers, List(topic, topic1, topic2), expandedSubscriptions)
-
-      // remove the topic we just added and validate re-assignment
-      changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers, List(topic, topic1), subscriptions)
-
-    } finally {
-      consumerPollers.foreach(_.shutdown())
-    }
-  }
-
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest
-  @CsvSource(Array(
-    "org.apache.kafka.clients.consumer.CooperativeStickyAssignor,   zk",
-    "org.apache.kafka.clients.consumer.RangeAssignor,               zk",
-    "org.apache.kafka.clients.consumer.CooperativeStickyAssignor,   kraft",
-    "org.apache.kafka.clients.consumer.RangeAssignor,               kraft"
-  ))
-  def testRebalanceAndRejoin(assignmentStrategy: String, quorum: String): Unit = {
-    // create 2 consumers
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, "classic")
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "rebalance-and-rejoin-group")
-    this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, assignmentStrategy)
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer1 = createConsumer()
-    val consumer2 = createConsumer()
-
-    // create a new topic, have 2 partitions
-    val topic = "topic1"
-    val producer = createProducer()
-    val expectedAssignment = createTopicAndSendRecords(producer, topic, 2, 100)
-
-    assertEquals(0, consumer1.assignment().size)
-    assertEquals(0, consumer2.assignment().size)
-
-    val lock = new ReentrantLock()
-    var generationId1 = -1
-    var memberId1 = ""
-    val customRebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
-      }
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
-        if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
-          fail(s"Time out while awaiting for lock.")
-        }
-        try {
-          generationId1 = consumer1.groupMetadata().generationId()
-          memberId1 = consumer1.groupMetadata().memberId()
-        } finally {
-          lock.unlock()
-        }
-      }
-    }
-    val consumerPoller1 = new ConsumerAssignmentPoller(consumer1, List(topic), Set.empty, customRebalanceListener)
-    consumerPoller1.start()
-    TestUtils.waitUntilTrue(() => consumerPoller1.consumerAssignment() == expectedAssignment,
-      s"Timed out while awaiting expected assignment change to $expectedAssignment.")
-
-    // Since the consumer1 already completed the rebalance,
-    // the `onPartitionsAssigned` rebalance listener will be invoked to set the generationId and memberId
-    var stableGeneration = -1
-    var stableMemberId1 = ""
-    if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
-      fail(s"Time out while awaiting for lock.")
-    }
-    try {
-      stableGeneration = generationId1
-      stableMemberId1 = memberId1
-    } finally {
-      lock.unlock()
-    }
-
-    val consumerPoller2 = subscribeConsumerAndStartPolling(consumer2, List(topic))
-    TestUtils.waitUntilTrue(() => consumerPoller1.consumerAssignment().size == 1,
-      s"Timed out while awaiting expected assignment size change to 1.")
-    TestUtils.waitUntilTrue(() => consumerPoller2.consumerAssignment().size == 1,
-      s"Timed out while awaiting expected assignment size change to 1.")
-
-    if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
-      fail(s"Time out while awaiting for lock.")
-    }
-    try {
-      if (assignmentStrategy.equals(classOf[CooperativeStickyAssignor].getName)) {
-        // cooperative rebalance should rebalance twice before finally stable
-        assertEquals(stableGeneration + 2, generationId1)
-      } else {
-        // eager rebalance should rebalance once before finally stable
-        assertEquals(stableGeneration + 1, generationId1)
-      }
-      assertEquals(stableMemberId1, memberId1)
-    } finally {
-      lock.unlock()
-    }
-
-    consumerPoller1.shutdown()
-    consumerPoller2.shutdown()
-  }
-
-  /**
-   * This test re-uses BaseConsumerTest's consumers.
-   * As a result, it is testing the default assignment strategy set by BaseConsumerTest
-   * It tests the assignment results is expected using default assignor (i.e. Range assignor)
-   */
-  // Only the classic group protocol supports client-side assignors
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersClassicGroupProtocolOnly"))
-  def testMultiConsumerDefaultAssignorAndVerifyAssignment(quorum: String, groupProtocol: String): Unit = {
-    // create two new topics, each having 3 partitions
-    val topic1 = "topic1"
-    val topic2 = "topic2"
-
-    createTopic(topic1, 3)
-    createTopic(topic2, 3)
-
-    val consumersInGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
-    consumersInGroup += createConsumer()
-    consumersInGroup += createConsumer()
-
-    val tp1_0 = new TopicPartition(topic1, 0)
-    val tp1_1 = new TopicPartition(topic1, 1)
-    val tp1_2 = new TopicPartition(topic1, 2)
-    val tp2_0 = new TopicPartition(topic2, 0)
-    val tp2_1 = new TopicPartition(topic2, 1)
-    val tp2_2 = new TopicPartition(topic2, 2)
-
-    val subscriptions = Set(tp1_0, tp1_1, tp1_2, tp2_0, tp2_1, tp2_2)
-    val consumerPollers = subscribeConsumers(consumersInGroup, List(topic1, topic2))
-
-    val expectedAssignment = Buffer(Set(tp1_0, tp1_1, tp2_0, tp2_1), Set(tp1_2, tp2_2))
-
-    try {
-      validateGroupAssignment(consumerPollers, subscriptions, expectedAssignment = expectedAssignment)
-    } finally {
-      consumerPollers.foreach(_.shutdown())
-    }
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMultiConsumerSessionTimeoutOnStopPolling(quorum: String, groupProtocol: String): Unit = {
-    runMultiConsumerSessionTimeoutTest(false)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testMultiConsumerSessionTimeoutOnClose(quorum: String, groupProtocol: String): Unit = {
-    runMultiConsumerSessionTimeoutTest(true)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testInterceptors(quorum: String, groupProtocol: String): Unit = {
@@ -1512,23 +604,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertEquals(2, topics.get(topic3).size)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testUnsubscribeTopic(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "100") // timeout quickly to avoid slow test
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "30")
-    val consumer = createConsumer()
-
-    val listener = new TestConsumerReassignmentListener()
-    consumer.subscribe(List(topic).asJava, listener)
-
-    // the initial subscription should cause a callback execution
-    awaitRebalance(consumer, listener)
-
-    consumer.subscribe(List[String]().asJava)
-    assertEquals(0, consumer.assignment.size())
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testPauseStateNotPreservedByRebalance(quorum: String, groupProtocol: String): Unit = {
@@ -1786,53 +861,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     assertNotNull(fetchLag)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testPerPartitionLeadWithMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
-    val numMessages = 1000
-    val maxPollRecords = 10
-    val producer = createProducer()
-    sendRecords(producer, numMessages, tp)
-
-    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
-    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
-    consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer = createConsumer()
-    consumer.assign(List(tp).asJava)
-    awaitNonEmptyRecords(consumer, tp)
-
-    val tags = new util.HashMap[String, String]()
-    tags.put("client-id", "testPerPartitionLeadWithMaxPollRecords")
-    tags.put("topic", tp.topic())
-    tags.put("partition", String.valueOf(tp.partition()))
-    val lead = consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags))
-    assertEquals(maxPollRecords, lead.metricValue().asInstanceOf[Double], s"The lead should be $maxPollRecords")
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testPerPartitionLagWithMaxPollRecords(quorum: String, groupProtocol: String): Unit = {
-    val numMessages = 1000
-    val maxPollRecords = 10
-    val producer = createProducer()
-    sendRecords(producer, numMessages, tp)
-
-    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
-    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
-    consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer = createConsumer()
-    consumer.assign(List(tp).asJava)
-    val records = awaitNonEmptyRecords(consumer, tp)
-
-    val tags = new util.HashMap[String, String]()
-    tags.put("client-id", "testPerPartitionLagWithMaxPollRecords")
-    tags.put("topic", tp.topic())
-    tags.put("partition", String.valueOf(tp.partition()))
-    val lag = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags))
-
-    assertEquals(numMessages - records.count, lag.metricValue.asInstanceOf[Double], epsilon, s"The lag should be ${numMessages - records.count}")
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testQuotaMetricsNotCreatedIfNoQuotasConfigured(quorum: String, groupProtocol: String): Unit = {
@@ -1865,100 +893,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     brokers.foreach(assertNoMetric(_, "throttle-time", QuotaType.Request, consumerClientId))
   }
 
-  def runMultiConsumerSessionTimeoutTest(closeConsumer: Boolean): Unit = {
-    // use consumers defined in this class plus one additional consumer
-    // Use topic defined in this class + one additional topic
-    val producer = createProducer()
-    sendRecords(producer, numRecords = 100, tp)
-    sendRecords(producer, numRecords = 100, tp2)
-    val topic1 = "topic1"
-    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 6, 100)
-
-    // first subscribe consumers that are defined in this class
-    val consumerPollers = Buffer[ConsumerAssignmentPoller]()
-    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
-    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
-
-    // create one more consumer and add it to the group; we will timeout this consumer
-    val timeoutConsumer = createConsumer()
-    val timeoutPoller = subscribeConsumerAndStartPolling(timeoutConsumer, List(topic, topic1))
-    consumerPollers += timeoutPoller
-
-    // validate the initial assignment
-    validateGroupAssignment(consumerPollers, subscriptions)
-
-    // stop polling and close one of the consumers, should trigger partition re-assignment among alive consumers
-    timeoutPoller.shutdown()
-    consumerPollers -= timeoutPoller
-    if (closeConsumer)
-      timeoutConsumer.close()
-
-    validateGroupAssignment(consumerPollers, subscriptions,
-      Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after one consumer left"), 3 * groupMaxSessionTimeoutMs)
-
-    // done with pollers and consumers
-    for (poller <- consumerPollers)
-      poller.shutdown()
-  }
-
-  /**
-    * Creates consumer pollers corresponding to a given consumer group, one per consumer; subscribes consumers to
-    * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
-    *
-    * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
-    *
-    * @param consumerGroup consumer group
-    * @param topicsToSubscribe topics to which consumers will subscribe to
-    * @return collection of consumer pollers
-    */
-  def subscribeConsumers(consumerGroup: mutable.Buffer[Consumer[Array[Byte], Array[Byte]]],
-                         topicsToSubscribe: List[String]): mutable.Buffer[ConsumerAssignmentPoller] = {
-    val consumerPollers = mutable.Buffer[ConsumerAssignmentPoller]()
-    for (consumer <- consumerGroup)
-      consumerPollers += subscribeConsumerAndStartPolling(consumer, topicsToSubscribe)
-    consumerPollers
-  }
-
-  /**
-    * Creates 'consumerCount' consumers and consumer pollers, one per consumer; subscribes consumers to
-    * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
-    *
-    * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
-    *
-    * @param consumerCount number of consumers to create
-    * @param topicsToSubscribe topics to which consumers will subscribe to
-    * @param subscriptions set of all topic partitions
-    * @return collection of created consumers and collection of corresponding consumer pollers
-    */
-  def createConsumerGroupAndWaitForAssignment(consumerCount: Int,
-                                              topicsToSubscribe: List[String],
-                                              subscriptions: Set[TopicPartition]): (Buffer[Consumer[Array[Byte], Array[Byte]]], Buffer[ConsumerAssignmentPoller]) = {
-    assertTrue(consumerCount <= subscriptions.size)
-    val consumerGroup = Buffer[Consumer[Array[Byte], Array[Byte]]]()
-    for (_ <- 0 until consumerCount)
-      consumerGroup += createConsumer()
-
-    // create consumer pollers, wait for assignment and validate it
-    val consumerPollers = subscribeConsumers(consumerGroup, topicsToSubscribe)
-    (consumerGroup, consumerPollers)
-  }
-
-  def changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers: Buffer[ConsumerAssignmentPoller],
-                                                           topicsToSubscribe: List[String],
-                                                           subscriptions: Set[TopicPartition]): Unit = {
-    for (poller <- consumerPollers)
-      poller.subscribe(topicsToSubscribe)
-
-    // since subscribe call to poller does not actually call consumer subscribe right away, wait
-    // until subscribe is called on all consumers
-    TestUtils.waitUntilTrue(() => {
-      consumerPollers.forall { poller => poller.isSubscribeRequestProcessed }
-    }, s"Failed to call subscribe on all consumers in the group for subscription $subscriptions", 1000L)
-
-    validateGroupAssignment(consumerPollers, subscriptions,
-      Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after we changed subscription"))
-  }
-
   def changeConsumerSubscriptionAndValidateAssignment[K, V](consumer: Consumer[K, V],
                                                             topicsToSubscribe: List[String],
                                                             expectedAssignment: Set[TopicPartition],
@@ -1967,21 +901,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     awaitAssignment(consumer, expectedAssignment)
   }
 
-  private def awaitNonEmptyRecords[K, V](consumer: Consumer[K, V], partition: TopicPartition): ConsumerRecords[K, V] = {
-    TestUtils.pollRecordsUntilTrue(consumer, (polledRecords: ConsumerRecords[K, V]) => {
-      if (polledRecords.records(partition).asScala.nonEmpty)
-        return polledRecords
-      false
-    }, s"Consumer did not consume any messages for partition $partition before timeout.")
-    throw new IllegalStateException("Should have timed out before reaching here")
-  }
-
-  private def awaitAssignment(consumer: Consumer[_, _], expectedAssignment: Set[TopicPartition]): Unit = {
-    TestUtils.pollUntilTrue(consumer, () => consumer.assignment() == expectedAssignment.asJava,
-      s"Timed out while awaiting expected assignment $expectedAssignment. " +
-        s"The current assignment is ${consumer.assignment()}")
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testConsumingWithNullGroupId(quorum: String, groupProtocol: String): Unit = {
@@ -2149,211 +1068,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumer2.close()
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndCommitAsyncNotCommitted(quorum:String, groupProtocol: String): Unit = {
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props)
-    val producer = createProducer()
-    val numRecords = 10000
-    val startingTimestamp = System.currentTimeMillis()
-    val cb = new CountConsumerCommitCallback
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-    consumer.assign(List(tp).asJava)
-    consumer.commitAsync(cb)
-    TestUtils.pollUntilTrue(consumer, () => cb.successCount >= 1 || cb.lastError.isDefined,
-      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
-    val committedOffset = consumer.committed(Set(tp).asJava)
-    assertNotNull(committedOffset)
-    // No valid fetch position due to the absence of consumer.poll; and therefore no offset was committed to
-    // tp. The committed offset should be null. This is intentional.
-    assertNull(committedOffset.get(tp))
-    assertTrue(consumer.assignment.contains(tp))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndCommitSyncNotCommitted(quorum:String, groupProtocol: String): Unit = {
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props)
-    val producer = createProducer()
-    val numRecords = 10000
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-    consumer.assign(List(tp).asJava)
-    consumer.commitSync()
-    val committedOffset = consumer.committed(Set(tp).asJava)
-    assertNotNull(committedOffset)
-    // No valid fetch position due to the absence of consumer.poll; and therefore no offset was committed to
-    // tp. The committed offset should be null. This is intentional.
-    assertNull(committedOffset.get(tp))
-    assertTrue(consumer.assignment.contains(tp))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndCommitSyncAllConsumed(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 10000
-
-    val producer = createProducer()
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props)
-    consumer.assign(List(tp).asJava)
-    consumer.seek(tp, 0)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
-
-    consumer.commitSync()
-    val committedOffset = consumer.committed(Set(tp).asJava)
-    assertNotNull(committedOffset)
-    assertNotNull(committedOffset.get(tp))
-    assertEquals(numRecords, committedOffset.get(tp).offset())
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndConsume(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 10
-
-    val producer = createProducer()
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props,
-      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
-    consumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
-
-    assertEquals(numRecords, consumer.position(tp))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndConsumeSkippingPosition(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 10
-
-    val producer = createProducer()
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props,
-      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
-    consumer.assign(List(tp).asJava)
-    val offset = 1
-    consumer.seek(tp, offset)
-    consumeAndVerifyRecords(consumer = consumer, numRecords - offset, startingOffset = offset,
-      startingKeyAndValueIndex = offset, startingTimestamp = startingTimestamp + offset)
-
-    assertEquals(numRecords, consumer.position(tp))
-  }
-
-  // partitionsFor not implemented in consumer group protocol and this test requires ZK also
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @CsvSource(Array(
-    "zk, classic"
-  ))
-  def testAssignAndConsumeWithLeaderChangeValidatingPositions(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 10
-    val producer = createProducer()
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props,
-      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
-    consumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0, startingTimestamp = startingTimestamp)
-
-    // Force leader epoch change to trigger position validation
-    var parts: mutable.Buffer[PartitionInfo] = null
-    while (parts == null)
-      parts = consumer.partitionsFor(tp.topic()).asScala
-    val leader = parts.head.leader().id()
-    this.servers(leader).shutdown()
-    this.servers(leader).startup()
-
-    // Consume after leader change
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 10,
-      startingTimestamp = startingTimestamp)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndFetchCommittedOffsets(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 100
-    val startingTimestamp = System.currentTimeMillis()
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props)
-    consumer.assign(List(tp).asJava)
-    // First consumer consumes and commits offsets
-    consumer.seek(tp, 0)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0,
-      startingTimestamp = startingTimestamp)
-    consumer.commitSync()
-    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
-    // We should see the committed offsets from another consumer
-    val anotherConsumer = createConsumer(configOverrides = props)
-    anotherConsumer.assign(List(tp).asJava)
-    assertEquals(numRecords, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndConsumeFromCommittedOffsets(quorum:String, groupProtocol: String): Unit = {
-    val producer = createProducer()
-    val numRecords = 100
-    val startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords = numRecords, tp, startingTimestamp = startingTimestamp)
-
-    // Commit offset with first consumer
-    val props = new Properties()
-    props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "group1")
-    val consumer = createConsumer(configOverrides = props)
-    consumer.assign(List(tp).asJava)
-    val offset = 10
-    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(offset)))
-      .asJava)
-    assertEquals(offset, consumer.committed(Set(tp).asJava).get(tp).offset)
-    consumer.close()
-
-    // Consume from committed offsets with another consumer in same group
-    val anotherConsumer = createConsumer(configOverrides = props)
-    assertEquals(offset, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
-    anotherConsumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = anotherConsumer, numRecords - offset,
-      startingOffset = offset, startingKeyAndValueIndex = offset,
-      startingTimestamp = startingTimestamp + offset)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAssignAndRetrievingCommittedOffsetsMultipleTimes(quorum:String, groupProtocol: String): Unit = {
-    val numRecords = 100
-    val startingTimestamp = System.currentTimeMillis()
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp, startingTimestamp = startingTimestamp)
-
-    val props = new Properties()
-    val consumer = createConsumer(configOverrides = props)
-    consumer.assign(List(tp).asJava)
-
-    // Consume and commit offsets
-    consumer.seek(tp, 0)
-    consumeAndVerifyRecords(consumer = consumer, numRecords, startingOffset = 0,
-      startingTimestamp = startingTimestamp)
-    consumer.commitSync()
-
-    // Check committed offsets twice with same consumer
-    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(numRecords, consumer.committed(Set(tp).asJava).get(tp).offset)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testSubscribeAndCommitSync(quorum: String, groupProtocol: String): Unit = {

From f8ce7feebcede6c6da93c649413a64695bc8d4c1 Mon Sep 17 00:00:00 2001
From: Igor Soarez <i@soarez.me>
Date: Mon, 25 Mar 2024 17:31:19 +0000
Subject: [PATCH 201/258] KAFKA-15950: Serialize heartbeat requests (#14903)

In between HeartbeatRequest being sent and the response being handled,
i.e. while a HeartbeatRequest is in flight, an extra request may be
immediately scheduled if propagateDirectoryFailure, setReadyToUnfence,
or beginControlledShutdown is called.

To prevent the extra request, we can avoid the extra requests by checking
whether a request is in flight, and delay the scheduling if necessary.

Some of the tests in BrokerLifecycleManagerTest are also improved to
remove race conditions and reduce flakiness.

Reviewers: Colin McCabe <colin@cmccabe.xyz>, Ron Dagostino <rdagostino@confluent.io>, Jun Rao <junrao@gmail.com>
---
 .../kafka/server/BrokerLifecycleManager.scala | 208 ++++++++++--------
 .../server/BrokerLifecycleManagerTest.scala   |  76 ++++---
 .../apache/kafka/queue/KafkaEventQueue.java   |  33 +++
 3 files changed, 197 insertions(+), 120 deletions(-)

diff --git a/core/src/main/scala/kafka/server/BrokerLifecycleManager.scala b/core/src/main/scala/kafka/server/BrokerLifecycleManager.scala
index 10c4ba05ce4d0..1a10c6f659574 100644
--- a/core/src/main/scala/kafka/server/BrokerLifecycleManager.scala
+++ b/core/src/main/scala/kafka/server/BrokerLifecycleManager.scala
@@ -23,7 +23,7 @@ import kafka.utils.Logging
 import org.apache.kafka.clients.ClientResponse
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.message.BrokerRegistrationRequestData.ListenerCollection
-import org.apache.kafka.common.message.{BrokerHeartbeatRequestData, BrokerHeartbeatResponseData, BrokerRegistrationRequestData, BrokerRegistrationResponseData}
+import org.apache.kafka.common.message.{BrokerHeartbeatRequestData, BrokerRegistrationRequestData}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{BrokerHeartbeatRequest, BrokerHeartbeatResponse, BrokerRegistrationRequest, BrokerRegistrationResponse}
 import org.apache.kafka.metadata.{BrokerState, VersionRange}
@@ -166,6 +166,19 @@ class BrokerLifecycleManager(
    */
   private var registered = false
 
+  /**
+   * True if a request has been sent and a response or timeout has not yet been processed.
+   * This variable can only be read or written from the event queue thread.
+   */
+  private var communicationInFlight = false
+
+  /**
+   * True if we should schedule the next communication immediately. This is used to delay
+   * an immediate scheduling of a communication event if one is already in flight.
+   * This variable can only be read or written from the event queue thread.
+   */
+  private var nextSchedulingShouldBeImmediate = false
+
   /**
    * True if the initial registration succeeded.  This variable can only be read or
    * written from the event queue thread.
@@ -377,10 +390,30 @@ class BrokerLifecycleManager(
     }
     _channelManager.sendRequest(new BrokerRegistrationRequest.Builder(data),
       new BrokerRegistrationResponseHandler())
+    communicationInFlight = true
   }
 
+  // the response handler is not invoked from the event handler thread,
+  // so it is not safe to update state here, instead, schedule an event
+  // to continue handling the response on the event handler thread
   private class BrokerRegistrationResponseHandler extends ControllerRequestCompletionHandler {
     override def onComplete(response: ClientResponse): Unit = {
+      eventQueue.prepend(new BrokerRegistrationResponseEvent(response, false))
+    }
+
+    override def onTimeout(): Unit = {
+      info(s"Unable to register the broker because the RPC got timed out before it could be sent.")
+      eventQueue.prepend(new BrokerRegistrationResponseEvent(null, true))
+    }
+  }
+
+  private class BrokerRegistrationResponseEvent(response: ClientResponse, timedOut: Boolean) extends EventQueue.Event {
+    override def run(): Unit = {
+      communicationInFlight = false
+      if (timedOut) {
+        scheduleNextCommunicationAfterFailure()
+        return
+      }
       if (response.authenticationException() != null) {
         error(s"Unable to register broker $nodeId because of an authentication exception.",
           response.authenticationException())
@@ -400,10 +433,12 @@ class BrokerLifecycleManager(
         val message = response.responseBody().asInstanceOf[BrokerRegistrationResponse]
         val errorCode = Errors.forCode(message.data().errorCode())
         if (errorCode == Errors.NONE) {
-          // this response handler is not invoked from the event handler thread,
-          // and processing a successful registration response requires updating
-          // state, so to continue we need to schedule an event
-          eventQueue.prepend(new BrokerRegistrationResponseEvent(message.data()))
+          failedAttempts = 0
+          _brokerEpoch = message.data().brokerEpoch()
+          registered = true
+          initialRegistrationSucceeded = true
+          info(s"Successfully registered broker $nodeId with broker epoch ${_brokerEpoch}")
+          scheduleNextCommunicationImmediately() // Immediately send a heartbeat
         } else {
           info(s"Unable to register broker $nodeId because the controller returned " +
             s"error $errorCode")
@@ -411,22 +446,6 @@ class BrokerLifecycleManager(
         }
       }
     }
-
-    override def onTimeout(): Unit = {
-      info(s"Unable to register the broker because the RPC got timed out before it could be sent.")
-      scheduleNextCommunicationAfterFailure()
-    }
-  }
-
-  private class BrokerRegistrationResponseEvent(response: BrokerRegistrationResponseData) extends EventQueue.Event {
-    override def run(): Unit = {
-      failedAttempts = 0
-      _brokerEpoch = response.brokerEpoch()
-      registered = true
-      initialRegistrationSucceeded = true
-      info(s"Successfully registered broker $nodeId with broker epoch ${_brokerEpoch}")
-      scheduleNextCommunicationImmediately() // Immediately send a heartbeat
-    }
   }
 
   private def sendBrokerHeartbeat(): Unit = {
@@ -443,10 +462,30 @@ class BrokerLifecycleManager(
     }
     val handler = new BrokerHeartbeatResponseHandler()
     _channelManager.sendRequest(new BrokerHeartbeatRequest.Builder(data), handler)
+    communicationInFlight = true
   }
 
+  // the response handler is not invoked from the event handler thread,
+  // so it is not safe to update state here, instead, schedule an event
+  // to continue handling the response on the event handler thread
   private class BrokerHeartbeatResponseHandler extends ControllerRequestCompletionHandler {
     override def onComplete(response: ClientResponse): Unit = {
+      eventQueue.prepend(new BrokerHeartbeatResponseEvent(response, false))
+    }
+
+    override def onTimeout(): Unit = {
+      info("Unable to send a heartbeat because the RPC got timed out before it could be sent.")
+      eventQueue.prepend(new BrokerHeartbeatResponseEvent(null, true))
+    }
+  }
+
+  private class BrokerHeartbeatResponseEvent(response: ClientResponse, timedOut: Boolean) extends EventQueue.Event {
+    override def run(): Unit = {
+      communicationInFlight = false
+      if (timedOut) {
+        scheduleNextCommunicationAfterFailure()
+        return
+      }
       if (response.authenticationException() != null) {
         error(s"Unable to send broker heartbeat for $nodeId because of an " +
           "authentication exception.", response.authenticationException())
@@ -466,82 +505,72 @@ class BrokerLifecycleManager(
         val message = response.responseBody().asInstanceOf[BrokerHeartbeatResponse]
         val errorCode = Errors.forCode(message.data().errorCode())
         if (errorCode == Errors.NONE) {
-          // this response handler is not invoked from the event handler thread,
-          // and processing a successful heartbeat response requires updating
-          // state, so to continue we need to schedule an event
-          eventQueue.prepend(new BrokerHeartbeatResponseEvent(message.data()))
+          val responseData = message.data()
+          failedAttempts = 0
+          _state match {
+            case BrokerState.STARTING =>
+              if (responseData.isCaughtUp) {
+                info(s"The broker has caught up. Transitioning from STARTING to RECOVERY.")
+                _state = BrokerState.RECOVERY
+                initialCatchUpFuture.complete(null)
+              } else {
+                debug(s"The broker is STARTING. Still waiting to catch up with cluster metadata.")
+              }
+              // Schedule the heartbeat after only 10 ms so that in the case where
+              // there is no recovery work to be done, we start up a bit quicker.
+              scheduleNextCommunication(NANOSECONDS.convert(10, MILLISECONDS))
+            case BrokerState.RECOVERY =>
+              if (!responseData.isFenced) {
+                info(s"The broker has been unfenced. Transitioning from RECOVERY to RUNNING.")
+                initialUnfenceFuture.complete(null)
+                _state = BrokerState.RUNNING
+              } else {
+                info(s"The broker is in RECOVERY.")
+              }
+              scheduleNextCommunicationAfterSuccess()
+            case BrokerState.RUNNING =>
+              debug(s"The broker is RUNNING. Processing heartbeat response.")
+              scheduleNextCommunicationAfterSuccess()
+            case BrokerState.PENDING_CONTROLLED_SHUTDOWN =>
+              if (!responseData.shouldShutDown()) {
+                info(s"The broker is in PENDING_CONTROLLED_SHUTDOWN state, still waiting " +
+                  "for the active controller.")
+                if (!gotControlledShutdownResponse) {
+                  // If this is the first pending controlled shutdown response we got,
+                  // schedule our next heartbeat a little bit sooner than we usually would.
+                  // In the case where controlled shutdown completes quickly, this will
+                  // speed things up a little bit.
+                  scheduleNextCommunication(NANOSECONDS.convert(50, MILLISECONDS))
+                } else {
+                  scheduleNextCommunicationAfterSuccess()
+                }
+              } else {
+                info(s"The controller has asked us to exit controlled shutdown.")
+                beginShutdown()
+              }
+              gotControlledShutdownResponse = true
+            case BrokerState.SHUTTING_DOWN =>
+              info(s"The broker is SHUTTING_DOWN. Ignoring heartbeat response.")
+            case _ =>
+              error(s"Unexpected broker state ${_state}")
+              scheduleNextCommunicationAfterSuccess()
+          }
         } else {
           warn(s"Broker $nodeId sent a heartbeat request but received error $errorCode.")
           scheduleNextCommunicationAfterFailure()
         }
       }
     }
-
-    override def onTimeout(): Unit = {
-      info("Unable to send a heartbeat because the RPC got timed out before it could be sent.")
-      scheduleNextCommunicationAfterFailure()
-    }
   }
 
-  private class BrokerHeartbeatResponseEvent(response: BrokerHeartbeatResponseData) extends EventQueue.Event {
-    override def run(): Unit = {
-      failedAttempts = 0
-      _state match {
-        case BrokerState.STARTING =>
-          if (response.isCaughtUp) {
-            info(s"The broker has caught up. Transitioning from STARTING to RECOVERY.")
-            _state = BrokerState.RECOVERY
-            initialCatchUpFuture.complete(null)
-          } else {
-            debug(s"The broker is STARTING. Still waiting to catch up with cluster metadata.")
-          }
-          // Schedule the heartbeat after only 10 ms so that in the case where
-          // there is no recovery work to be done, we start up a bit quicker.
-          scheduleNextCommunication(NANOSECONDS.convert(10, MILLISECONDS))
-        case BrokerState.RECOVERY =>
-          if (!response.isFenced) {
-            info(s"The broker has been unfenced. Transitioning from RECOVERY to RUNNING.")
-            initialUnfenceFuture.complete(null)
-            _state = BrokerState.RUNNING
-          } else {
-            info(s"The broker is in RECOVERY.")
-          }
-          scheduleNextCommunicationAfterSuccess()
-        case BrokerState.RUNNING =>
-          debug(s"The broker is RUNNING. Processing heartbeat response.")
-          scheduleNextCommunicationAfterSuccess()
-        case BrokerState.PENDING_CONTROLLED_SHUTDOWN =>
-          if (!response.shouldShutDown()) {
-            info(s"The broker is in PENDING_CONTROLLED_SHUTDOWN state, still waiting " +
-              "for the active controller.")
-            if (!gotControlledShutdownResponse) {
-              // If this is the first pending controlled shutdown response we got,
-              // schedule our next heartbeat a little bit sooner than we usually would.
-              // In the case where controlled shutdown completes quickly, this will
-              // speed things up a little bit.
-              scheduleNextCommunication(NANOSECONDS.convert(50, MILLISECONDS))
-            } else {
-              scheduleNextCommunicationAfterSuccess()
-            }
-          } else {
-            info(s"The controller has asked us to exit controlled shutdown.")
-            beginShutdown()
-          }
-          gotControlledShutdownResponse = true
-        case BrokerState.SHUTTING_DOWN =>
-          info(s"The broker is SHUTTING_DOWN. Ignoring heartbeat response.")
-        case _ =>
-          error(s"Unexpected broker state ${_state}")
-          scheduleNextCommunicationAfterSuccess()
-      }
-    }
+  private def scheduleNextCommunicationImmediately(): Unit = {
+      scheduleNextCommunication(0)
   }
 
-  private def scheduleNextCommunicationImmediately(): Unit = scheduleNextCommunication(0)
-
   private def scheduleNextCommunicationAfterFailure(): Unit = {
     val delayMs = resendExponentialBackoff.backoff(failedAttempts)
     failedAttempts = failedAttempts + 1
+    nextSchedulingShouldBeImmediate = false // never immediately reschedule after a failure
     scheduleNextCommunication(NANOSECONDS.convert(delayMs, MILLISECONDS))
   }
 
@@ -551,9 +580,11 @@ class BrokerLifecycleManager(
   }
 
   private def scheduleNextCommunication(intervalNs: Long): Unit = {
-    trace(s"Scheduling next communication at ${MILLISECONDS.convert(intervalNs, NANOSECONDS)} " +
+    val adjustedIntervalNs = if (nextSchedulingShouldBeImmediate) 0 else intervalNs
+    nextSchedulingShouldBeImmediate = false
+    trace(s"Scheduling next communication at ${MILLISECONDS.convert(adjustedIntervalNs, NANOSECONDS)} " +
       "ms from now.")
-    val deadlineNs = time.nanoseconds() + intervalNs
+    val deadlineNs = time.nanoseconds() + adjustedIntervalNs
     eventQueue.scheduleDeferred("communication",
       new DeadlineFunction(deadlineNs),
       new CommunicationEvent())
@@ -570,7 +601,10 @@ class BrokerLifecycleManager(
 
   private class CommunicationEvent extends EventQueue.Event {
     override def run(): Unit = {
-      if (registered) {
+      if (communicationInFlight) {
+        trace("Delaying communication because there is already one in flight.")
+        nextSchedulingShouldBeImmediate = true
+      } else if (registered) {
         sendBrokerHeartbeat()
       } else {
         sendBrokerRegistration()
diff --git a/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala b/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala
index d7925ef8b56ab..a36813986e8d9 100644
--- a/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala
@@ -197,11 +197,18 @@ class BrokerLifecycleManagerTest {
     result
   }
 
-  def poll[T](context: RegistrationTestContext, manager: BrokerLifecycleManager, future: Future[T]): T = {
-    while (!future.isDone || context.mockClient.hasInFlightRequests) {
-      context.poll()
+  def poll[T](ctx: RegistrationTestContext, manager: BrokerLifecycleManager, future: Future[T]): T = {
+    while (ctx.mockChannelManager.unsentQueue.isEmpty) {
+      // Cancel a potential timeout event, so it doesn't get activated if we advance the clock too much
+      manager.eventQueue.cancelDeferred("initialRegistrationTimeout")
+
+      // If the manager is idling until scheduled events we need to advance the clock
+      if (manager.eventQueue.firstDeferredIfIdling().isPresent)
+        ctx.time.sleep(5)
       manager.eventQueue.wakeup()
-      context.time.sleep(5)
+    }
+    while (!future.isDone) {
+      ctx.poll()
     }
     future.get
   }
@@ -219,15 +226,16 @@ class BrokerLifecycleManagerTest {
       Collections.emptyMap(), OptionalLong.empty())
     poll(ctx, manager, registration)
 
+    def nextHeartbeatDirs(): Set[String] =
+      poll(ctx, manager, prepareResponse[BrokerHeartbeatRequest](ctx, new BrokerHeartbeatResponse(new BrokerHeartbeatResponseData())))
+        .data().offlineLogDirs().asScala.map(_.toString).toSet
+    assertEquals(Set.empty, nextHeartbeatDirs())
     manager.propagateDirectoryFailure(Uuid.fromString("h3sC4Yk-Q9-fd0ntJTocCA"))
+    assertEquals(Set("h3sC4Yk-Q9-fd0ntJTocCA"), nextHeartbeatDirs())
     manager.propagateDirectoryFailure(Uuid.fromString("ej8Q9_d2Ri6FXNiTxKFiow"))
+    assertEquals(Set("h3sC4Yk-Q9-fd0ntJTocCA", "ej8Q9_d2Ri6FXNiTxKFiow"), nextHeartbeatDirs())
     manager.propagateDirectoryFailure(Uuid.fromString("1iF76HVNRPqC7Y4r6647eg"))
-    val latestHeartbeat = Seq.fill(10)(
-      prepareResponse[BrokerHeartbeatRequest](ctx, new BrokerHeartbeatResponse(new BrokerHeartbeatResponseData()))
-    ).map(poll(ctx, manager, _)).last
-    assertEquals(
-      Set("h3sC4Yk-Q9-fd0ntJTocCA", "ej8Q9_d2Ri6FXNiTxKFiow", "1iF76HVNRPqC7Y4r6647eg").map(Uuid.fromString),
-      latestHeartbeat.data().offlineLogDirs().asScala.toSet)
+    assertEquals(Set("h3sC4Yk-Q9-fd0ntJTocCA", "ej8Q9_d2Ri6FXNiTxKFiow", "1iF76HVNRPqC7Y4r6647eg"), nextHeartbeatDirs())
     manager.close()
   }
 
@@ -254,33 +262,35 @@ class BrokerLifecycleManagerTest {
 
   @Test
   def testKraftJBODMetadataVersionUpdateEvent(): Unit = {
-    val context = new RegistrationTestContext(configProperties)
-    val manager = new BrokerLifecycleManager(context.config, context.time, "successful-registration-", isZkBroker = false, Set(Uuid.fromString("gCpDJgRlS2CBCpxoP2VMsQ")))
+    val ctx = new RegistrationTestContext(configProperties)
+    val manager = new BrokerLifecycleManager(ctx.config, ctx.time, "jbod-metadata-version-update", isZkBroker = false, Set(Uuid.fromString("gCpDJgRlS2CBCpxoP2VMsQ")))
     val controllerNode = new Node(3000, "localhost", 8021)
-    context.controllerNodeProvider.node.set(controllerNode)
-    manager.start(() => context.highestMetadataOffset.get(),
-      context.mockChannelManager, context.clusterId, context.advertisedListeners,
+    ctx.controllerNodeProvider.node.set(controllerNode)
+
+    manager.start(() => ctx.highestMetadataOffset.get(),
+      ctx.mockChannelManager, ctx.clusterId, ctx.advertisedListeners,
       Collections.emptyMap(), OptionalLong.of(10L))
-    TestUtils.retry(60000) {
-      assertEquals(1, context.mockChannelManager.unsentQueue.size)
-      assertEquals(10L, context.mockChannelManager.unsentQueue.getFirst.request.build().asInstanceOf[BrokerRegistrationRequest].data().previousBrokerEpoch())
-    }
-    context.mockClient.prepareResponseFrom(new BrokerRegistrationResponse(
-      new BrokerRegistrationResponseData().setBrokerEpoch(1000)), controllerNode)
-    TestUtils.retry(10000) {
-      context.poll()
-      assertEquals(1000L, manager.brokerEpoch)
-    }
 
+    def doPoll[T<:AbstractRequest](response: AbstractResponse) = poll(ctx, manager, prepareResponse[T](ctx, response))
+    def nextHeartbeatRequest() = doPoll[AbstractRequest](new BrokerHeartbeatResponse(new BrokerHeartbeatResponseData()))
+    def nextRegistrationRequest(epoch: Long) =
+      doPoll[BrokerRegistrationRequest](new BrokerRegistrationResponse(new BrokerRegistrationResponseData().setBrokerEpoch(epoch)))
+
+    // Broker registers and response sets epoch to 1000L
+    assertEquals(10L, nextRegistrationRequest(1000L).data().previousBrokerEpoch())
+
+    nextHeartbeatRequest() // poll for next request as way to synchronize with the new value into brokerEpoch
+    assertEquals(1000L, manager.brokerEpoch)
+
+    // Trigger JBOD MV update
     manager.handleKraftJBODMetadataVersionUpdate()
-    context.mockClient.prepareResponseFrom(new BrokerRegistrationResponse(
-      new BrokerRegistrationResponseData().setBrokerEpoch(1200)), controllerNode)
-    TestUtils.retry(60000) {
-      context.time.sleep(100)
-      context.poll()
-      manager.eventQueue.wakeup()
-      assertEquals(1200, manager.brokerEpoch)
-    }
+
+    // Accept new registration, response sets epoch to 1200
+    nextRegistrationRequest(1200L)
+
+    nextHeartbeatRequest()
+    assertEquals(1200L, manager.brokerEpoch)
+
     manager.close()
   }
 }
diff --git a/server-common/src/main/java/org/apache/kafka/queue/KafkaEventQueue.java b/server-common/src/main/java/org/apache/kafka/queue/KafkaEventQueue.java
index 36284ed7f3e94..972ceb1e9346f 100644
--- a/server-common/src/main/java/org/apache/kafka/queue/KafkaEventQueue.java
+++ b/server-common/src/main/java/org/apache/kafka/queue/KafkaEventQueue.java
@@ -20,6 +20,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.TreeMap;
 import java.util.concurrent.RejectedExecutionException;
@@ -513,4 +514,36 @@ public void close() throws InterruptedException {
         eventHandlerThread.join();
         log.info("closed event queue.");
     }
+
+    /**
+     * Returns the deferred event that the queue is waiting for, idling until
+     * its deadline comes, if there is any.
+     * If the queue has immediate work to do, this returns empty.
+     * This is useful for unit tests, where to make progress, we need to
+     * speed the clock up until the next scheduled event is ready to run.
+     */
+    public Optional<Event> firstDeferredIfIdling() {
+        lock.lock();
+        try {
+            if (eventHandler.head.next != eventHandler.head) {
+                // There are events ready to run immediately. The queue is not idling.
+                return Optional.empty();
+            }
+            Map.Entry<Long, EventContext> entry = eventHandler.deadlineMap.firstEntry();
+            if (entry == null) {
+                // The queue is idling, but not waiting for any deadline.
+                return Optional.empty();
+            }
+            EventContext eventContext = entry.getValue();
+            if (eventContext.insertionType != EventInsertionType.DEFERRED) {
+                // Any event with a deadline is put in `deadlineMap`.
+                // But events of type other than DEFERRED will run immediately,
+                // so the queue will not idle waiting for their deadline.
+                return Optional.empty();
+            }
+            return Optional.of(eventContext.event);
+        } finally {
+            lock.unlock();
+        }
+    }
 }

From da6bb365bc68a5d46ab37963429c39d972e011cd Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <cadonna@apache.org>
Date: Mon, 25 Mar 2024 20:59:41 +0100
Subject: [PATCH 202/258] KAFKA-16224: Do not retry committing if topic or
 partition deleted (#15581)

Current logic for auto-committing offsets when partitions are revoked
will retry continuously when getting UNKNOWN_TOPIC_OR_PARTITION,
leading to the member not completing the revocation in time.

This commit considers error UNKNOWN_TOPIC_OR_PARTITION to be fatal
in the context of an auto-commit of offsets before a revocation,
even though the error is defined as retriable. This ensures that
the revocation can finish in time.

Reviewers: Andrew Schofield <aschofield@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>, Lianet Magrans <lianetmr@gmail.com>
---
 .../internals/CommitRequestManager.java       | 24 ++++++++++++-------
 .../internals/MembershipManagerImpl.java      |  2 +-
 .../internals/CommitRequestManagerTest.java   |  4 ++--
 .../internals/MembershipManagerImplTest.java  |  4 ++--
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
index f7acbde60bdd5..67ea6eb3134c0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/CommitRequestManager.java
@@ -30,6 +30,7 @@
 import org.apache.kafka.common.errors.StaleMemberEpochException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.errors.UnstableOffsetCommitException;
 import org.apache.kafka.common.message.OffsetCommitRequestData;
 import org.apache.kafka.common.message.OffsetCommitResponseData;
@@ -272,14 +273,18 @@ private void maybeResetTimerWithBackoff(final CompletableFuture<Map<TopicPartiti
      * Commit consumed offsets if auto-commit is enabled, regardless of the auto-commit interval.
      * This is used for committing offsets before revoking partitions. This will retry committing
      * the latest offsets until the request succeeds, fails with a fatal error, or the timeout
-     * expires. Note that this considers {@link Errors#STALE_MEMBER_EPOCH} as a retriable error,
+     * expires. Note that:
+     * - Considers {@link Errors#STALE_MEMBER_EPOCH} as a retriable error,
      * and will retry it including the latest member ID and epoch received from the broker.
+     * - Considers {@link Errors#UNKNOWN_TOPIC_OR_PARTITION} as a fatal error, and will not retry
+     * it although the error extends RetriableException. The reason is that if a topic or partition
+     * is deleted, revocation would not finish in time since the auto commit would keep retrying.
      *
      * @return Future that will complete when the offsets are successfully committed. It will
      * complete exceptionally if the commit fails with a non-retriable error, or if the retry
      * timeout expires.
      */
-    public CompletableFuture<Void> maybeAutoCommitSyncNow(final long retryExpirationTimeMs) {
+    public CompletableFuture<Void> maybeAutoCommitSyncBeforeRevocation(final long retryExpirationTimeMs) {
         if (!autoCommitEnabled()) {
             return CompletableFuture.completedFuture(null);
         }
@@ -287,12 +292,12 @@ public CompletableFuture<Void> maybeAutoCommitSyncNow(final long retryExpiration
         CompletableFuture<Void> result = new CompletableFuture<>();
         OffsetCommitRequestState requestState =
             createOffsetCommitRequest(subscriptions.allConsumed(), Optional.of(retryExpirationTimeMs));
-        autoCommitSyncNowWithRetries(requestState, result);
+        autoCommitSyncBeforeRevocationWithRetries(requestState, result);
         return result;
     }
 
-    private void autoCommitSyncNowWithRetries(OffsetCommitRequestState requestAttempt,
-                                              CompletableFuture<Void> result) {
+    private void autoCommitSyncBeforeRevocationWithRetries(OffsetCommitRequestState requestAttempt,
+                                                           CompletableFuture<Void> result) {
         CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> commitAttempt = requestAutoCommit(requestAttempt);
         commitAttempt.whenComplete((committedOffsets, error) -> {
             if (error == null) {
@@ -300,16 +305,19 @@ private void autoCommitSyncNowWithRetries(OffsetCommitRequestState requestAttemp
             } else {
                 if (error instanceof RetriableException || isStaleEpochErrorAndValidEpochAvailable(error)) {
                     if (error instanceof TimeoutException && requestAttempt.isExpired) {
-                        log.debug("Auto-commit sync timed out and won't be retried anymore");
+                        log.debug("Auto-commit sync before revocation timed out and won't be retried anymore");
+                        result.completeExceptionally(error);
+                    } else if (error instanceof UnknownTopicOrPartitionException) {
+                        log.debug("Auto-commit sync before revocation failed because topic or partition were deleted");
                         result.completeExceptionally(error);
                     } else {
                         // Make sure the auto-commit is retries with the latest offsets
                         requestAttempt.offsets = subscriptions.allConsumed();
                         requestAttempt.resetFuture();
-                        autoCommitSyncNowWithRetries(requestAttempt, result);
+                        autoCommitSyncBeforeRevocationWithRetries(requestAttempt, result);
                     }
                 } else {
-                    log.debug("Auto-commit sync failed with non-retriable error", error);
+                    log.debug("Auto-commit sync before revocation failed with non-retriable error", error);
                     result.completeExceptionally(error);
                 }
             }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
index 574db9ebcd0fe..4710746995935 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
@@ -961,7 +961,7 @@ void maybeReconcile() {
         // best effort to commit the offsets in the case where the epoch might have changed while
         // the current reconciliation is in process. Note this is using the rebalance timeout as
         // it is the limit enforced by the broker to complete the reconciliation process.
-        commitResult = commitRequestManager.maybeAutoCommitSyncNow(getExpirationTimeForTimeout(rebalanceTimeoutMs));
+        commitResult = commitRequestManager.maybeAutoCommitSyncBeforeRevocation(getExpirationTimeForTimeout(rebalanceTimeoutMs));
 
         // Execute commit -> onPartitionsRevoked -> onPartitionsAssigned.
         commitResult.whenComplete((__, commitReqError) -> {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
index d22d618904232..e5b276adf74f1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/CommitRequestManagerTest.java
@@ -912,7 +912,7 @@ public void testAutoCommitSyncBeforeRevocationRetriesOnRetriableAndStaleEpoch(Er
         long expirationTimeMs = time.milliseconds() + retryBackoffMs * 2;
 
         // Send commit request expected to be retried on STALE_MEMBER_EPOCH error while it does not expire
-        commitRequestManager.maybeAutoCommitSyncNow(expirationTimeMs);
+        commitRequestManager.maybeAutoCommitSyncBeforeRevocation(expirationTimeMs);
 
         int newEpoch = 8;
         String memberId = "member1";
@@ -923,7 +923,7 @@ public void testAutoCommitSyncBeforeRevocationRetriesOnRetriableAndStaleEpoch(Er
 
         completeOffsetCommitRequestWithError(commitRequestManager, error);
 
-        if (error.exception() instanceof RetriableException || error == Errors.STALE_MEMBER_EPOCH) {
+        if ((error.exception() instanceof RetriableException || error == Errors.STALE_MEMBER_EPOCH) && error != Errors.UNKNOWN_TOPIC_OR_PARTITION) {
             assertEquals(1, commitRequestManager.pendingRequests.unsentOffsetCommits.size(),
                 "Request to be retried should be added to the outbound queue");
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index b2c75f495dd7e..cef608a902d1f 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -597,7 +597,7 @@ public void testSameAssignmentReconciledAgainWhenFenced() {
         membershipManager.onHeartbeatSuccess(createConsumerGroupHeartbeatResponse(assignment2).data());
         assertEquals(MemberState.RECONCILING, membershipManager.state());
         CompletableFuture<Void> commitResult = new CompletableFuture<>();
-        when(commitRequestManager.maybeAutoCommitSyncNow(anyLong())).thenReturn(commitResult);
+        when(commitRequestManager.maybeAutoCommitSyncBeforeRevocation(anyLong())).thenReturn(commitResult);
         membershipManager.poll(time.milliseconds());
 
         // Get fenced, commit completes
@@ -2457,7 +2457,7 @@ private CompletableFuture<Void> mockRevocationNoCallbacks(boolean withAutoCommit
         if (withAutoCommit) {
             when(commitRequestManager.autoCommitEnabled()).thenReturn(true);
             CompletableFuture<Void> commitResult = new CompletableFuture<>();
-            when(commitRequestManager.maybeAutoCommitSyncNow(anyLong())).thenReturn(commitResult);
+            when(commitRequestManager.maybeAutoCommitSyncBeforeRevocation(anyLong())).thenReturn(commitResult);
             return commitResult;
         } else {
             return CompletableFuture.completedFuture(null);

From ad960635a90ead44f4a696c0da7777b5d06f5837 Mon Sep 17 00:00:00 2001
From: Sean Quah <squah@confluent.io>
Date: Mon, 25 Mar 2024 23:08:23 +0000
Subject: [PATCH 203/258] KAFKA-16386: Convert NETWORK_EXCEPTIONs from KIP-890
 transaction verification (#15559)

KIP-890 Part 1 introduced verification of transactions with the
transaction coordinator on the `Produce` and `TxnOffsetCommit` paths.
This introduced the possibility of new errors when responding to those
requests. For backwards compatibility with older clients, a choice was
made to convert some of the new retriable errors to existing errors that
are expected and retried correctly by older clients.

`NETWORK_EXCEPTION` was forgotten about and not converted, but can occur
if, for example, the transaction coordinator is temporarily refusing
connections. Now, we convert it to:
 * `NOT_ENOUGH_REPLICAS` on the `Produce` path, just like the other
   retriable errors that can arise from transaction verification.
 * `COORDINATOR_LOAD_IN_PROGRESS` on the `TxnOffsetCommit` path. This
   error does not force coordinator lookup on clients, unlike
   `COORDINATOR_NOT_AVAILABLE`. Note that this deviates from KIP-890,
   which says that retriable errors should be converted to
   `COORDINATOR_NOT_AVAILABLE`.

Reviewers: Artem Livshits <alivshits@confluent.io>, David Jacot <djacot@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../coordinator/group/GroupMetadataManager.scala |  8 ++++++++
 .../main/scala/kafka/server/ReplicaManager.scala |  5 +++++
 .../coordinator/group/GroupCoordinatorTest.scala |  1 +
 .../unit/kafka/server/ReplicaManagerTest.scala   | 11 ++++++++++-
 .../group/GroupCoordinatorService.java           |  8 ++++++++
 .../group/GroupCoordinatorServiceTest.java       | 16 ++++++++++++----
 6 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
index 6012722332d56..b9dda91e25c16 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
@@ -1365,6 +1365,14 @@ object GroupMetadataManager {
 
   def maybeConvertOffsetCommitError(error: Errors) : Errors = {
     error match {
+      case Errors.NETWORK_EXCEPTION =>
+        // When committing offsets transactionally, we now verify the transaction with the
+        // transaction coordinator. Verification can fail with `NETWORK_EXCEPTION`, a retriable
+        // error which older clients may not expect and retry correctly. We translate the error to
+        // `COORDINATOR_LOAD_IN_PROGRESS` because it causes clients to retry the request without an
+        // unnecessary coordinator lookup.
+        Errors.COORDINATOR_LOAD_IN_PROGRESS
+
       case Errors.UNKNOWN_TOPIC_OR_PARTITION
            | Errors.NOT_ENOUGH_REPLICAS
            | Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND =>
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index ab9003b0a4e18..779ba575d4d6c 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -832,7 +832,12 @@ class ReplicaManager(val config: KafkaConfig,
           val customException =
             error match {
               case Errors.INVALID_TXN_STATE => Some(error.exception("Partition was not added to the transaction"))
+              // Transaction verification can fail with a retriable error that older clients may not
+              // retry correctly. Translate these to an error which will cause such clients to retry
+              // the produce request. We pick `NOT_ENOUGH_REPLICAS` because it does not trigger a
+              // metadata refresh.
               case Errors.CONCURRENT_TRANSACTIONS |
+                   Errors.NETWORK_EXCEPTION |
                    Errors.COORDINATOR_LOAD_IN_PROGRESS |
                    Errors.COORDINATOR_NOT_AVAILABLE |
                    Errors.NOT_COORDINATOR => Some(new NotEnoughReplicasException(
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
index c8417f678670b..0f8190b17cab6 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
@@ -3844,6 +3844,7 @@ class GroupCoordinatorTest {
 
     verifyErrors(Errors.INVALID_PRODUCER_ID_MAPPING, Errors.INVALID_PRODUCER_ID_MAPPING)
     verifyErrors(Errors.INVALID_TXN_STATE, Errors.INVALID_TXN_STATE)
+    verifyErrors(Errors.NETWORK_EXCEPTION, Errors.COORDINATOR_LOAD_IN_PROGRESS)
     verifyErrors(Errors.NOT_ENOUGH_REPLICAS, Errors.COORDINATOR_NOT_AVAILABLE)
     verifyErrors(Errors.NOT_LEADER_OR_FOLLOWER, Errors.NOT_COORDINATOR)
     verifyErrors(Errors.KAFKA_STORAGE_ERROR, Errors.NOT_COORDINATOR)
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index bfb822983ec81..ed64d986c269e 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -2533,7 +2533,16 @@ class ReplicaManagerTest {
   }
 
   @ParameterizedTest
-  @EnumSource(value = classOf[Errors], names = Array("NOT_COORDINATOR", "CONCURRENT_TRANSACTIONS", "COORDINATOR_LOAD_IN_PROGRESS", "COORDINATOR_NOT_AVAILABLE"))
+  @EnumSource(
+    value = classOf[Errors],
+    names = Array(
+      "NOT_COORDINATOR",
+      "CONCURRENT_TRANSACTIONS",
+      "NETWORK_EXCEPTION",
+      "COORDINATOR_LOAD_IN_PROGRESS",
+      "COORDINATOR_NOT_AVAILABLE"
+    )
+  )
   def testVerificationErrorConversions(error: Errors): Unit = {
     val tp0 = new TopicPartition(topic, 0)
     val producerId = 24L
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
index d6193ce799523..fb4dc782b6c44 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java
@@ -1096,6 +1096,14 @@ private <IN, OUT> OUT handleOperationException(
                     operationName, operationInput, exception.getMessage(), exception);
                 return handler.apply(Errors.UNKNOWN_SERVER_ERROR, null);
 
+            case NETWORK_EXCEPTION:
+                // When committing offsets transactionally, we now verify the transaction with the
+                // transaction coordinator. Verification can fail with `NETWORK_EXCEPTION`, a
+                // retriable error which older clients may not expect and retry correctly. We
+                // translate the error to `COORDINATOR_LOAD_IN_PROGRESS` because it causes clients
+                // to retry the request without an unnecessary coordinator lookup.
+                return handler.apply(Errors.COORDINATOR_LOAD_IN_PROGRESS, null);
+
             case UNKNOWN_TOPIC_OR_PARTITION:
             case NOT_ENOUGH_REPLICAS:
             case REQUEST_TIMED_OUT:
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
index 121e624f2b053..ee2defb5ae647 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorServiceTest.java
@@ -73,6 +73,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.NullSource;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -1900,8 +1901,15 @@ public void testCommitTransactionalOffsets() throws ExecutionException, Interrup
         assertEquals(response, future.get());
     }
 
-    @Test
-    public void testCommitTransactionalOffsetsWithWrappedError() throws ExecutionException, InterruptedException {
+    @ParameterizedTest
+    @CsvSource({
+        "NOT_ENOUGH_REPLICAS, COORDINATOR_NOT_AVAILABLE",
+        "NETWORK_EXCEPTION, COORDINATOR_LOAD_IN_PROGRESS"
+    })
+    public void testCommitTransactionalOffsetsWithWrappedError(
+        Errors error,
+        Errors expectedError
+    ) throws ExecutionException, InterruptedException {
         CoordinatorRuntime<GroupCoordinatorShard, Record> runtime = mockRuntime();
         GroupCoordinatorService service = new GroupCoordinatorService(
             new LogContext(),
@@ -1929,7 +1937,7 @@ public void testCommitTransactionalOffsetsWithWrappedError() throws ExecutionExc
                 .setName("topic")
                 .setPartitions(Collections.singletonList(new TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition()
                     .setPartitionIndex(0)
-                    .setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code())))));
+                    .setErrorCode(expectedError.code())))));
 
         when(runtime.scheduleTransactionalWriteOperation(
             ArgumentMatchers.eq("txn-commit-offset"),
@@ -1940,7 +1948,7 @@ public void testCommitTransactionalOffsetsWithWrappedError() throws ExecutionExc
             ArgumentMatchers.eq(Duration.ofMillis(5000)),
             ArgumentMatchers.any(),
             ArgumentMatchers.any()
-        )).thenReturn(FutureUtils.failedFuture(new CompletionException(Errors.NOT_ENOUGH_REPLICAS.exception())));
+        )).thenReturn(FutureUtils.failedFuture(new CompletionException(error.exception())));
 
         CompletableFuture<TxnOffsetCommitResponseData> future = service.commitTransactionalOffsets(
             requestContext(ApiKeys.TXN_OFFSET_COMMIT),

From fa1cf7975e4005f59bf7a1ffc49e9d018c0eb572 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Tue, 26 Mar 2024 08:44:59 +0800
Subject: [PATCH 204/258] KAFKA-16409: DeleteRecordsCommand should use standard
 exception handling (#15586)

DeleteRecordsCommand should use standard exception handling

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../apache/kafka/tools/DeleteRecordsCommand.java   | 14 +++++++++++++-
 .../kafka/tools/DeleteRecordsCommandTest.java      | 13 +++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/DeleteRecordsCommand.java b/tools/src/main/java/org/apache/kafka/tools/DeleteRecordsCommand.java
index 5e44865b200f8..6bf82bedc5173 100644
--- a/tools/src/main/java/org/apache/kafka/tools/DeleteRecordsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/DeleteRecordsCommand.java
@@ -24,6 +24,7 @@
 import org.apache.kafka.clients.admin.DeleteRecordsResult;
 import org.apache.kafka.clients.admin.RecordsToDelete;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.common.AdminCommandFailedException;
 import org.apache.kafka.server.common.AdminOperationException;
@@ -61,7 +62,18 @@ public class DeleteRecordsCommand {
     private static final DecodeJson.DecodeString STRING = new DecodeJson.DecodeString();
 
     public static void main(String[] args) throws Exception {
-        execute(args, System.out);
+        Exit.exit(mainNoExit(args));
+    }
+
+    static int mainNoExit(String... args) {
+        try {
+            execute(args, System.out);
+            return 0;
+        } catch (Exception e) {
+            System.err.println(e.getMessage());
+            System.err.println(Utils.stackTrace(e));
+            return 1;
+        }
     }
 
     static Map<TopicPartition, List<Long>> parseOffsetJsonStringWithoutDedup(String jsonData) throws JsonProcessingException {
diff --git a/tools/src/test/java/org/apache/kafka/tools/DeleteRecordsCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/DeleteRecordsCommandTest.java
index b540c8ffc17ef..0f4e551a48a40 100644
--- a/tools/src/test/java/org/apache/kafka/tools/DeleteRecordsCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/DeleteRecordsCommandTest.java
@@ -120,7 +120,11 @@ private static void executeAndAssertOutput(String json, String expOut, Admin adm
 class DeleteRecordsCommandUnitTest {
     @Test
     public void testOffsetFileNotExists() {
-        assertThrows(IOException.class, () -> DeleteRecordsCommand.main(new String[]{
+        assertThrows(IOException.class, () -> DeleteRecordsCommand.execute(new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--offset-json-file", "/not/existing/file"
+        }, System.out));
+        assertEquals(1, DeleteRecordsCommand.mainNoExit(new String[]{
             "--bootstrap-server", "localhost:9092",
             "--offset-json-file", "/not/existing/file"
         }));
@@ -128,7 +132,12 @@ public void testOffsetFileNotExists() {
 
     @Test
     public void testCommandConfigNotExists() {
-        assertThrows(NoSuchFileException.class, () -> DeleteRecordsCommand.main(new String[] {
+        assertThrows(NoSuchFileException.class, () -> DeleteRecordsCommand.execute(new String[] {
+            "--bootstrap-server", "localhost:9092",
+            "--offset-json-file", "/not/existing/file",
+            "--command-config", "/another/not/existing/file"
+        }, System.out));
+        assertEquals(1, DeleteRecordsCommand.mainNoExit(new String[] {
             "--bootstrap-server", "localhost:9092",
             "--offset-json-file", "/not/existing/file",
             "--command-config", "/another/not/existing/file"

From 2d4abb85bf4a3afb1e3359a05786ab8f3fda127e Mon Sep 17 00:00:00 2001
From: Dmitry Werner <grimekillah@gmail.com>
Date: Tue, 26 Mar 2024 12:44:23 +0500
Subject: [PATCH 205/258] KAFKA-16415 Fix handling of "--version" option in
 ConsumerGroupCommand (#15592)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../consumer/group/ConsumerGroupCommand.java      |  2 --
 .../group/ConsumerGroupCommandOptions.java        |  2 ++
 .../consumer/group/DescribeConsumerGroupTest.java | 15 +++++++++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
index 35371ce75aa65..5aafde1c0fe36 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommand.java
@@ -92,8 +92,6 @@ public class ConsumerGroupCommand {
     public static void main(String[] args) {
         ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(args);
         try {
-            CommandLineUtils.maybePrintHelpOrVersion(opts, "This tool helps to list all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.");
-
             // should have exactly one action
             long actions = Stream.of(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt, opts.deleteOffsetsOpt).filter(opts.options::has).count();
             if (actions != 1)
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
index 3eb811bcaf829..121594be4adab 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandOptions.java
@@ -206,6 +206,8 @@ private ConsumerGroupCommandOptions(String[] args) {
 
     @SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"})
     void checkArgs() {
+        CommandLineUtils.maybePrintHelpOrVersion(this, "This tool helps to list all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.");
+
         CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt);
 
         if (options.has(describeOpt)) {
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
index f1b6bba9ac1af..4b5b0b93f53ac 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/DescribeConsumerGroupTest.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.utils.AppInfoParser;
 import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.test.TestUtils;
 import org.apache.kafka.tools.ToolsTestUtils;
@@ -126,6 +127,20 @@ public void testDescribeWithStateValue(String quorum) {
         assertTrue(exitMessage.get().contains("Option [describe] does not take a value for [state]"));
     }
 
+    @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
+    @ValueSource(strings = {"zk", "kraft"})
+    public void testPrintVersion(String quorum) {
+        ToolsTestUtils.MockExitProcedure exitProcedure = new ToolsTestUtils.MockExitProcedure();
+        Exit.setExitProcedure(exitProcedure);
+        try {
+            String out = ToolsTestUtils.captureStandardOut(() -> ConsumerGroupCommandOptions.fromArgs(new String[]{"--version"}));
+            assertEquals(0, exitProcedure.statusCode());
+            assertEquals(AppInfoParser.getVersion(), out);
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_AND_GROUP_PROTOCOL_NAMES)
     @MethodSource({"getTestQuorumAndGroupProtocolParametersAll"})
     public void testDescribeOffsetsOfNonExistingGroup(String quorum, String groupProtocol) throws Exception {

From d70cfda02519a985ff881b2776f7393c10c7696c Mon Sep 17 00:00:00 2001
From: Avneesh Kumar Singhal <116718699+avn-confluent@users.noreply.github.com>
Date: Tue, 26 Mar 2024 14:59:27 +0530
Subject: [PATCH 206/258] update tags and added seamphore tags as part of
 semaphore migration (#1102)

* added new names

* fixed URL

* added semaphorebuildurl for tracking which url job spin up that vagrant worker

* fixed syntax

* Added new SemaphoreWorkflowUrl and SemaphoreJobId in ec2 tags
---
 Vagrantfile | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Vagrantfile b/Vagrantfile
index 36267c6f18cc6..9318cf0c7f30c 100644
--- a/Vagrantfile
+++ b/Vagrantfile
@@ -164,6 +164,8 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config|
         'role' => 'ce-kafka',
         'Owner' => 'ce-kafka',
         'JenkinsBuildUrl' => ENV['BUILD_URL'],
+        'SemaphoreWorkflowUrl' => ENV['SEMAPHORE_WORKFLOW_URL'],
+        'SemaphoreJobId' => ENV['SEMAPHORE_JOB_ID'],
         'cflt_environment' => 'devel',
         'cflt_partition' => 'onprem',
         'cflt_managed_by' => 'iac',

From 6f8d4fe26b7ceaa369c5e974a8cf54cc2714f9a7 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Tue, 26 Mar 2024 20:09:29 +0800
Subject: [PATCH 207/258] KAFKA-15949: Unify metadata.version format in log and
 error message (#15505)

There were different words for metadata.version like metadata version or metadataVersion. Unify format as metadata.version.

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../kafka/server/ControllerRegistrationManager.scala   |  2 +-
 core/src/main/scala/kafka/server/KafkaApis.scala       |  2 +-
 core/src/main/scala/kafka/server/KafkaConfig.scala     |  2 +-
 core/src/main/scala/kafka/tools/StorageTool.scala      | 10 +++++-----
 .../integration/kafka/server/KRaftClusterTest.scala    |  4 ++--
 .../server/AlterUserScramCredentialsRequestTest.scala  |  2 +-
 .../test/scala/unit/kafka/tools/StorageToolTest.scala  |  4 ++--
 .../apache/kafka/controller/FeatureControlManager.java |  2 +-
 .../org/apache/kafka/controller/QuorumFeatures.java    |  2 +-
 .../apache/kafka/controller/ScramControlManager.java   |  4 ++--
 .../kafka/image/MetadataVersionChangeException.java    |  2 +-
 .../image/writer/UnwritableMetadataException.java      |  2 +-
 .../kafka/metadata/bootstrap/BootstrapMetadata.java    |  2 +-
 .../apache/kafka/controller/QuorumControllerTest.java  |  4 ++--
 .../apache/kafka/controller/QuorumFeaturesTest.java    |  2 +-
 .../apache/kafka/image/MetadataVersionChangeTest.java  |  8 ++++----
 .../kafka/image/writer/ImageWriterOptionsTest.java     |  2 +-
 .../metadata/bootstrap/BootstrapMetadataTest.java      |  2 +-
 .../apache/kafka/server/common/MetadataVersion.java    |  2 +-
 .../internals/assignment/AssignorConfiguration.java    |  4 ++--
 .../java/org/apache/kafka/tools/FeatureCommand.java    |  4 ++--
 21 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/core/src/main/scala/kafka/server/ControllerRegistrationManager.scala b/core/src/main/scala/kafka/server/ControllerRegistrationManager.scala
index 20edae23ca829..4b5843f8d31c5 100644
--- a/core/src/main/scala/kafka/server/ControllerRegistrationManager.scala
+++ b/core/src/main/scala/kafka/server/ControllerRegistrationManager.scala
@@ -205,7 +205,7 @@ class ControllerRegistrationManager(
       debug("maybeSendControllerRegistration: cannot register yet because the channel manager has " +
           "not been initialized.")
     } else if (!metadataVersion.isControllerRegistrationSupported) {
-      info("maybeSendControllerRegistration: cannot register yet because the metadata version is " +
+      info("maybeSendControllerRegistration: cannot register yet because the metadata.version is " +
           s"still $metadataVersion, which does not support KIP-919 controller registration.")
     } else if (pendingRpc) {
       info("maybeSendControllerRegistration: waiting for the previous RPC to complete.")
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index 3b76696f199f1..fad6890df4050 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -2516,7 +2516,7 @@ class KafkaApis(val requestChannel: RequestChannel,
 
   def ensureInterBrokerVersion(version: MetadataVersion): Unit = {
     if (config.interBrokerProtocolVersion.isLessThan(version))
-      throw new UnsupportedVersionException(s"inter.broker.protocol.version: ${config.interBrokerProtocolVersion.version} is less than the required version: ${version.version}")
+      throw new UnsupportedVersionException(s"inter.broker.protocol.version: ${config.interBrokerProtocolVersion} is less than the required version: ${version}")
   }
 
   def handleAddPartitionsToTxnRequest(request: RequestChannel.Request, requestLocal: RequestLocal): Unit = {
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index eaddb047da8d5..4290ce1da8d45 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -1733,7 +1733,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
       } else {
         warn(s"${KafkaConfig.InterBrokerProtocolVersionProp} is deprecated in KRaft mode as of 3.3 and will only " +
           s"be read when first upgrading from a KRaft prior to 3.3. See kafka-storage.sh help for details on setting " +
-          s"the metadata version for a new KRaft cluster.")
+          s"the metadata.version for a new KRaft cluster.")
       }
     }
     // In KRaft mode, we pin this value to the minimum KRaft-supported version. This prevents inadvertent usage of
diff --git a/core/src/main/scala/kafka/tools/StorageTool.scala b/core/src/main/scala/kafka/tools/StorageTool.scala
index d328d1df8718f..be31b1648b195 100644
--- a/core/src/main/scala/kafka/tools/StorageTool.scala
+++ b/core/src/main/scala/kafka/tools/StorageTool.scala
@@ -62,13 +62,13 @@ object StorageTool extends Logging {
           val metadataVersion = getMetadataVersion(namespace,
             Option(config.get.originals.get(KafkaConfig.InterBrokerProtocolVersionProp)).map(_.toString))
           if (!metadataVersion.isKRaftSupported) {
-            throw new TerseFailure(s"Must specify a valid KRaft metadata version of at least 3.0.")
+            throw new TerseFailure(s"Must specify a valid KRaft metadata.version of at least ${MetadataVersion.IBP_3_0_IV0}.")
           }
           if (!metadataVersion.isProduction) {
             if (config.get.unstableMetadataVersionsEnabled) {
-              System.out.println(s"WARNING: using pre-production metadata version $metadataVersion.")
+              System.out.println(s"WARNING: using pre-production metadata.version $metadataVersion.")
             } else {
-              throw new TerseFailure(s"Metadata version $metadataVersion is not ready for production use yet.")
+              throw new TerseFailure(s"The metadata.version $metadataVersion is not ready for production use yet.")
             }
           }
           val metaProperties = new MetaProperties.Builder().
@@ -79,7 +79,7 @@ object StorageTool extends Logging {
           val metadataRecords : ArrayBuffer[ApiMessageAndVersion] = ArrayBuffer()
           getUserScramCredentialRecords(namespace).foreach(userScramCredentialRecords => {
             if (!metadataVersion.isScramSupported) {
-              throw new TerseFailure(s"SCRAM is only supported in metadataVersion IBP_3_5_IV2 or later.")
+              throw new TerseFailure(s"SCRAM is only supported in metadata.version ${MetadataVersion.IBP_3_5_IV2} or later.")
             }
             for (record <- userScramCredentialRecords) {
               metadataRecords.append(new ApiMessageAndVersion(record, 0.toShort))
@@ -139,7 +139,7 @@ object StorageTool extends Logging {
       action(storeTrue())
     formatParser.addArgument("--release-version", "-r").
       action(store()).
-      help(s"A KRaft release version to use for the initial metadata version. The minimum is 3.0, the default is ${MetadataVersion.LATEST_PRODUCTION.version()}")
+      help(s"A KRaft release version to use for the initial metadata.version. The minimum is ${MetadataVersion.IBP_3_0_IV0}, the default is ${MetadataVersion.LATEST_PRODUCTION}")
 
     parser.parseArgsOrFail(args)
   }
diff --git a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
index 3be1b400ab564..75f07b8efe853 100644
--- a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
@@ -358,7 +358,7 @@ class KRaftClusterTest {
 
   @Test
   def testCreateClusterInvalidMetadataVersion(): Unit = {
-    assertEquals("Bootstrap metadata versions before 3.3-IV0 are not supported. Can't load " +
+    assertEquals("Bootstrap metadata.version before 3.3-IV0 are not supported. Can't load " +
       "metadata from testkit", assertThrows(classOf[RuntimeException], () => {
         new KafkaClusterTestKit.Builder(
           new TestKitNodes.Builder().
@@ -963,7 +963,7 @@ class KRaftClusterTest {
         admin.close()
       }
       TestUtils.waitUntilTrue(() => cluster.brokers().get(1).metadataCache.currentImage().features().metadataVersion().equals(MetadataVersion.latestTesting()),
-        "Timed out waiting for metadata version update.")
+        "Timed out waiting for metadata.version update.")
     } finally {
       cluster.close()
     }
diff --git a/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala
index b20478cf6f2e2..f45299e9a3e41 100644
--- a/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala
@@ -405,7 +405,7 @@ class AlterUserScramCredentialsRequestTest extends BaseRequestTest {
     assertEquals(1, results.size)
     checkAllErrorsAlteringCredentials(results, Errors.UNSUPPORTED_VERSION,
                                       "when altering the credentials on unsupported IBP version")
-    assertEquals("The current metadata version does not support SCRAM", results.get(0).errorMessage)
+    assertEquals("The current metadata.version does not support SCRAM", results.get(0).errorMessage)
   }
 
   private def sendAlterUserScramCredentialsRequest(request: AlterUserScramCredentialsRequest, socketServer: SocketServer = adminSocketServer): AlterUserScramCredentialsResponse = {
diff --git a/core/src/test/scala/unit/kafka/tools/StorageToolTest.scala b/core/src/test/scala/unit/kafka/tools/StorageToolTest.scala
index c3cdce65cbe15..4470d877536c5 100644
--- a/core/src/test/scala/unit/kafka/tools/StorageToolTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/StorageToolTest.scala
@@ -333,7 +333,7 @@ Found problem:
     try {
       assertEquals(1, StorageTool.main(args))
     } catch {
-      case e: StorageToolTestException => assertEquals(s"SCRAM is only supported in metadataVersion IBP_3_5_IV2 or later.", exitString)
+      case e: StorageToolTestException => assertEquals(s"SCRAM is only supported in metadata.version ${MetadataVersion.IBP_3_5_IV2} or later.", exitString)
     } finally {
       Exit.resetExitProcedure()
     }
@@ -428,7 +428,7 @@ Found problem:
       assertEquals("", exitString)
       assertEquals(0, exitStatus)
     } else {
-      assertEquals(s"Metadata version ${MetadataVersion.latestTesting().toString} is not ready for " +
+      assertEquals(s"The metadata.version ${MetadataVersion.latestTesting().toString} is not ready for " +
         "production use yet.", exitString)
       assertEquals(1, exitStatus)
     }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java
index d5911d4f54f8a..cf1be493591c2 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/FeatureControlManager.java
@@ -397,7 +397,7 @@ public void replay(FeatureLevelRecord record) {
         if (record.name().equals(MetadataVersion.FEATURE_NAME)) {
             MetadataVersion mv = MetadataVersion.fromFeatureLevel(record.featureLevel());
             metadataVersion.set(mv);
-            log.info("Replayed a FeatureLevelRecord setting metadata version to {}", mv);
+            log.info("Replayed a FeatureLevelRecord setting metadata.version to {}", mv);
         } else {
             if (record.featureLevel() == 0) {
                 finalizedVersions.remove(record.name());
diff --git a/metadata/src/main/java/org/apache/kafka/controller/QuorumFeatures.java b/metadata/src/main/java/org/apache/kafka/controller/QuorumFeatures.java
index 9b0355c0db71f..b0183295b7b46 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/QuorumFeatures.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/QuorumFeatures.java
@@ -108,7 +108,7 @@ public Optional<String> reasonAllControllersZkMigrationNotReady(
         Map<Integer, ControllerRegistration> controllers
     ) {
         if (!metadataVersion.isMigrationSupported()) {
-            return Optional.of("Metadata version too low at " + metadataVersion);
+            return Optional.of("The metadata.version too low at " + metadataVersion);
         } else if (!metadataVersion.isControllerRegistrationSupported()) {
             return Optional.empty();
         }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/ScramControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ScramControlManager.java
index 2a876053e4021..91fda0779847e 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/ScramControlManager.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/ScramControlManager.java
@@ -188,7 +188,7 @@ public ControllerResult<AlterUserScramCredentialsResponseData> alterCredentials(
                 } else {
                     if (!scramIsSupported) {
                         userToError.put(deletion.name(), new ApiError(UNSUPPORTED_VERSION,
-                            "The current metadata version does not support SCRAM"));
+                            "The current metadata.version does not support SCRAM"));
                     } else {
                         ApiError error = validateDeletion(deletion);
                         if (error.isFailure()) {
@@ -209,7 +209,7 @@ public ControllerResult<AlterUserScramCredentialsResponseData> alterCredentials(
                 } else {
                     if (!scramIsSupported) {
                         userToError.put(upsertion.name(), new ApiError(UNSUPPORTED_VERSION,
-                            "The current metadata version does not support SCRAM"));
+                            "The current metadata.version does not support SCRAM"));
                     } else {
                         ApiError error = validateUpsertion(upsertion);
                         if (error.isFailure()) {
diff --git a/metadata/src/main/java/org/apache/kafka/image/MetadataVersionChangeException.java b/metadata/src/main/java/org/apache/kafka/image/MetadataVersionChangeException.java
index a4c931f77b7f6..a2678aba5c865 100644
--- a/metadata/src/main/java/org/apache/kafka/image/MetadataVersionChangeException.java
+++ b/metadata/src/main/java/org/apache/kafka/image/MetadataVersionChangeException.java
@@ -26,7 +26,7 @@ public final class MetadataVersionChangeException extends RuntimeException {
     private final MetadataVersionChange change;
 
     public MetadataVersionChangeException(MetadataVersionChange change) {
-        super("The metadata version is changing from " + change.oldVersion() + " to " +
+        super("The metadata.version is changing from " + change.oldVersion() + " to " +
                 change.newVersion());
         this.change = change;
     }
diff --git a/metadata/src/main/java/org/apache/kafka/image/writer/UnwritableMetadataException.java b/metadata/src/main/java/org/apache/kafka/image/writer/UnwritableMetadataException.java
index 46ecbcff92ba7..bf3c514922d03 100644
--- a/metadata/src/main/java/org/apache/kafka/image/writer/UnwritableMetadataException.java
+++ b/metadata/src/main/java/org/apache/kafka/image/writer/UnwritableMetadataException.java
@@ -33,7 +33,7 @@ public UnwritableMetadataException(
         String loss
     ) {
         super("Metadata has been lost because the following could not be represented " +
-            "in metadata version " + metadataVersion + ": " + loss);
+            "in metadata.version " + metadataVersion + ": " + loss);
         this.metadataVersion = metadataVersion;
         this.loss = loss;
     }
diff --git a/metadata/src/main/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadata.java b/metadata/src/main/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadata.java
index 53c6de27bc6e1..5f0265179b89c 100644
--- a/metadata/src/main/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadata.java
+++ b/metadata/src/main/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadata.java
@@ -80,7 +80,7 @@ public static Optional<MetadataVersion> recordToMetadataVersion(ApiMessage recor
     ) {
         this.records = Objects.requireNonNull(records);
         if (metadataVersion.isLessThan(MINIMUM_BOOTSTRAP_VERSION)) {
-            throw new RuntimeException("Bootstrap metadata versions before " +
+            throw new RuntimeException("Bootstrap metadata.version before " +
                     MINIMUM_BOOTSTRAP_VERSION + " are not supported. Can't load metadata from " +
                     source);
         }
diff --git a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
index bcdb3ed791923..e907514e05b90 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
@@ -1257,7 +1257,7 @@ public void testUpgradeFromPreProductionVersion() throws Exception {
             QuorumController active = controlEnv.activeController();
             TestUtils.waitForCondition(() ->
                 active.featureControl().metadataVersion().equals(MetadataVersion.IBP_3_0_IV1),
-                "Failed to get a metadata version of " + MetadataVersion.IBP_3_0_IV1);
+                "Failed to get a metadata.version of " + MetadataVersion.IBP_3_0_IV1);
             // The ConfigRecord in our bootstrap should not have been applied, since there
             // were already records present.
             assertEquals(Collections.emptyMap(), active.configurationControl().
@@ -1288,7 +1288,7 @@ public void testInsertBootstrapRecordsToEmptyLog() throws Exception {
                 FinalizedControllerFeatures features = active.finalizedFeatures(ctx).get();
                 Optional<Short> metadataVersionOpt = features.get(MetadataVersion.FEATURE_NAME);
                 return Optional.of(MetadataVersion.IBP_3_3_IV1.featureLevel()).equals(metadataVersionOpt);
-            }, "Failed to see expected metadata version from bootstrap metadata");
+            }, "Failed to see expected metadata.version from bootstrap metadata");
 
             TestUtils.waitForCondition(() -> {
                 ConfigResource defaultBrokerResource = new ConfigResource(BROKER, "");
diff --git a/metadata/src/test/java/org/apache/kafka/controller/QuorumFeaturesTest.java b/metadata/src/test/java/org/apache/kafka/controller/QuorumFeaturesTest.java
index e0bab331e673c..099acea2de4be 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/QuorumFeaturesTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/QuorumFeaturesTest.java
@@ -95,7 +95,7 @@ public void testIsControllerId() {
 
     @Test
     public void testZkMigrationNotReadyIfMetadataVersionTooLow() {
-        assertEquals(Optional.of("Metadata version too low at 3.0-IV1"),
+        assertEquals(Optional.of("The metadata.version too low at 3.0-IV1"),
             QUORUM_FEATURES.reasonAllControllersZkMigrationNotReady(
                 MetadataVersion.IBP_3_0_IV1, Collections.emptyMap()));
     }
diff --git a/metadata/src/test/java/org/apache/kafka/image/MetadataVersionChangeTest.java b/metadata/src/test/java/org/apache/kafka/image/MetadataVersionChangeTest.java
index 5356f59b82e4d..decf6d315f6aa 100644
--- a/metadata/src/test/java/org/apache/kafka/image/MetadataVersionChangeTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/MetadataVersionChangeTest.java
@@ -50,11 +50,11 @@ public void testIsDowngrade() {
 
     @Test
     public void testMetadataVersionChangeExceptionToString() {
-        assertEquals("org.apache.kafka.image.MetadataVersionChangeException: The metadata " +
-            "version is changing from 3.0-IV1 to 3.3-IV0",
+        assertEquals("org.apache.kafka.image.MetadataVersionChangeException: The metadata.version " +
+            "is changing from 3.0-IV1 to 3.3-IV0",
                 new MetadataVersionChangeException(CHANGE_3_0_IV1_TO_3_3_IV0).toString());
-        assertEquals("org.apache.kafka.image.MetadataVersionChangeException: The metadata " +
-            "version is changing from 3.3-IV0 to 3.0-IV1",
+        assertEquals("org.apache.kafka.image.MetadataVersionChangeException: The metadata.version " +
+            "is changing from 3.3-IV0 to 3.0-IV1",
                 new MetadataVersionChangeException(CHANGE_3_3_IV0_TO_3_0_IV1).toString());
     }
 }
diff --git a/metadata/src/test/java/org/apache/kafka/image/writer/ImageWriterOptionsTest.java b/metadata/src/test/java/org/apache/kafka/image/writer/ImageWriterOptionsTest.java
index c8653c46d8aec..8db386117ab51 100644
--- a/metadata/src/test/java/org/apache/kafka/image/writer/ImageWriterOptionsTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/writer/ImageWriterOptionsTest.java
@@ -61,7 +61,7 @@ public void testHandleLoss() {
              i < MetadataVersion.VERSIONS.length;
              i++) {
             MetadataVersion version = MetadataVersion.VERSIONS[i];
-            String formattedMessage = String.format("Metadata has been lost because the following could not be represented in metadata version %s: %s", version, expectedMessage);
+            String formattedMessage = String.format("Metadata has been lost because the following could not be represented in metadata.version %s: %s", version, expectedMessage);
             Consumer<UnwritableMetadataException> customLossHandler = e -> assertEquals(formattedMessage, e.getMessage());
             ImageWriterOptions options = new ImageWriterOptions.Builder()
                     .setMetadataVersion(version)
diff --git a/metadata/src/test/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadataTest.java b/metadata/src/test/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadataTest.java
index 3c75f4cadcb2f..f22351f10fbaf 100644
--- a/metadata/src/test/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadataTest.java
+++ b/metadata/src/test/java/org/apache/kafka/metadata/bootstrap/BootstrapMetadataTest.java
@@ -85,7 +85,7 @@ public void testCopyWithOnlyVersion() {
     public void testFromRecordsListWithOldMetadataVersion() {
         RuntimeException exception = assertThrows(RuntimeException.class,
             () -> BootstrapMetadata.fromRecords(RECORDS_WITH_OLD_METADATA_VERSION, "quux"));
-        assertEquals("Bootstrap metadata versions before 3.3-IV0 are not supported. Can't load " +
+        assertEquals("Bootstrap metadata.version before 3.3-IV0 are not supported. Can't load " +
             "metadata from quux", exception.getMessage());
     }
 }
diff --git a/server-common/src/main/java/org/apache/kafka/server/common/MetadataVersion.java b/server-common/src/main/java/org/apache/kafka/server/common/MetadataVersion.java
index d38ecff19577e..22e545c11292c 100644
--- a/server-common/src/main/java/org/apache/kafka/server/common/MetadataVersion.java
+++ b/server-common/src/main/java/org/apache/kafka/server/common/MetadataVersion.java
@@ -563,7 +563,7 @@ public static MetadataVersion fromFeatureLevel(short version) {
                 return metadataVersion;
             }
         }
-        throw new IllegalArgumentException("No MetadataVersion with metadata version " + version);
+        throw new IllegalArgumentException("No MetadataVersion with feature level " + version);
     }
 
     /**
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
index 6d99d93536e5d..a917b989ba602 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java
@@ -155,7 +155,7 @@ public int configuredMetadataVersion(final int priorVersion) {
             switch (UpgradeFromValues.getValueFromString(upgradeFrom)) {
                 case UPGRADE_FROM_0100:
                     log.info(
-                        "Downgrading metadata version from {} to 1 for upgrade from 0.10.0.x.",
+                        "Downgrading metadata.version from {} to 1 for upgrade from 0.10.0.x.",
                         LATEST_SUPPORTED_VERSION
                     );
                     return 1;
@@ -165,7 +165,7 @@ public int configuredMetadataVersion(final int priorVersion) {
                 case UPGRADE_FROM_10:
                 case UPGRADE_FROM_11:
                     log.info(
-                        "Downgrading metadata version from {} to 2 for upgrade from {}.x.",
+                        "Downgrading metadata.version from {} to 2 for upgrade from {}.x.",
                         LATEST_SUPPORTED_VERSION,
                         upgradeFrom
                     );
diff --git a/tools/src/main/java/org/apache/kafka/tools/FeatureCommand.java b/tools/src/main/java/org/apache/kafka/tools/FeatureCommand.java
index ef3e4363a24e3..34cc3b8411546 100644
--- a/tools/src/main/java/org/apache/kafka/tools/FeatureCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/FeatureCommand.java
@@ -242,8 +242,8 @@ private static void handleUpgradeOrDowngrade(String op, Namespace namespace, Adm
             try {
                 version = MetadataVersion.fromVersionString(metadata);
             } catch (Throwable e) {
-                throw new TerseException("Unsupported metadata version " + metadata +
-                        ". Supported metadata versions are " + metadataVersionsToString(
+                throw new TerseException("Unsupported metadata.version " + metadata +
+                        ". Supported metadata.version are " + metadataVersionsToString(
                         MetadataVersion.MINIMUM_BOOTSTRAP_VERSION, MetadataVersion.latestProduction()));
             }
             updates.put(MetadataVersion.FEATURE_NAME, new FeatureUpdate(version.featureLevel(), upgradeType));

From b667e61ea7c7ed53406338df2434e97ef684d0dd Mon Sep 17 00:00:00 2001
From: "Cheng-Kai, Zhang" <kai821104@gmail.com>
Date: Wed, 27 Mar 2024 01:06:47 +0800
Subject: [PATCH 208/258] KAFKA-16416 Use NetworkClientTest to replace
 RequestResponseTest to be the example of log4j output (#15596)

Reviewers: Kirk True <kirk@kirktrue.pro>, Chia-Ping Tsai <chia7712@gmail.com>
---
 README.md | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index b857e7c1843e8..d8ee34b6188a3 100644
--- a/README.md
+++ b/README.md
@@ -54,9 +54,14 @@ Follow instructions in https://kafka.apache.org/quickstart
     ./gradlew clients:test --tests org.apache.kafka.clients.MetadataTest.testTimeToNextUpdate
 
 ### Running a particular unit/integration test with log4j output ###
-Change the log4j setting in either `clients/src/test/resources/log4j.properties` or `core/src/test/resources/log4j.properties`
+By default, there will be only small number of logs output while testing. You can adjust it by changing the `log4j.properties` file in the module's `src/test/resources` directory.
 
-    ./gradlew clients:test --tests RequestResponseTest
+For example, if you want to see more logs for clients project tests, you can modify [the line](https://github.com/apache/kafka/blob/trunk/clients/src/test/resources/log4j.properties#L21) in `clients/src/test/resources/log4j.properties` 
+to `log4j.logger.org.apache.kafka=INFO` and then run:
+    
+    ./gradlew cleanTest clients:test --tests NetworkClientTest   
+
+And you should see `INFO` level logs in the file under the `clients/build/test-results/test` directory.
 
 ### Specifying test retries ###
 By default, each failed test is retried once up to a maximum of five retries per test run. Tests are retried at the end of the test task. Adjust these parameters in the following way:

From 4099774da9ccb9ebb2c53c37e15d7f58b1a66a4d Mon Sep 17 00:00:00 2001
From: Ahmed Sobeh <ahmed.sobeh@aiven.io>
Date: Tue, 26 Mar 2024 22:24:37 +0100
Subject: [PATCH 209/258] KAFKA-16084: Simplify and deduplicate standalone
 herder test mocking (#15389)

Reviewers: Greg Harris <greg.harris@aiven.io>
---
 .../standalone/StandaloneHerderTest.java      | 340 +++++++-----------
 1 file changed, 136 insertions(+), 204 deletions(-)

diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
index 41091bbdf467d..bed08ffa2ebf7 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/standalone/StandaloneHerderTest.java
@@ -68,7 +68,6 @@
 import org.mockito.junit.MockitoJUnitRunner;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -115,6 +114,7 @@ public class StandaloneHerderTest {
     private static final String TOPICS_LIST_STR = "topic1,topic2";
     private static final String WORKER_ID = "localhost:8083";
     private static final String KAFKA_CLUSTER_ID = "I4ZmrWqfT2e-upky_4fdPA";
+    private static final Long WAIT_TIME_MS = 15000L;
 
     private enum SourceSink {
         SOURCE, SINK
@@ -122,30 +122,28 @@ private enum SourceSink {
 
     private StandaloneHerder herder;
 
-    private Connector connector;
     @Mock
     protected Worker worker;
-    @Mock protected WorkerConfigTransformer transformer;
-    @Mock private Plugins plugins;
+    @Mock
+    protected WorkerConfigTransformer transformer;
+    @Mock
+    private Plugins plugins;
     @Mock
     private PluginClassLoader pluginLoader;
     @Mock
     private LoaderSwap loaderSwap;
     protected FutureCallback<Herder.Created<ConnectorInfo>> createCallback;
-    @Mock protected StatusBackingStore statusBackingStore;
+    @Mock
+    protected StatusBackingStore statusBackingStore;
     private final SampleConnectorClientConfigOverridePolicy
-        noneConnectorClientConfigOverridePolicy = new SampleConnectorClientConfigOverridePolicy();
+            noneConnectorClientConfigOverridePolicy = new SampleConnectorClientConfigOverridePolicy();
 
     @Before
     public void setup() throws ExecutionException, InterruptedException {
-        worker = mock(Worker.class);
         herder = mock(StandaloneHerder.class, withSettings()
-            .useConstructor(worker, WORKER_ID, KAFKA_CLUSTER_ID, statusBackingStore, new MemoryConfigBackingStore(transformer), noneConnectorClientConfigOverridePolicy, new MockTime())
-            .defaultAnswer(CALLS_REAL_METHODS));
+                .useConstructor(worker, WORKER_ID, KAFKA_CLUSTER_ID, statusBackingStore, new MemoryConfigBackingStore(transformer), noneConnectorClientConfigOverridePolicy, new MockTime())
+                .defaultAnswer(CALLS_REAL_METHODS));
         createCallback = new FutureCallback<>();
-        plugins = mock(Plugins.class);
-        pluginLoader = mock(PluginClassLoader.class);
-        loaderSwap = mock(LoaderSwap.class);
         final ArgumentCaptor<Map<String, String>> configCapture = ArgumentCaptor.forClass(Map.class);
         when(transformer.transform(eq(CONNECTOR_NAME), configCapture.capture())).thenAnswer(invocation -> configCapture.getValue());
     }
@@ -153,26 +151,24 @@ public void setup() throws ExecutionException, InterruptedException {
     @After
     public void tearDown() {
         verifyNoMoreInteractions(worker, statusBackingStore);
+        herder.stop();
     }
 
     @Test
     public void testCreateSourceConnector() throws Exception {
-        connector = mock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SOURCE, config);
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
     }
 
     @Test
     public void testCreateConnectorFailedValidation() {
         // Basic validation should be performed and return an error, but should still evaluate the connector's config
-        connector = mock(BogusSourceConnector.class);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
         config.remove(ConnectorConfig.NAME_CONFIG);
@@ -193,7 +189,7 @@ public void testCreateConnectorFailedValidation() {
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
 
-        ExecutionException exception = assertThrows(ExecutionException.class, () -> createCallback.get(1000L, TimeUnit.SECONDS));
+        ExecutionException exception = assertThrows(ExecutionException.class, () -> createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         if (BadRequestException.class != exception.getCause().getClass()) {
             throw new AssertionError(exception.getCause());
         }
@@ -202,23 +198,21 @@ public void testCreateConnectorFailedValidation() {
 
     @Test
     public void testCreateConnectorAlreadyExists() throws Throwable {
-        connector = mock(BogusSourceConnector.class);
         // First addition should succeed
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config, config);
+        expectConfigValidation(SourceSink.SOURCE, config, config);
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         // Second should fail
         FutureCallback<Herder.Created<ConnectorInfo>> failedCreateCallback = new FutureCallback<>();
         // No new connector is created
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, failedCreateCallback);
-        ExecutionException exception = assertThrows(ExecutionException.class, () -> failedCreateCallback.get(1000L, TimeUnit.SECONDS));
+        ExecutionException exception = assertThrows(ExecutionException.class, () -> failedCreateCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         if (AlreadyExistsException.class != exception.getCause().getClass()) {
             throw new AssertionError(exception.getCause());
         }
@@ -227,59 +221,47 @@ public void testCreateConnectorAlreadyExists() throws Throwable {
 
     @Test
     public void testCreateSinkConnector() throws Exception {
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> config = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SINK, config);
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
     }
 
     @Test
     public void testCreateConnectorWithStoppedInitialState() throws Exception {
-        connector = mock(BogusSinkConnector.class);
         Map<String, String> config = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, false, config);
-        when(plugins.newConnector(anyString())).thenReturn(connectorMock);
+        expectConfigValidation(SourceSink.SINK, config);
 
         // Only the connector should be created; we expect no tasks to be spawned for a connector created with a paused or stopped initial state
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STOPPED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(config), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STOPPED), onStart.capture());
+        mockStartConnector(config, null, TargetState.STOPPED, null);
 
         when(worker.isRunning(CONNECTOR_NAME)).thenReturn(true);
         when(herder.connectorType(any())).thenReturn(ConnectorType.SINK);
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, TargetState.STOPPED, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(
-            new ConnectorInfo(CONNECTOR_NAME, connectorConfig(SourceSink.SINK), Collections.emptyList(), ConnectorType.SINK),
-            connectorInfo.result()
+                new ConnectorInfo(CONNECTOR_NAME, connectorConfig(SourceSink.SINK), Collections.emptyList(), ConnectorType.SINK),
+                connectorInfo.result()
         );
         verify(loaderSwap).close();
     }
 
     @Test
     public void testDestroyConnector() throws Exception {
-        connector = mock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SOURCE, config);
 
         when(statusBackingStore.getAll(CONNECTOR_NAME)).thenReturn(Collections.emptyList());
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         FutureCallback<Herder.Created<ConnectorInfo>> deleteCallback = new FutureCallback<>();
@@ -288,7 +270,7 @@ public void testDestroyConnector() throws Exception {
         verify(herder).onDeletion(CONNECTOR_NAME);
         verify(statusBackingStore).put(new TaskStatus(new ConnectorTaskId(CONNECTOR_NAME, 0), TaskStatus.State.DESTROYED, WORKER_ID, 0));
         verify(statusBackingStore).put(new ConnectorStatus(CONNECTOR_NAME, ConnectorStatus.State.DESTROYED, WORKER_ID, 0));
-        deleteCallback.get(1000L, TimeUnit.MILLISECONDS);
+        deleteCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
 
         // Second deletion should fail since the connector is gone
         FutureCallback<Herder.Created<ConnectorInfo>> failedDeleteCallback = new FutureCallback<>();
@@ -296,7 +278,7 @@ public void testDestroyConnector() throws Exception {
         ExecutionException e = assertThrows(
                 "Should have thrown NotFoundException",
                 ExecutionException.class,
-                () -> failedDeleteCallback.get(1000L, TimeUnit.MILLISECONDS)
+                () -> failedDeleteCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS)
         );
         assertInstanceOf(NotFoundException.class, e.getCause());
     }
@@ -306,29 +288,23 @@ public void testRestartConnectorSameTaskConfigs() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SOURCE, config);
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(config), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(config, TargetState.STARTED, TargetState.STARTED, null);
 
         when(worker.connectorNames()).thenReturn(Collections.singleton(CONNECTOR_NAME));
         when(worker.getPlugins()).thenReturn(plugins);
         // same task configs as earlier, so don't expect a new set of tasks to be brought up
         when(worker.connectorTaskConfigs(CONNECTOR_NAME, new SourceConnectorConfig(plugins, config, true)))
-            .thenReturn(Collections.singletonList(taskConfig(SourceSink.SOURCE)));
+                .thenReturn(Collections.singletonList(taskConfig(SourceSink.SOURCE)));
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         FutureCallback<Void> restartCallback = new FutureCallback<>();
         herder.restartConnector(CONNECTOR_NAME, restartCallback);
-        restartCallback.get(1000L, TimeUnit.MILLISECONDS);
+        restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         verify(worker).stopAndAwaitConnector(eq(CONNECTOR_NAME));
     }
 
@@ -338,19 +314,13 @@ public void testRestartConnectorNewTaskConfigs() throws Exception {
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
         ConnectorTaskId taskId = new ConnectorTaskId(CONNECTOR_NAME, 0);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SOURCE, config);
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(config), any(HerderConnectorContext.class),
-                eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(config, TargetState.STARTED, TargetState.STARTED, null);
 
         when(worker.connectorNames()).thenReturn(Collections.singleton(CONNECTOR_NAME));
         when(worker.getPlugins()).thenReturn(plugins);
@@ -366,7 +336,7 @@ public void testRestartConnectorNewTaskConfigs() throws Exception {
         expectStop();
         herder.restartConnector(CONNECTOR_NAME, restartCallback);
         verify(statusBackingStore).put(new TaskStatus(taskId, TaskStatus.State.DESTROYED, WORKER_ID, 0));
-        restartCallback.get(1000L, TimeUnit.MILLISECONDS);
+        restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
     }
 
     @Test
@@ -374,28 +344,21 @@ public void testRestartConnectorFailureOnStart() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
+        expectConfigValidation(SourceSink.SOURCE, config);
 
         doNothing().when(worker).stopAndAwaitConnector(CONNECTOR_NAME);
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
 
         Exception exception = new ConnectException("Failed to start connector");
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
-
         FutureCallback<Void> restartCallback = new FutureCallback<>();
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(exception, null);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(config), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(config, null, TargetState.STARTED, exception);
         herder.restartConnector(CONNECTOR_NAME, restartCallback);
         try {
-            restartCallback.get(1000L, TimeUnit.MILLISECONDS);
+            restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
             fail();
         } catch (ExecutionException e) {
             assertEquals(exception, e.getCause());
@@ -408,8 +371,8 @@ public void testRestartTask() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+
+        expectConfigValidation(SourceSink.SOURCE, connectorConfig);
 
         doNothing().when(worker).stopAndAwaitTask(taskId);
 
@@ -426,16 +389,16 @@ public void testRestartTask() throws Exception {
                 new HashSet<>(),
                 transformer);
         when(worker.startSourceTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED))
-            .thenReturn(true);
+                .thenReturn(true);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
-        createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
+        createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         FutureCallback<Void> restartTaskCallback = new FutureCallback<>();
         herder.restartTask(taskId, restartTaskCallback);
-        restartTaskCallback.get(1000L, TimeUnit.MILLISECONDS);
+        restartTaskCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
     }
 
     @Test
@@ -444,8 +407,7 @@ public void testRestartTaskFailureOnStart() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SOURCE, connectorConfig);
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
@@ -460,17 +422,17 @@ public void testRestartTaskFailureOnStart() throws Exception {
                 new HashSet<>(),
                 transformer);
         when(worker.startSourceTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED))
-            .thenReturn(false);
+                .thenReturn(false);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.MILLISECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         FutureCallback<Void> cb = new FutureCallback<>();
         herder.restartTask(taskId, cb);
         verify(worker).stopAndAwaitTask(taskId);
         try {
-            cb.get(1000L, TimeUnit.MILLISECONDS);
+            cb.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
             fail("Expected restart callback to raise an exception");
         } catch (ExecutionException exception) {
             assertEquals(ConnectException.class, exception.getCause().getClass());
@@ -482,7 +444,7 @@ public void testRestartConnectorAndTasksUnknownConnector() {
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         RestartRequest restartRequest = new RestartRequest("UnknownConnector", false, true);
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(NotFoundException.class, ee.getCause());
     }
 
@@ -491,20 +453,18 @@ public void testRestartConnectorAndTasksNoStatus() throws Exception {
         RestartRequest restartRequest = new RestartRequest(CONNECTOR_NAME, false, true);
         doReturn(Optional.empty()).when(herder).buildRestartPlan(restartRequest);
 
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SINK, connectorConfig);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        ExecutionException ee = assertThrows(ExecutionException.class, () -> restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(NotFoundException.class, ee.getCause());
         assertTrue(ee.getMessage().contains("Status for connector"));
     }
@@ -519,20 +479,18 @@ public void testRestartConnectorAndTasksNoRestarts() throws Exception {
         when(restartPlan.restartConnectorStateInfo()).thenReturn(connectorStateInfo);
         doReturn(Optional.of(restartPlan)).when(herder).buildRestartPlan(restartRequest);
 
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SINK, connectorConfig);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        assertEquals(connectorStateInfo, restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        assertEquals(connectorStateInfo, restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
     }
 
     @Test
@@ -545,29 +503,22 @@ public void testRestartConnectorAndTasksOnlyConnector() throws Exception {
         when(restartPlan.restartConnectorStateInfo()).thenReturn(connectorStateInfo);
         doReturn(Optional.of(restartPlan)).when(herder).buildRestartPlan(restartRequest);
 
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SINK, connectorConfig);
 
         doNothing().when(worker).stopAndAwaitConnector(CONNECTOR_NAME);
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(connectorConfig), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(connectorConfig, null, TargetState.STARTED, null);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        assertEquals(connectorStateInfo, restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        assertEquals(connectorStateInfo, restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
 
         verifyConnectorStatusRestart();
     }
@@ -586,12 +537,10 @@ public void testRestartConnectorAndTasksOnlyTasks() throws Exception {
         when(restartPlan.restartConnectorStateInfo()).thenReturn(connectorStateInfo);
         doReturn(Optional.of(restartPlan)).when(herder).buildRestartPlan(restartRequest);
 
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SINK, connectorConfig);
 
         doNothing().when(worker).stopAndAwaitTasks(Collections.singletonList(taskId));
 
@@ -608,15 +557,15 @@ public void testRestartConnectorAndTasksOnlyTasks() throws Exception {
                 new HashSet<>(),
                 transformer);
         when(worker.startSinkTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SINK), herder, TargetState.STARTED))
-            .thenReturn(true);
+                .thenReturn(true);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        assertEquals(connectorStateInfo, restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        assertEquals(connectorStateInfo, restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         ArgumentCaptor<TaskStatus> taskStatus = ArgumentCaptor.forClass(TaskStatus.class);
         verify(statusBackingStore).put(taskStatus.capture());
         assertEquals(AbstractStatus.State.RESTARTING, taskStatus.getValue().state());
@@ -639,22 +588,15 @@ public void testRestartConnectorAndTasksBoth() throws Exception {
 
         ArgumentCaptor<TaskStatus> taskStatus = ArgumentCaptor.forClass(TaskStatus.class);
 
-        connector = mock(BogusSinkConnector.class);
         expectAdd(SourceSink.SINK);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SINK);
-        Connector connectorMock = mock(SinkConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SINK, connectorConfig);
 
         doNothing().when(worker).stopAndAwaitConnector(CONNECTOR_NAME);
         doNothing().when(worker).stopAndAwaitTasks(Collections.singletonList(taskId));
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(connectorConfig), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(connectorConfig, null, TargetState.STARTED, null);
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
@@ -669,15 +611,15 @@ public void testRestartConnectorAndTasksBoth() throws Exception {
                 new HashSet<>(),
                 transformer);
         when(worker.startSinkTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SINK), herder, TargetState.STARTED))
-            .thenReturn(true);
+                .thenReturn(true);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         FutureCallback<ConnectorStateInfo> restartCallback = new FutureCallback<>();
         herder.restartConnectorAndTasks(restartRequest, restartCallback);
-        assertEquals(connectorStateInfo, restartCallback.get(1000L, TimeUnit.MILLISECONDS));
+        assertEquals(connectorStateInfo, restartCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
 
         verifyConnectorStatusRestart();
         verify(statusBackingStore).put(taskStatus.capture());
@@ -687,17 +629,15 @@ public void testRestartConnectorAndTasksBoth() throws Exception {
 
     @Test
     public void testCreateAndStop() throws Exception {
-        connector = mock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SOURCE, connectorConfig);
 
         expectStop();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         // herder.stop() should stop any running connectors and tasks even if destroyConnector was not invoked
@@ -726,22 +666,21 @@ public void testAccessors() throws Exception {
         doNothing().when(tasksConfigCb).onCompletion(any(NotFoundException.class), isNull());
         doNothing().when(connectorConfigCb).onCompletion(any(NotFoundException.class), isNull());
 
-        // Create connector
-        connector = mock(BogusSourceConnector.class);
+
         expectAdd(SourceSink.SOURCE);
-        expectConfigValidation(connector, true, connConfig);
+        expectConfigValidation(SourceSink.SOURCE, connConfig);
 
         // Validate accessors with 1 connector
         doNothing().when(listConnectorsCb).onCompletion(null, singleton(CONNECTOR_NAME));
-        ConnectorInfo connInfo = new ConnectorInfo(CONNECTOR_NAME, connConfig, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
-            ConnectorType.SOURCE);
+        ConnectorInfo connInfo = new ConnectorInfo(CONNECTOR_NAME, connConfig, singletonList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
+                ConnectorType.SOURCE);
         doNothing().when(connectorInfoCb).onCompletion(null, connInfo);
 
         TaskInfo taskInfo = new TaskInfo(new ConnectorTaskId(CONNECTOR_NAME, 0), taskConfig(SourceSink.SOURCE));
-        doNothing().when(taskConfigsCb).onCompletion(null, Arrays.asList(taskInfo));
+        doNothing().when(taskConfigsCb).onCompletion(null, singletonList(taskInfo));
 
         Map<ConnectorTaskId, Map<String, String>> tasksConfig = Collections.singletonMap(new ConnectorTaskId(CONNECTOR_NAME, 0),
-            taskConfig(SourceSink.SOURCE));
+                taskConfig(SourceSink.SOURCE));
         doNothing().when(tasksConfigCb).onCompletion(null, tasksConfig);
 
         // All operations are synchronous for StandaloneHerder, so we don't need to actually wait after making each call
@@ -752,7 +691,7 @@ public void testAccessors() throws Exception {
         herder.tasksConfig(CONNECTOR_NAME, tasksConfigCb);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         reset(transformer);
@@ -773,11 +712,9 @@ public void testPutConnectorConfig() throws Exception {
 
         Callback<Map<String, String>> connectorConfigCb = mock(Callback.class);
 
-        // Create
-        connector = mock(BogusSourceConnector.class);
+
         expectAdd(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, connConfig);
+        expectConfigValidation(SourceSink.SOURCE, connConfig, newConnConfig);
 
         // Should get first config
         doNothing().when(connectorConfigCb).onCompletion(null, connConfig);
@@ -789,15 +726,14 @@ public void testPutConnectorConfig() throws Exception {
             onStart.getValue().onCompletion(null, TargetState.STARTED);
             return true;
         }).when(worker).startConnector(eq(CONNECTOR_NAME), capturedConfig.capture(), any(),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+                eq(herder), eq(TargetState.STARTED), onStart.capture());
         // Generate same task config, which should result in no additional action to restart tasks
         when(worker.connectorTaskConfigs(CONNECTOR_NAME, new SourceConnectorConfig(plugins, newConnConfig, true)))
-            .thenReturn(singletonList(taskConfig(SourceSink.SOURCE)));
+                .thenReturn(singletonList(taskConfig(SourceSink.SOURCE)));
 
-        expectConfigValidation(connectorMock, false, newConnConfig);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         herder.connectorConfig(CONNECTOR_NAME, connectorConfigCb);
@@ -806,8 +742,8 @@ public void testPutConnectorConfig() throws Exception {
         doNothing().when(connectorConfigCb).onCompletion(null, newConnConfig);
         herder.putConnectorConfig(CONNECTOR_NAME, newConnConfig, true, reconfigureCallback);
         Herder.Created<ConnectorInfo> newConnectorInfo = reconfigureCallback.get(1000L, TimeUnit.SECONDS);
-        ConnectorInfo newConnInfo = new ConnectorInfo(CONNECTOR_NAME, newConnConfig, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
-            ConnectorType.SOURCE);
+        ConnectorInfo newConnInfo = new ConnectorInfo(CONNECTOR_NAME, newConnConfig, singletonList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
+                ConnectorType.SOURCE);
         assertEquals(newConnInfo, newConnectorInfo.result());
 
         assertEquals("bar", capturedConfig.getValue().get("foo"));
@@ -834,9 +770,9 @@ public void testCorruptConfig() {
         List<String> errors = new ArrayList<>(singletonList(error));
         String key = "foo.invalid.key";
         when(connectorMock.validate(config)).thenReturn(
-            new Config(
-                Arrays.asList(new ConfigValue(key, null, Collections.emptyList(), errors))
-            )
+                new Config(
+                        singletonList(new ConfigValue(key, null, Collections.emptyList(), errors))
+                )
         );
         ConfigDef configDef = new ConfigDef();
         configDef.define(key, ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "");
@@ -853,7 +789,7 @@ public void testCorruptConfig() {
         ExecutionException e = assertThrows(
                 "Should have failed to configure connector",
                 ExecutionException.class,
-                () -> createCallback.get(1000L, TimeUnit.SECONDS)
+                () -> createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS)
         );
         assertNotNull(e.getCause());
         Throwable cause = e.getCause();
@@ -869,12 +805,10 @@ public void testCorruptConfig() {
 
     @Test
     public void testTargetStates() throws Exception {
-        connector = mock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
 
         Map<String, String> connectorConfig = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, connectorConfig);
+        expectConfigValidation(SourceSink.SOURCE, connectorConfig);
 
         // We pause, then stop, the connector
         expectTargetState(CONNECTOR_NAME, TargetState.PAUSED);
@@ -887,14 +821,14 @@ public void testTargetStates() throws Exception {
         FutureCallback<List<TaskInfo>> taskConfigsCallback = new FutureCallback<>();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
         herder.pauseConnector(CONNECTOR_NAME);
         herder.stopConnector(CONNECTOR_NAME, stopCallback);
         verify(statusBackingStore).put(new TaskStatus(new ConnectorTaskId(CONNECTOR_NAME, 0), AbstractStatus.State.DESTROYED, WORKER_ID, 0));
-        stopCallback.get(10L, TimeUnit.SECONDS);
+        stopCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         herder.taskConfigs(CONNECTOR_NAME, taskConfigsCallback);
-        assertEquals(Collections.emptyList(), taskConfigsCallback.get(1, TimeUnit.SECONDS));
+        assertEquals(Collections.emptyList(), taskConfigsCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
 
         // herder.stop() should stop any running connectors and tasks even if destroyConnector was not invoked
         herder.stop();
@@ -910,12 +844,12 @@ public void testModifyConnectorOffsetsUnknownConnector() {
         herder.alterConnectorOffsets("unknown-connector",
                 Collections.singletonMap(Collections.singletonMap("partitionKey", "partitionValue"), Collections.singletonMap("offsetKey", "offsetValue")),
                 alterOffsetsCallback);
-        ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
+        ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(NotFoundException.class, e.getCause());
 
         FutureCallback<Message> resetOffsetsCallback = new FutureCallback<>();
         herder.resetConnectorOffsets("unknown-connector", resetOffsetsCallback);
-        e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
+        e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(NotFoundException.class, e.getCause());
     }
 
@@ -938,12 +872,12 @@ public void testModifyConnectorOffsetsConnectorNotInStoppedState() {
         herder.alterConnectorOffsets(CONNECTOR_NAME,
                 Collections.singletonMap(Collections.singletonMap("partitionKey", "partitionValue"), Collections.singletonMap("offsetKey", "offsetValue")),
                 alterOffsetsCallback);
-        ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
+        ExecutionException e = assertThrows(ExecutionException.class, () -> alterOffsetsCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(BadRequestException.class, e.getCause());
 
         FutureCallback<Message> resetOffsetsCallback = new FutureCallback<>();
         herder.resetConnectorOffsets(CONNECTOR_NAME, resetOffsetsCallback);
-        e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(1000L, TimeUnit.MILLISECONDS));
+        e = assertThrows(ExecutionException.class, () -> resetOffsetsCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS));
         assertInstanceOf(BadRequestException.class, e.getCause());
     }
 
@@ -1004,33 +938,25 @@ public void testResetConnectorOffsets() throws Exception {
 
     @Test()
     public void testRequestTaskReconfigurationDoesNotDeadlock() throws Exception {
-        connector = mock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
 
         // Start the connector
         Map<String, String> config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = mock(SourceConnector.class);
-        expectConfigValidation(connectorMock, true, config);
-
+        // Prepare for connector and task config update
+        Map<String, String> newConfig = connectorConfig(SourceSink.SOURCE);
+        newConfig.put("dummy-connector-property", "yes");
+        expectConfigValidation(SourceSink.SOURCE, config, newConfig);
+        mockStartConnector(newConfig, TargetState.STARTED, TargetState.STARTED, null);
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
 
         // Wait on connector to start
-        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(1000L, TimeUnit.MILLISECONDS);
+        Herder.Created<ConnectorInfo> connectorInfo = createCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
         assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         // Prepare for task config update
         when(worker.connectorNames()).thenReturn(Collections.singleton(CONNECTOR_NAME));
         expectStop();
 
-        // Prepare for connector and task config update
-        Map<String, String> newConfig = connectorConfig(SourceSink.SOURCE);
-        newConfig.put("dummy-connector-property", "yes");
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(newConfig), any(HerderConnectorContext.class),
-                eq(herder), eq(TargetState.STARTED), onStart.capture());
 
         // Common invocations
         when(worker.startSourceTask(eq(new ConnectorTaskId(CONNECTOR_NAME, 0)), any(), any(), any(), eq(herder), eq(TargetState.STARTED))).thenReturn(true);
@@ -1045,15 +971,14 @@ public void testRequestTaskReconfigurationDoesNotDeadlock() throws Exception {
 
         // Set new config on the connector and tasks
         FutureCallback<Herder.Created<ConnectorInfo>> reconfigureCallback = new FutureCallback<>();
-        expectConfigValidation(connectorMock, false, newConfig);
         herder.putConnectorConfig(CONNECTOR_NAME, newConfig, true, reconfigureCallback);
 
         // Reconfigure the tasks
         herder.requestTaskReconfiguration(CONNECTOR_NAME);
 
         // Wait on connector update
-        Herder.Created<ConnectorInfo> updatedConnectorInfo = reconfigureCallback.get(1000L, TimeUnit.MILLISECONDS);
-        ConnectorInfo expectedConnectorInfo = new ConnectorInfo(CONNECTOR_NAME, newConfig, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)), ConnectorType.SOURCE);
+        Herder.Created<ConnectorInfo> updatedConnectorInfo = reconfigureCallback.get(WAIT_TIME_MS, TimeUnit.MILLISECONDS);
+        ConnectorInfo expectedConnectorInfo = new ConnectorInfo(CONNECTOR_NAME, newConfig, singletonList(new ConnectorTaskId(CONNECTOR_NAME, 0)), ConnectorType.SOURCE);
         assertEquals(expectedConnectorInfo, updatedConnectorInfo.result());
 
         verify(statusBackingStore, times(2)).put(new TaskStatus(new ConnectorTaskId(CONNECTOR_NAME, 0), TaskStatus.State.DESTROYED, WORKER_ID, 0));
@@ -1062,15 +987,10 @@ public void testRequestTaskReconfigurationDoesNotDeadlock() throws Exception {
     private void expectAdd(SourceSink sourceSink) {
         Map<String, String> connectorProps = connectorConfig(sourceSink);
         ConnectorConfig connConfig = sourceSink == SourceSink.SOURCE ?
-            new SourceConnectorConfig(plugins, connectorProps, true) :
-            new SinkConnectorConfig(plugins, connectorProps);
+                new SourceConnectorConfig(plugins, connectorProps, true) :
+                new SinkConnectorConfig(plugins, connectorProps);
 
-        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
-        doAnswer(invocation -> {
-            onStart.getValue().onCompletion(null, TargetState.STARTED);
-            return true;
-        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(connectorProps), any(HerderConnectorContext.class),
-            eq(herder), eq(TargetState.STARTED), onStart.capture());
+        mockStartConnector(connectorProps, TargetState.STARTED, TargetState.STARTED, null);
         when(worker.isRunning(CONNECTOR_NAME)).thenReturn(true);
         if (sourceSink == SourceSink.SOURCE) {
             when(worker.isTopicCreationEnabled()).thenReturn(true);
@@ -1081,7 +1001,7 @@ private void expectAdd(SourceSink sourceSink) {
         Map<String, String> generatedTaskProps = taskConfig(sourceSink);
 
         when(worker.connectorTaskConfigs(CONNECTOR_NAME, connConfig))
-            .thenReturn(singletonList(generatedTaskProps));
+                .thenReturn(singletonList(generatedTaskProps));
 
         ClusterConfigState configState = new ClusterConfigState(
                 -1,
@@ -1127,8 +1047,8 @@ private void expectTargetState(String connector, TargetState state) {
 
     private ConnectorInfo createdInfo(SourceSink sourceSink) {
         return new ConnectorInfo(CONNECTOR_NAME, connectorConfig(sourceSink),
-            Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
-            SourceSink.SOURCE == sourceSink ? ConnectorType.SOURCE : ConnectorType.SINK);
+                singletonList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
+                SourceSink.SOURCE == sourceSink ? ConnectorType.SOURCE : ConnectorType.SINK);
     }
 
     private void expectStop() {
@@ -1168,38 +1088,40 @@ private static Map<String, String> taskConfig(SourceSink sourceSink) {
     }
 
     private void expectConfigValidation(
-            Connector connectorMock,
-            boolean shouldCreateConnector,
+            SourceSink sourceSink,
             Map<String, String>... configs
     ) {
         // config validation
+        Connector connectorMock = sourceSink == SourceSink.SOURCE ? mock(SourceConnector.class) : mock(SinkConnector.class);
         when(worker.configTransformer()).thenReturn(transformer);
         final ArgumentCaptor<Map<String, String>> configCapture = ArgumentCaptor.forClass(Map.class);
         when(transformer.transform(configCapture.capture())).thenAnswer(invocation -> configCapture.getValue());
         when(worker.getPlugins()).thenReturn(plugins);
         when(plugins.connectorLoader(anyString())).thenReturn(pluginLoader);
         when(plugins.withClassLoader(pluginLoader)).thenReturn(loaderSwap);
-        if (shouldCreateConnector) {
-            when(worker.getPlugins()).thenReturn(plugins);
-            when(plugins.newConnector(anyString())).thenReturn(connectorMock);
-        }
+
+        // Assume the connector should always be created
+        when(worker.getPlugins()).thenReturn(plugins);
+        when(plugins.newConnector(anyString())).thenReturn(connectorMock);
         when(connectorMock.config()).thenReturn(new ConfigDef());
 
-        for (Map<String, String> config : configs)
+        // Set up validation for each config
+        for (Map<String, String> config : configs) {
             when(connectorMock.validate(config)).thenReturn(new Config(Collections.emptyList()));
+        }
     }
 
     // We need to use a real class here due to some issue with mocking java.lang.Class
-    private abstract class BogusSourceConnector extends SourceConnector {
+    private static abstract class BogusSourceConnector extends SourceConnector {
     }
 
-    private abstract class BogusSourceTask extends SourceTask {
+    private static abstract class BogusSourceTask extends SourceTask {
     }
 
-    private abstract class BogusSinkConnector extends SinkConnector {
+    private static abstract class BogusSinkConnector extends SinkConnector {
     }
 
-    private abstract class BogusSinkTask extends SourceTask {
+    private static abstract class BogusSinkTask extends SourceTask {
     }
 
     private void verifyConnectorStatusRestart() {
@@ -1208,4 +1130,14 @@ private void verifyConnectorStatusRestart() {
         assertEquals(CONNECTOR_NAME, connectorStatus.getValue().id());
         assertEquals(AbstractStatus.State.RESTARTING, connectorStatus.getValue().state());
     }
+
+    private void mockStartConnector(Map<String, String> config, TargetState result, TargetState targetState, Exception exception) {
+        final ArgumentCaptor<Callback<TargetState>> onStart = ArgumentCaptor.forClass(Callback.class);
+        doAnswer(invocation -> {
+            onStart.getValue().onCompletion(exception, result);
+            return true;
+        }).when(worker).startConnector(eq(CONNECTOR_NAME), eq(config),
+                any(HerderConnectorContext.class),
+                eq(herder), eq(targetState), onStart.capture());
+    }
 }

From 8d914b543d5d23031fe178d424f45789eaa8d1fc Mon Sep 17 00:00:00 2001
From: Colin Patrick McCabe <cmccabe@apache.org>
Date: Tue, 26 Mar 2024 16:49:38 -0700
Subject: [PATCH 210/258] KAFKA-16411: Correctly migrate default client quota
 entities (#15584)

KAFKA-16222 fixed a bug whereby we didn't undo the name sanitization used on client quota entity names
stored in ZooKeeper. However, it incorrectly claimed to fix the handling of default client quota
entities. It also failed to correctly re-sanitize when syncronizing the data back to ZooKeeper.

This PR fixes ZkConfigMigrationClient to do the sanitization correctly on both the read and write
paths. We do de-sanitization before invoking the visitors, since after all it does not make sense to
do the same de-sanitization step in each and every visitor.

Additionally, this PR fixes a bug causing default entities to be converted incorrectly. For example,
ClientQuotaEntity(user -> null) is stored under the /config/users/<default> znode in ZooKeeper. In
KRaft it appears as a ClientQuotaRecord with EntityData(entityType=users, entityName=null).
Prior to this PR, this was being converted to a ClientQuotaRecord with EntityData(entityType=users,
entityName=""). That represents a quota on the user whose name is the empty string (yes, we allow
users to name themselves with the empty string, sadly.)

The confusion appears to have arisen because for TOPIC and BROKER configurations, the default
ConfigResource is indeed the one named with the empty (not null) string. For example, the default
topic configuration resource is ConfigResource(name="", type=TOPIC).  However, things are different
for client quotas. Default client quota entities in KRaft (and also in AdminClient) are represented
by maps with null values. For example, the default User entity is represented by Map("user" ->
null).  In retrospect, using a map with null values was a poor choice; a Map<String,
Optional<String>> would have made more sense. However, this is the way the API currently is and we
have to convert correctly.

There was an additional level of confusion present in KAFKA-16222 where someone thought that using
the ZooKeeper placeholder string "<default>" in the AdminClient API would yield a default client
quota entity. Thise seems to have been suggested by the ConfigEntityName class that was created
recently. In fact, <default> is not part of any public API in Kafka. Accordingly, this PR also
renames ConfigEntityName.DEFAULT to ZooKeeperInternals.DEFAULT_STRING, to make it clear that the
string <default> is just a detail of the ZooKeeper implementation.  It is not used in the Kafka API
to indicate defaults. Hopefully this will avoid confusion in the future.

Finally, the PR also creates KRaftClusterTest.testDefaultClientQuotas to get extra test coverage of
setting default client quotas.

Reviewers: Manikumar Reddy <manikumar.reddy@gmail.com>, Igor Soarez <soarez@apple.com>
---
 .../scala/kafka/admin/ConfigCommand.scala     | 12 +--
 .../kafka/server/ClientQuotaManager.scala     | 13 ++-
 .../scala/kafka/server/ConfigHandler.scala    |  8 +-
 .../kafka/server/DynamicBrokerConfig.scala    |  4 +-
 .../scala/kafka/server/DynamicConfig.scala    |  4 +-
 .../scala/kafka/server/ZkAdminManager.scala   |  8 +-
 .../metadata/ClientQuotaMetadataManager.scala | 14 +--
 .../metadata/DynamicConfigPublisher.scala     |  4 +-
 .../server/metadata/ZkConfigRepository.scala  |  4 +-
 .../main/scala/kafka/zk/AdminZkClient.scala   |  8 +-
 .../scala/kafka/zk/ZkMigrationClient.scala    |  4 -
 .../migration/ZkConfigMigrationClient.scala   | 99 +++++++++++++------
 .../admin/ConfigCommandIntegrationTest.scala  |  4 +-
 .../kafka/server/KRaftClusterTest.scala       | 65 +++++++++++-
 .../kafka/zk/ZkMigrationIntegrationTest.scala | 46 +++++----
 .../unit/kafka/admin/ConfigCommandTest.scala  | 12 +--
 .../kafka/server/ClientQuotaManagerTest.scala | 42 ++++----
 .../server/ClientQuotasRequestTest.scala      |  5 +-
 .../server/DynamicConfigChangeTest.scala      |  4 +-
 .../ZkConfigMigrationClientTest.scala         | 22 +++--
 ...ntityName.java => ZooKeeperInternals.java} | 10 +-
 21 files changed, 251 insertions(+), 141 deletions(-)
 rename server-common/src/main/java/org/apache/kafka/server/config/{ConfigEntityName.java => ZooKeeperInternals.java} (63%)

diff --git a/core/src/main/scala/kafka/admin/ConfigCommand.scala b/core/src/main/scala/kafka/admin/ConfigCommand.scala
index bcbeedc468e4e..d03d7dfb9b52a 100644
--- a/core/src/main/scala/kafka/admin/ConfigCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConfigCommand.scala
@@ -35,7 +35,7 @@ import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity,
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter, ScramMechanism}
 import org.apache.kafka.common.utils.{Sanitizer, Time, Utils}
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.kafka.security.{PasswordEncoder, PasswordEncoderConfigs}
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 import org.apache.kafka.storage.internals.log.LogConfig
@@ -153,7 +153,7 @@ object ConfigCommand extends Logging {
       if (!configsToBeAdded.isEmpty || configsToBeDeleted.nonEmpty) {
         validateBrokersNotRunning(entityName, adminZkClient, zkClient, errorMessage)
 
-        val perBrokerConfig = entityName != ConfigEntityName.DEFAULT
+        val perBrokerConfig = entityName != ZooKeeperInternals.DEFAULT_STRING
         preProcessBrokerConfigs(configsToBeAdded, perBrokerConfig)
       }
     }
@@ -178,7 +178,7 @@ object ConfigCommand extends Logging {
                                         adminZkClient: AdminZkClient,
                                         zkClient: KafkaZkClient,
                                         errorMessage: String): Unit = {
-    val perBrokerConfig = entityName != ConfigEntityName.DEFAULT
+    val perBrokerConfig = entityName != ZooKeeperInternals.DEFAULT_STRING
     val info = "Broker configuration operations using ZooKeeper are only supported if the affected broker(s) are not running."
     if (perBrokerConfig) {
       adminZkClient.parseBroker(entityName).foreach { brokerId =>
@@ -697,7 +697,7 @@ object ConfigCommand extends Logging {
         case t => t
       }
       sanitizedName match {
-        case Some(ConfigEntityName.DEFAULT) => "default " + typeName
+        case Some(ZooKeeperInternals.DEFAULT_STRING) => "default " + typeName
         case Some(n) =>
           val desanitized = if (entityType == ConfigType.USER || entityType == ConfigType.CLIENT) Sanitizer.desanitize(n) else n
           s"$typeName '$desanitized'"
@@ -758,7 +758,7 @@ object ConfigCommand extends Logging {
     else {
       // Exactly one entity type and at-most one entity name expected for other entities
       val name = entityNames.headOption match {
-        case Some("") => Some(ConfigEntityName.DEFAULT)
+        case Some("") => Some(ZooKeeperInternals.DEFAULT_STRING)
         case v => v
       }
       ConfigEntity(Entity(entityTypes.head, name), None)
@@ -775,7 +775,7 @@ object ConfigCommand extends Logging {
 
     def sanitizeName(entityType: String, name: String) = {
       if (name.isEmpty)
-        ConfigEntityName.DEFAULT
+        ZooKeeperInternals.DEFAULT_STRING
       else {
         entityType match {
           case ConfigType.USER | ConfigType.CLIENT => Sanitizer.sanitize(name)
diff --git a/core/src/main/scala/kafka/server/ClientQuotaManager.scala b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
index a10832d969071..a8c6a2fc557a1 100644
--- a/core/src/main/scala/kafka/server/ClientQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
@@ -19,7 +19,6 @@ package kafka.server
 import java.{lang, util}
 import java.util.concurrent.{ConcurrentHashMap, DelayQueue, TimeUnit}
 import java.util.concurrent.locks.ReentrantReadWriteLock
-
 import kafka.network.RequestChannel
 import kafka.server.ClientQuotaManager._
 import kafka.utils.{Logging, QuotaUtils}
@@ -29,7 +28,7 @@ import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.metrics.stats.{Avg, CumulativeSum, Rate}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{Sanitizer, Time}
-import org.apache.kafka.server.config.{ConfigEntityName, ClientQuotaManagerConfig}
+import org.apache.kafka.server.config.{ClientQuotaManagerConfig, ZooKeeperInternals}
 import org.apache.kafka.server.quota.{ClientQuotaCallback, ClientQuotaEntity, ClientQuotaType}
 import org.apache.kafka.server.util.ShutdownableThread
 import org.apache.kafka.network.Session
@@ -76,13 +75,13 @@ object ClientQuotaManager {
 
   case object DefaultUserEntity extends BaseUserEntity {
     override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.DEFAULT_USER
-    override def name: String = ConfigEntityName.DEFAULT
+    override def name: String = ZooKeeperInternals.DEFAULT_STRING
     override def toString: String = "default user"
   }
 
   case object DefaultClientIdEntity extends ClientQuotaEntity.ConfigEntity {
     override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.DEFAULT_CLIENT_ID
-    override def name: String = ConfigEntityName.DEFAULT
+    override def name: String = ZooKeeperInternals.DEFAULT_STRING
     override def toString: String = "default client-id"
   }
 
@@ -93,7 +92,7 @@ object ClientQuotaManager {
 
     def sanitizedUser: String = userEntity.map {
       case entity: UserEntity => entity.sanitizedUser
-      case DefaultUserEntity => ConfigEntityName.DEFAULT
+      case DefaultUserEntity => ZooKeeperInternals.DEFAULT_STRING
     }.getOrElse("")
 
     def clientId: String = clientIdEntity.map(_.name).getOrElse("")
@@ -419,11 +418,11 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
     lock.writeLock().lock()
     try {
       val userEntity = sanitizedUser.map {
-        case ConfigEntityName.DEFAULT => DefaultUserEntity
+        case ZooKeeperInternals.DEFAULT_STRING => DefaultUserEntity
         case user => UserEntity(user)
       }
       val clientIdEntity = sanitizedClientId.map {
-        case ConfigEntityName.DEFAULT => DefaultClientIdEntity
+        case ZooKeeperInternals.DEFAULT_STRING => DefaultClientIdEntity
         case _ => ClientIdEntity(clientId.getOrElse(throw new IllegalStateException("Client-id not provided")))
       }
       val quotaEntity = KafkaQuotaEntity(userEntity, clientIdEntity)
diff --git a/core/src/main/scala/kafka/server/ConfigHandler.scala b/core/src/main/scala/kafka/server/ConfigHandler.scala
index 66e1c2394e85f..05db9b3bc15a9 100644
--- a/core/src/main/scala/kafka/server/ConfigHandler.scala
+++ b/core/src/main/scala/kafka/server/ConfigHandler.scala
@@ -34,7 +34,7 @@ import org.apache.kafka.common.metrics.Quota
 import org.apache.kafka.common.metrics.Quota._
 import org.apache.kafka.common.utils.Sanitizer
 import org.apache.kafka.server.ClientMetricsManager
-import org.apache.kafka.server.config.ConfigEntityName
+import org.apache.kafka.server.config.ZooKeeperInternals
 import org.apache.kafka.storage.internals.log.{LogConfig, ThrottledReplicaListValidator}
 import org.apache.kafka.storage.internals.log.LogConfig.MessageFormatVersion
 
@@ -208,7 +208,7 @@ class UserConfigHandler(private val quotaManagers: QuotaManagers, val credential
     val sanitizedUser = entities(0)
     val sanitizedClientId = if (entities.length == 3) Some(entities(2)) else None
     updateQuotaConfig(Some(sanitizedUser), sanitizedClientId, config)
-    if (sanitizedClientId.isEmpty && sanitizedUser != ConfigEntityName.DEFAULT)
+    if (sanitizedClientId.isEmpty && sanitizedUser != ZooKeeperInternals.DEFAULT_STRING)
       credentialProvider.updateCredentials(Sanitizer.desanitize(sanitizedUser), config)
   }
 }
@@ -218,7 +218,7 @@ class IpConfigHandler(private val connectionQuotas: ConnectionQuotas) extends Co
   def processConfigChanges(ip: String, config: Properties): Unit = {
     val ipConnectionRateQuota = Option(config.getProperty(QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG)).map(_.toInt)
     val updatedIp = {
-      if (ip != ConfigEntityName.DEFAULT) {
+      if (ip != ZooKeeperInternals.DEFAULT_STRING) {
         try {
           Some(InetAddress.getByName(ip))
         } catch {
@@ -246,7 +246,7 @@ class BrokerConfigHandler(private val brokerConfig: KafkaConfig,
       else
         DefaultReplicationThrottledRate
     }
-    if (brokerId == ConfigEntityName.DEFAULT)
+    if (brokerId == ZooKeeperInternals.DEFAULT_STRING)
       brokerConfig.dynamicConfig.updateDefaultConfig(properties)
     else if (brokerConfig.brokerId == brokerId.trim.toInt) {
       brokerConfig.dynamicConfig.updateBrokerConfig(brokerConfig.brokerId, properties)
diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
index 74880971bc382..fefa385558206 100755
--- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
@@ -37,7 +37,7 @@ import org.apache.kafka.common.security.authenticator.LoginManager
 import org.apache.kafka.common.utils.{ConfigUtils, Utils}
 import org.apache.kafka.security.PasswordEncoder
 import org.apache.kafka.server.ProcessRole
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType, ServerTopicConfigSynonyms}
+import org.apache.kafka.server.config.{ConfigType, ServerTopicConfigSynonyms, ZooKeeperInternals}
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.metrics.ClientMetricsReceiverPlugin
 import org.apache.kafka.server.telemetry.ClientTelemetry
@@ -233,7 +233,7 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging
 
     zkClientOpt.foreach { zkClient =>
       val adminZkClient = new AdminZkClient(zkClient)
-      updateDefaultConfig(adminZkClient.fetchEntityConfig(ConfigType.BROKER, ConfigEntityName.DEFAULT), false)
+      updateDefaultConfig(adminZkClient.fetchEntityConfig(ConfigType.BROKER, ZooKeeperInternals.DEFAULT_STRING), false)
       val props = adminZkClient.fetchEntityConfig(ConfigType.BROKER, kafkaConfig.brokerId.toString)
       val brokerConfig = maybeReEncodePasswords(props, adminZkClient)
       updateBrokerConfig(kafkaConfig.brokerId, brokerConfig)
diff --git a/core/src/main/scala/kafka/server/DynamicConfig.scala b/core/src/main/scala/kafka/server/DynamicConfig.scala
index 6ce5863cd5638..4fbbacd8f3451 100644
--- a/core/src/main/scala/kafka/server/DynamicConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicConfig.scala
@@ -23,7 +23,7 @@ import org.apache.kafka.common.config.ConfigDef
 import org.apache.kafka.common.config.ConfigDef.Importance._
 import org.apache.kafka.common.config.ConfigDef.Range._
 import org.apache.kafka.common.config.ConfigDef.Type._
-import org.apache.kafka.server.config.{ConfigEntityName, ReplicationQuotaManagerConfig}
+import org.apache.kafka.server.config.{ReplicationQuotaManagerConfig, ZooKeeperInternals}
 import org.apache.kafka.storage.internals.log.LogConfig
 
 import java.util
@@ -102,7 +102,7 @@ object DynamicConfig {
     def validate(props: Properties) = DynamicConfig.validate(ipConfigs, props, customPropsAllowed = false)
 
     def isValidIpEntity(ip: String): Boolean = {
-      if (ip != ConfigEntityName.DEFAULT) {
+      if (ip != ZooKeeperInternals.DEFAULT_STRING) {
         try {
           InetAddress.getByName(ip)
         } catch {
diff --git a/core/src/main/scala/kafka/server/ZkAdminManager.scala b/core/src/main/scala/kafka/server/ZkAdminManager.scala
index dce32ae3fcc13..6acb4c518ceef 100644
--- a/core/src/main/scala/kafka/server/ZkAdminManager.scala
+++ b/core/src/main/scala/kafka/server/ZkAdminManager.scala
@@ -48,7 +48,7 @@ import org.apache.kafka.common.requests.{AlterConfigsRequest, ApiError}
 import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter}
 import org.apache.kafka.common.utils.Sanitizer
 import org.apache.kafka.server.common.AdminOperationException
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.kafka.storage.internals.log.LogConfig
 
 import scala.collection.{Map, mutable, _}
@@ -518,7 +518,7 @@ class ZkAdminManager(val config: KafkaConfig,
             val perBrokerConfig = brokerId.nonEmpty
 
             val persistentProps = if (perBrokerConfig) adminZkClient.fetchEntityConfig(ConfigType.BROKER, brokerId.get.toString)
-            else adminZkClient.fetchEntityConfig(ConfigType.BROKER, ConfigEntityName.DEFAULT)
+            else adminZkClient.fetchEntityConfig(ConfigType.BROKER, ZooKeeperInternals.DEFAULT_STRING)
 
             val configProps = this.config.dynamicConfig.fromPersistentProps(persistentProps, perBrokerConfig)
             prepareIncrementalConfigs(alterConfigOps, configProps, KafkaConfig.configKeys)
@@ -559,13 +559,13 @@ class ZkAdminManager(val config: KafkaConfig,
 
   private def sanitizeEntityName(entityName: String): String =
     Option(entityName) match {
-      case None => ConfigEntityName.DEFAULT
+      case None => ZooKeeperInternals.DEFAULT_STRING
       case Some(name) => Sanitizer.sanitize(name)
     }
 
   private def desanitizeEntityName(sanitizedEntityName: String): String =
     sanitizedEntityName match {
-      case ConfigEntityName.DEFAULT => null
+      case ZooKeeperInternals.DEFAULT_STRING => null
       case name => Sanitizer.desanitize(name)
     }
 
diff --git a/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala
index c269069d5a4eb..919d203a4d92b 100644
--- a/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala
+++ b/core/src/main/scala/kafka/server/metadata/ClientQuotaMetadataManager.scala
@@ -24,10 +24,10 @@ import org.apache.kafka.common.config.internals.QuotaConfigs
 import org.apache.kafka.common.metrics.Quota
 import org.apache.kafka.common.quota.ClientQuotaEntity
 import org.apache.kafka.common.utils.Sanitizer
-import java.net.{InetAddress, UnknownHostException}
 
+import java.net.{InetAddress, UnknownHostException}
 import org.apache.kafka.image.{ClientQuotaDelta, ClientQuotasDelta}
-import org.apache.kafka.server.config.ConfigEntityName
+import org.apache.kafka.server.config.ZooKeeperInternals
 
 import scala.compat.java8.OptionConverters._
 
@@ -147,13 +147,13 @@ class ClientQuotaMetadataManager(private[metadata] val quotaManagers: QuotaManag
     // Convert entity into Options with sanitized values for QuotaManagers
     val (sanitizedUser, sanitizedClientId) = quotaEntity match {
       case UserEntity(user) => (Some(Sanitizer.sanitize(user)), None)
-      case DefaultUserEntity => (Some(ConfigEntityName.DEFAULT), None)
+      case DefaultUserEntity => (Some(ZooKeeperInternals.DEFAULT_STRING), None)
       case ClientIdEntity(clientId) => (None, Some(Sanitizer.sanitize(clientId)))
-      case DefaultClientIdEntity => (None, Some(ConfigEntityName.DEFAULT))
+      case DefaultClientIdEntity => (None, Some(ZooKeeperInternals.DEFAULT_STRING))
       case ExplicitUserExplicitClientIdEntity(user, clientId) => (Some(Sanitizer.sanitize(user)), Some(Sanitizer.sanitize(clientId)))
-      case ExplicitUserDefaultClientIdEntity(user) => (Some(Sanitizer.sanitize(user)), Some(ConfigEntityName.DEFAULT))
-      case DefaultUserExplicitClientIdEntity(clientId) => (Some(ConfigEntityName.DEFAULT), Some(Sanitizer.sanitize(clientId)))
-      case DefaultUserDefaultClientIdEntity => (Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT))
+      case ExplicitUserDefaultClientIdEntity(user) => (Some(Sanitizer.sanitize(user)), Some(ZooKeeperInternals.DEFAULT_STRING))
+      case DefaultUserExplicitClientIdEntity(clientId) => (Some(ZooKeeperInternals.DEFAULT_STRING), Some(Sanitizer.sanitize(clientId)))
+      case DefaultUserDefaultClientIdEntity => (Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING))
       case IpEntity(_) | DefaultIpEntity => throw new IllegalStateException("Should not see IP quota entities here")
     }
 
diff --git a/core/src/main/scala/kafka/server/metadata/DynamicConfigPublisher.scala b/core/src/main/scala/kafka/server/metadata/DynamicConfigPublisher.scala
index 498c26b32ba95..dfe3716bf1e8d 100644
--- a/core/src/main/scala/kafka/server/metadata/DynamicConfigPublisher.scala
+++ b/core/src/main/scala/kafka/server/metadata/DynamicConfigPublisher.scala
@@ -24,7 +24,7 @@ import kafka.utils.Logging
 import org.apache.kafka.common.config.ConfigResource.Type.{BROKER, CLIENT_METRICS, TOPIC}
 import org.apache.kafka.image.loader.LoaderManifest
 import org.apache.kafka.image.{MetadataDelta, MetadataImage}
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.kafka.server.fault.FaultHandler
 
 
@@ -78,7 +78,7 @@ class DynamicConfigPublisher(
                     // These are stored in KRaft with an empty name field.
                     info("Updating cluster configuration : " +
                       toLoggableProps(resource, props).mkString(","))
-                    nodeConfigHandler.processConfigChanges(ConfigEntityName.DEFAULT, props)
+                    nodeConfigHandler.processConfigChanges(ZooKeeperInternals.DEFAULT_STRING, props)
                   } catch {
                     case t: Throwable => faultHandler.handleFault("Error updating " +
                       s"cluster with new configuration: ${toLoggableProps(resource, props).mkString(",")} " +
diff --git a/core/src/main/scala/kafka/server/metadata/ZkConfigRepository.scala b/core/src/main/scala/kafka/server/metadata/ZkConfigRepository.scala
index f4b98a0835731..58eadcbf21b44 100644
--- a/core/src/main/scala/kafka/server/metadata/ZkConfigRepository.scala
+++ b/core/src/main/scala/kafka/server/metadata/ZkConfigRepository.scala
@@ -22,7 +22,7 @@ import kafka.zk.{AdminZkClient, KafkaZkClient}
 import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.common.config.ConfigResource.Type
 import org.apache.kafka.common.errors.InvalidRequestException
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 
 
 object ZkConfigRepository {
@@ -41,7 +41,7 @@ class ZkConfigRepository(adminZkClient: AdminZkClient) extends ConfigRepository
     // ZK stores cluster configs under "<default>".
     val effectiveName = if (configResource.`type`.equals(Type.BROKER) &&
         configResource.name.isEmpty) {
-      ConfigEntityName.DEFAULT
+      ZooKeeperInternals.DEFAULT_STRING
     } else {
       configResource.name
     }
diff --git a/core/src/main/scala/kafka/zk/AdminZkClient.scala b/core/src/main/scala/kafka/zk/AdminZkClient.scala
index 72925a036b10b..efecfe854bbf2 100644
--- a/core/src/main/scala/kafka/zk/AdminZkClient.scala
+++ b/core/src/main/scala/kafka/zk/AdminZkClient.scala
@@ -28,7 +28,7 @@ import org.apache.kafka.common.{TopicPartition, Uuid}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.server.common.AdminOperationException
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.kafka.storage.internals.log.LogConfig
 import org.apache.zookeeper.KeeperException.NodeExistsException
 
@@ -345,7 +345,7 @@ class AdminZkClient(zkClient: KafkaZkClient,
    */
   def parseBroker(broker: String): Option[Int] = {
     broker match {
-      case ConfigEntityName.DEFAULT => None
+      case ZooKeeperInternals.DEFAULT_STRING => None
       case _ =>
         try Some(broker.toInt)
         catch {
@@ -440,7 +440,7 @@ class AdminZkClient(zkClient: KafkaZkClient,
    *
    */
   def changeUserOrUserClientIdConfig(sanitizedEntityName: String, configs: Properties, isUserClientId: Boolean = false): Unit = {
-    if (sanitizedEntityName == ConfigEntityName.DEFAULT || sanitizedEntityName.contains("/clients"))
+    if (sanitizedEntityName == ZooKeeperInternals.DEFAULT_STRING || sanitizedEntityName.contains("/clients"))
       DynamicConfig.Client.validate(configs)
     else
       DynamicConfig.User.validate(configs)
@@ -520,7 +520,7 @@ class AdminZkClient(zkClient: KafkaZkClient,
     */
   def changeBrokerConfig(broker: Option[Int], configs: Properties): Unit = {
     validateBrokerConfig(configs)
-    changeEntityConfig(ConfigType.BROKER, broker.map(_.toString).getOrElse(ConfigEntityName.DEFAULT), configs)
+    changeEntityConfig(ConfigType.BROKER, broker.map(_.toString).getOrElse(ZooKeeperInternals.DEFAULT_STRING), configs)
   }
 
   /**
diff --git a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
index e69719e7501a3..ee960bd35f8bf 100644
--- a/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/ZkMigrationClient.scala
@@ -27,7 +27,6 @@ import org.apache.kafka.common.errors.ControllerMovedException
 import org.apache.kafka.common.metadata._
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.scram.ScramCredential
-import org.apache.kafka.common.utils.Sanitizer
 import org.apache.kafka.common.{TopicIdPartition, Uuid}
 import org.apache.kafka.metadata.DelegationTokenData
 import org.apache.kafka.metadata.PartitionRegistration
@@ -227,9 +226,6 @@ class ZkMigrationClient(
         entityDataList: util.List[ClientQuotaRecord.EntityData],
         quotas: util.Map[String, lang.Double]
       ): Unit = {
-        entityDataList.forEach(entityData => {
-          entityData.setEntityName(Sanitizer.desanitize(entityData.entityName()))
-        })
         val batch = new util.ArrayList[ApiMessageAndVersion]()
         quotas.forEach((key, value) => {
           batch.add(new ApiMessageAndVersion(new ClientQuotaRecord()
diff --git a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
index dffdc138aa138..fda6e93c2492d 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
@@ -21,6 +21,7 @@ import kafka.server.{DynamicBrokerConfig, DynamicConfig, ZkAdminManager}
 import kafka.utils.Logging
 import kafka.zk.ZkMigrationClient.{logAndRethrow, wrapZkException}
 import kafka.zk._
+import kafka.zk.migration.ZkConfigMigrationClient.getSanitizedClientQuotaZNodeName
 import kafka.zookeeper.{CreateRequest, DeleteRequest, SetDataRequest}
 import org.apache.kafka.clients.admin.ScramMechanism
 import org.apache.kafka.common.config.types.Password
@@ -29,14 +30,14 @@ import org.apache.kafka.common.errors.InvalidRequestException
 import org.apache.kafka.common.metadata.ClientQuotaRecord.EntityData
 import org.apache.kafka.common.quota.ClientQuotaEntity
 import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils
+import org.apache.kafka.common.utils.Sanitizer
 import org.apache.kafka.metadata.migration.ConfigMigrationClient.ClientQuotaVisitor
 import org.apache.kafka.metadata.migration.{ConfigMigrationClient, MigrationClientException, ZkMigrationLeadershipState}
 import org.apache.kafka.security.PasswordEncoder
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.zookeeper.KeeperException.Code
 import org.apache.zookeeper.{CreateMode, KeeperException}
 
-
 import java.{lang, util}
 import java.util.Properties
 import java.util.function.{BiConsumer, Consumer}
@@ -50,44 +51,54 @@ class ZkConfigMigrationClient(
 
   val adminZkClient = new AdminZkClient(zkClient)
 
-
   /**
-   * In ZK, we use the special string "&lt;default&gt;" to represent the default entity.
-   * In KRaft, we use an empty string. This method builds an EntityData that converts the special ZK string
-   * to the special KRaft string.
+   * In ZK, we use the special string "&lt;default&gt;" to represent the default config entity.
+   * In KRaft, we use an empty string. This method converts the between the two conventions.
    */
-  private def fromZkEntityName(entityName: String): String = {
-    if (entityName.equals(ConfigEntityName.DEFAULT)) {
+  private def fromZkConfigfEntityName(entityName: String): String = {
+    if (entityName.equals(ZooKeeperInternals.DEFAULT_STRING)) {
       ""
     } else {
       entityName
     }
   }
 
-  private def toZkEntityName(entityName: String): String = {
+  private def toZkConfigEntityName(entityName: String): String = {
     if (entityName.isEmpty) {
-      ConfigEntityName.DEFAULT
+      ZooKeeperInternals.DEFAULT_STRING
     } else {
       entityName
     }
   }
 
-  private def buildEntityData(entityType: String, entityName: String): EntityData = {
-    new EntityData().setEntityType(entityType).setEntityName(fromZkEntityName(entityName))
+  private def buildClientQuotaEntityData(
+    entityType: String,
+    znodeName: String
+  ): EntityData = {
+    val result = new EntityData().setEntityType(entityType)
+    if (znodeName.equals(ZooKeeperInternals.DEFAULT_STRING)) {
+      // Default __client quota__ entity names are null. This is different than default __configs__,
+      // which have their names set to the empty string instead.
+      result.setEntityName(null)
+    } else {
+      // ZNode names are sanitized before being stored in ZooKeeper.
+      // For example, @ is turned into %40. Undo the sanitization here.
+      result.setEntityName(Sanitizer.desanitize(znodeName))
+    }
+    result
   }
 
-
   override def iterateClientQuotas(visitor: ClientQuotaVisitor): Unit = {
     def migrateEntityType(zkEntityType: String, entityType: String): Unit = {
-      adminZkClient.fetchAllEntityConfigs(zkEntityType).foreach { case (name, props) =>
-        val entity = List(buildEntityData(entityType, name)).asJava
+      adminZkClient.fetchAllEntityConfigs(zkEntityType).foreach { case (znodeName, props) =>
+        val entity = List(buildClientQuotaEntityData(entityType, znodeName)).asJava
 
         ScramMechanism.values().filter(_ != ScramMechanism.UNKNOWN).foreach { mechanism =>
           val propertyValue = props.getProperty(mechanism.mechanismName)
           if (propertyValue != null) {
             val scramCredentials = ScramCredentialUtils.credentialFromString(propertyValue)
             logAndRethrow(this, s"Error in client quota visitor for SCRAM credential. User was $entity.") {
-              visitor.visitScramCredential(name, mechanism, scramCredentials)
+              visitor.visitScramCredential(Sanitizer.desanitize(znodeName), mechanism, scramCredentials)
             }
             props.remove(mechanism.mechanismName)
           }
@@ -108,14 +119,14 @@ class ZkConfigMigrationClient(
     migrateEntityType(ConfigType.USER, ClientQuotaEntity.USER)
     migrateEntityType(ConfigType.CLIENT, ClientQuotaEntity.CLIENT_ID)
 
-    adminZkClient.fetchAllChildEntityConfigs(ConfigType.USER, ConfigType.CLIENT).foreach { case (name, props) =>
+    adminZkClient.fetchAllChildEntityConfigs(ConfigType.USER, ConfigType.CLIENT).foreach { case (znodePath, props) =>
       // Taken from ZkAdminManager
-      val components = name.split("/")
+      val components = znodePath.split("/")
       if (components.size != 3 || components(1) != "clients")
-        throw new IllegalArgumentException(s"Unexpected config path: $name")
+        throw new IllegalArgumentException(s"Unexpected config path: $znodePath")
       val entity = List(
-        buildEntityData(ClientQuotaEntity.USER, components(0)),
-        buildEntityData(ClientQuotaEntity.CLIENT_ID, components(2))
+        buildClientQuotaEntityData(ClientQuotaEntity.USER, components(0)),
+        buildClientQuotaEntityData(ClientQuotaEntity.CLIENT_ID, components(2))
       )
       val quotaMap = props.asScala.map { case (key, value) =>
         val doubleValue = try lang.Double.valueOf(value) catch {
@@ -135,7 +146,7 @@ class ZkConfigMigrationClient(
   override def iterateBrokerConfigs(configConsumer: BiConsumer[String, util.Map[String, String]]): Unit = {
     val brokerEntities = zkClient.getAllEntitiesWithConfig(ConfigType.BROKER)
     zkClient.getEntitiesConfigs(ConfigType.BROKER, brokerEntities.toSet).foreach { case (broker, props) =>
-      val brokerResource = fromZkEntityName(broker)
+      val brokerResource = fromZkConfigfEntityName(broker)
       val decodedProps = props.asScala.map { case (key, value) =>
         if (DynamicBrokerConfig.isPasswordConfig(key))
           key -> passwordEncoder.decode(value).value
@@ -157,7 +168,7 @@ class ZkConfigMigrationClient(
   }
 
   override def readTopicConfigs(topicName: String, configConsumer: Consumer[util.Map[String, String]]): Unit = {
-    val topicResource = fromZkEntityName(topicName)
+    val topicResource = fromZkConfigfEntityName(topicName)
     val props = zkClient.getEntityConfigs(ConfigType.TOPIC, topicResource)
     val decodedProps = props.asScala.map { case (key, value) =>
       if (DynamicBrokerConfig.isPasswordConfig(key))
@@ -182,7 +193,7 @@ class ZkConfigMigrationClient(
       case _ => None
     }
 
-    val configName = toZkEntityName(configResource.name())
+    val configName = toZkConfigEntityName(configResource.name())
     if (configType.isDefined) {
       val props = new Properties()
       configMap.forEach { case (key, value) =>
@@ -221,7 +232,7 @@ class ZkConfigMigrationClient(
       case _ => None
     }
 
-    val configName = toZkEntityName(configResource.name())
+    val configName = toZkConfigEntityName(configResource.name())
     if (configType.isDefined) {
       val path = ConfigEntityZNode.path(configType.get, configName)
       val requests = Seq(DeleteRequest(path, ZkVersion.MatchAnyVersion))
@@ -250,10 +261,9 @@ class ZkConfigMigrationClient(
     scram: util.Map[String, String],
     state: ZkMigrationLeadershipState
   ): ZkMigrationLeadershipState = wrapZkException {
-    val entityMap = entity.asScala
-    val user = entityMap.get(ClientQuotaEntity.USER).map(toZkEntityName)
-    val client = entityMap.get(ClientQuotaEntity.CLIENT_ID).map(toZkEntityName)
-    val ip = entityMap.get(ClientQuotaEntity.IP).map(toZkEntityName)
+    val user: Option[String] = getSanitizedClientQuotaZNodeName(entity, ClientQuotaEntity.USER)
+    val client: Option[String] = getSanitizedClientQuotaZNodeName(entity, ClientQuotaEntity.CLIENT_ID)
+    val ip: Option[String] = getSanitizedClientQuotaZNodeName(entity, ClientQuotaEntity.IP)
     val props = new Properties()
 
     val (configType, path, configKeys) = if (user.isDefined && client.isEmpty) {
@@ -351,3 +361,34 @@ class ZkConfigMigrationClient(
   }
 }
 
+object ZkConfigMigrationClient {
+  /**
+   * Find the znode name to use for a ClientQuotaEntity.
+   *
+   * @param entity      The client quota entity map. See org.apache.kafka.common.ClientQuotaEntity.
+   * @param component   The component that we want a znode name for.
+   * @return            Some(znodeName) if there is a znode path; None otherwise.
+   */
+  def getSanitizedClientQuotaZNodeName(
+    entity: util.Map[String, String],
+    component: String
+  ): Option[String] = {
+    if (!entity.containsKey(component)) {
+      // There is no znode path, because the component wasn't found. For example, if the
+      // entity was (user -> "bob") and our component was "ip", we would return None here.
+      None
+    } else {
+      val rawValue = entity.get(component)
+      if (rawValue == null) {
+        // A raw value of null means this is a default entity. For example, (user -> null) means
+        // the default user. Yes, this means we stored a null value in the map and it did not mean
+        // "not present." This is an unfortunate API that should be revisited at some point.
+        Some(ZooKeeperInternals.DEFAULT_STRING)
+      } else {
+        // We found a non-null value, and now we need to sanitize it. For example, "c@@ldude" will
+        // turn into c%40%40ldude, so that we can use it as a znode name in ZooKeeper.
+        Some(Sanitizer.sanitize(rawValue))
+      }
+    }
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/admin/ConfigCommandIntegrationTest.scala b/core/src/test/scala/integration/kafka/admin/ConfigCommandIntegrationTest.scala
index 6335e978533b7..22e4be47579bc 100644
--- a/core/src/test/scala/integration/kafka/admin/ConfigCommandIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/ConfigCommandIntegrationTest.scala
@@ -25,7 +25,7 @@ import org.apache.kafka.common.config.ConfigException
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.server.common.MetadataVersion
-import org.apache.kafka.server.config.ConfigEntityName
+import org.apache.kafka.server.config.ZooKeeperInternals
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
@@ -93,7 +93,7 @@ class ConfigCommandIntegrationTest extends QuorumTestHarness with Logging {
     }
 
     def verifyConfig(configs: Map[String, String], brokerId: Option[String]): Unit = {
-      val entityConfigs = zkClient.getEntityConfigs("brokers", brokerId.getOrElse(ConfigEntityName.DEFAULT))
+      val entityConfigs = zkClient.getEntityConfigs("brokers", brokerId.getOrElse(ZooKeeperInternals.DEFAULT_STRING))
       assertEquals(configs, entityConfigs.asScala)
     }
 
diff --git a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
index 75f07b8efe853..50be178f76035 100644
--- a/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KRaftClusterTest.scala
@@ -26,12 +26,13 @@ import org.apache.kafka.clients.admin._
 import org.apache.kafka.common.acl.{AclBinding, AclBindingFilter}
 import org.apache.kafka.common.config.{ConfigException, ConfigResource}
 import org.apache.kafka.common.config.ConfigResource.Type
-import org.apache.kafka.common.errors.{InvalidPartitionsException,PolicyViolationException, UnsupportedVersionException}
+import org.apache.kafka.common.errors.{InvalidPartitionsException, PolicyViolationException, UnsupportedVersionException}
 import org.apache.kafka.common.message.DescribeClusterRequestData
 import org.apache.kafka.common.metadata.{ConfigRecord, FeatureLevelRecord}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.Errors._
+import org.apache.kafka.common.quota.ClientQuotaAlteration.Op
 import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
 import org.apache.kafka.common.requests.{ApiError, DescribeClusterRequest, DescribeClusterResponse}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
@@ -324,6 +325,68 @@ class KRaftClusterTest {
     }
   }
 
+  def setConsumerByteRate(
+    admin: Admin,
+    entity: ClientQuotaEntity,
+    value: Long
+  ): Unit = {
+    admin.alterClientQuotas(Collections.singletonList(
+      new ClientQuotaAlteration(entity, Collections.singletonList(
+        new Op("consumer_byte_rate", value.doubleValue()))))).
+        all().get()
+  }
+
+  def getConsumerByteRates(admin: Admin): Map[ClientQuotaEntity, Long] = {
+    val allFilter = ClientQuotaFilter.contains(Collections.emptyList())
+    val results = new java.util.HashMap[ClientQuotaEntity, Long]
+    admin.describeClientQuotas(allFilter).entities().get().forEach {
+      case (entity, entityMap) =>
+        Option(entityMap.get("consumer_byte_rate")).foreach {
+          case value => results.put(entity, value.longValue())
+        }
+    }
+    results.asScala.toMap
+  }
+
+  @Test
+  def testDefaultClientQuotas(): Unit = {
+    val cluster = new KafkaClusterTestKit.Builder(
+      new TestKitNodes.Builder().
+        setNumBrokerNodes(1).
+        setNumControllerNodes(1).build()).build()
+    try {
+      cluster.format()
+      cluster.startup()
+      TestUtils.waitUntilTrue(() => cluster.brokers().get(0).brokerState == BrokerState.RUNNING,
+        "Broker never made it to RUNNING state.")
+      val admin = Admin.create(cluster.clientProperties())
+      try {
+        val defaultUser = new ClientQuotaEntity(Collections.singletonMap[String, String]("user", null))
+        val bobUser = new ClientQuotaEntity(Collections.singletonMap[String, String]("user", "bob"))
+        TestUtils.retry(30000) {
+          assertEquals(Map(), getConsumerByteRates(admin))
+        }
+        setConsumerByteRate(admin, defaultUser, 100L)
+        TestUtils.retry(30000) {
+          assertEquals(Map(
+              defaultUser -> 100L
+            ), getConsumerByteRates(admin))
+        }
+        setConsumerByteRate(admin, bobUser, 1000L)
+        TestUtils.retry(30000) {
+          assertEquals(Map(
+            defaultUser -> 100L,
+            bobUser -> 1000L
+          ), getConsumerByteRates(admin))
+        }
+      } finally {
+        admin.close()
+      }
+    } finally {
+      cluster.close()
+    }
+  }
+
   @Test
   def testCreateClusterWithAdvertisedPortZero(): Unit = {
     val brokerPropertyOverrides: (TestKitNodes, BrokerNode) => Map[String, String] = (nodes, _) => Map(
diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
index d442df33adfdf..a3e5ecc0814d9 100644
--- a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -40,7 +40,7 @@ import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.resource.ResourceType.TOPIC
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils
-import org.apache.kafka.common.utils.SecurityUtils
+import org.apache.kafka.common.utils.{Sanitizer, SecurityUtils}
 import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
 import org.apache.kafka.metadata.authorizer.StandardAcl
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
@@ -48,7 +48,7 @@ import org.apache.kafka.raft.RaftConfig
 import org.apache.kafka.security.PasswordEncoder
 import org.apache.kafka.server.ControllerRequestCompletionHandler
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.ConfigType
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotEquals, assertNotNull, assertTrue, fail}
 import org.junit.jupiter.api.{Assumptions, Timeout}
 import org.junit.jupiter.api.extension.ExtendWith
@@ -227,11 +227,11 @@ class ZkMigrationIntegrationTest {
     createTopicResult.all().get(60, TimeUnit.SECONDS)
 
     val quotas = new util.ArrayList[ClientQuotaAlteration]()
-    val defaultUserEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> ConfigEntityName.DEFAULT).asJava)
+    val defaultUserEntity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.USER, null))
     quotas.add(new ClientQuotaAlteration(defaultUserEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 900.0)).asJava))
-    val defaultClientIdEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.CLIENT_ID -> ConfigEntityName.DEFAULT).asJava)
+    val defaultClientIdEntity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.CLIENT_ID, null))
     quotas.add(new ClientQuotaAlteration(defaultClientIdEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 900.0)).asJava))
-    val defaultIpEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> null.asInstanceOf[String]).asJava)
+    val defaultIpEntity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.IP, null))
     quotas.add(new ClientQuotaAlteration(defaultIpEntity, List(new ClientQuotaAlteration.Op("connection_creation_rate", 9.0)).asJava))
     val userEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> "user/1@prod").asJava)
     quotas.add(new ClientQuotaAlteration(userEntity, List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
@@ -275,15 +275,15 @@ class ZkMigrationIntegrationTest {
       assertEquals(10, image.topics().getTopic("test-topic-3").partitions().size())
 
       val clientQuotas = image.clientQuotas().entities()
-      assertEquals(6, clientQuotas.size())
-      assertEquals(true, clientQuotas.containsKey(defaultUserEntity))
-      assertEquals(true, clientQuotas.containsKey(defaultClientIdEntity))
-      assertEquals(true, clientQuotas.containsKey(new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> "").asJava))) // default ip
-      assertEquals(true, clientQuotas.containsKey(userEntity))
-      assertEquals(true, clientQuotas.containsKey(userClientEntity))
-      assertEquals(true, clientQuotas.containsKey(ipEntity))
+      assertEquals(new java.util.HashSet[ClientQuotaEntity](java.util.Arrays.asList(
+        defaultUserEntity,
+        defaultClientIdEntity,
+        defaultIpEntity,
+        userEntity,
+        userClientEntity,
+        ipEntity
+      )), clientQuotas.keySet())
     }
-
     migrationState = migrationClient.releaseControllerLeadership(migrationState)
   }
 
@@ -881,11 +881,14 @@ class ZkMigrationIntegrationTest {
   def alterClientQuotas(admin: Admin): AlterClientQuotasResult = {
     val quotas = new util.ArrayList[ClientQuotaAlteration]()
     quotas.add(new ClientQuotaAlteration(
-      new ClientQuotaEntity(Map("user" -> "user1").asJava),
+      new ClientQuotaEntity(Map("user" -> "user@1").asJava),
       List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
     quotas.add(new ClientQuotaAlteration(
-      new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava),
+      new ClientQuotaEntity(Map("user" -> "user@1", "client-id" -> "clientA").asJava),
       List(new ClientQuotaAlteration.Op("consumer_byte_rate", 800.0), new ClientQuotaAlteration.Op("producer_byte_rate", 100.0)).asJava))
+    quotas.add(new ClientQuotaAlteration(
+      new ClientQuotaEntity(Collections.singletonMap("user", null)),
+      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 900.0), new ClientQuotaAlteration.Op("producer_byte_rate", 100.0)).asJava))
     quotas.add(new ClientQuotaAlteration(
       new ClientQuotaEntity(Map("ip" -> "8.8.8.8").asJava),
       List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
@@ -903,7 +906,7 @@ class ZkMigrationIntegrationTest {
     val alterations = new util.ArrayList[UserScramCredentialAlteration]()
     alterations.add(new UserScramCredentialUpsertion("user1",
         new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 8191), "password1"))
-    alterations.add(new UserScramCredentialUpsertion("user2",
+    alterations.add(new UserScramCredentialUpsertion("user@2",
         new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 8192), "password2"))
     admin.alterUserScramCredentials(alterations)
   }
@@ -918,20 +921,21 @@ class ZkMigrationIntegrationTest {
 
   def verifyClientQuotas(zkClient: KafkaZkClient): Unit = {
     TestUtils.retry(10000) {
-      assertEquals("1000", zkClient.getEntityConfigs(ConfigType.USER, "user1").getProperty("consumer_byte_rate"))
-      assertEquals("800", zkClient.getEntityConfigs("users/user1/clients", "clientA").getProperty("consumer_byte_rate"))
-      assertEquals("100", zkClient.getEntityConfigs("users/user1/clients", "clientA").getProperty("producer_byte_rate"))
+      assertEquals("1000", zkClient.getEntityConfigs(ConfigType.USER, Sanitizer.sanitize("user@1")).getProperty("consumer_byte_rate"))
+      assertEquals("900", zkClient.getEntityConfigs(ConfigType.USER, "<default>").getProperty("consumer_byte_rate"))
+      assertEquals("800", zkClient.getEntityConfigs("users/" + Sanitizer.sanitize("user@1") + "/clients", "clientA").getProperty("consumer_byte_rate"))
+      assertEquals("100", zkClient.getEntityConfigs("users/" + Sanitizer.sanitize("user@1") + "/clients", "clientA").getProperty("producer_byte_rate"))
       assertEquals("10", zkClient.getEntityConfigs(ConfigType.IP, "8.8.8.8").getProperty("connection_creation_rate"))
     }
   }
 
   def verifyUserScramCredentials(zkClient: KafkaZkClient): Unit = {
     TestUtils.retry(10000) {
-      val propertyValue1 = zkClient.getEntityConfigs(ConfigType.USER, "user1").getProperty("SCRAM-SHA-256")
+      val propertyValue1 = zkClient.getEntityConfigs(ConfigType.USER, Sanitizer.sanitize("user1")).getProperty("SCRAM-SHA-256")
       val scramCredentials1 = ScramCredentialUtils.credentialFromString(propertyValue1)
       assertEquals(8191, scramCredentials1.iterations)
 
-      val propertyValue2 = zkClient.getEntityConfigs(ConfigType.USER, "user2").getProperty("SCRAM-SHA-256")
+      val propertyValue2 = zkClient.getEntityConfigs(ConfigType.USER, Sanitizer.sanitize("user@2")).getProperty("SCRAM-SHA-256")
       assertNotNull(propertyValue2)
       val scramCredentials2 = ScramCredentialUtils.credentialFromString(propertyValue2)
       assertEquals(8192, scramCredentials2.iterations)
diff --git a/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
index 085a2fd649a21..3d2514428b248 100644
--- a/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
@@ -30,7 +30,7 @@ import org.apache.kafka.common.internals.KafkaFutureImpl
 import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
 import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils
 import org.apache.kafka.common.utils.Sanitizer
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
 import org.mockito.ArgumentMatchers.anyString
@@ -1475,7 +1475,7 @@ class ConfigCommandTest extends Logging {
     val clientId = "client-1"
     for (opts <- Seq(describeOpts, alterOpts)) {
       checkEntity("clients", Some(clientId), clientId, opts)
-      checkEntity("clients", Some(""), ConfigEntityName.DEFAULT, opts)
+      checkEntity("clients", Some(""), ZooKeeperInternals.DEFAULT_STRING, opts)
     }
     checkEntity("clients", None, "", describeOpts)
     checkInvalidArgs("clients", None, alterOpts)
@@ -1487,7 +1487,7 @@ class ConfigCommandTest extends Logging {
     assertEquals(principal, Sanitizer.desanitize(sanitizedPrincipal))
     for (opts <- Seq(describeOpts, alterOpts)) {
       checkEntity("users", Some(principal), sanitizedPrincipal, opts)
-      checkEntity("users", Some(""), ConfigEntityName.DEFAULT, opts)
+      checkEntity("users", Some(""), ZooKeeperInternals.DEFAULT_STRING, opts)
     }
     checkEntity("users", None, "", describeOpts)
     checkInvalidArgs("users", None, alterOpts)
@@ -1497,9 +1497,9 @@ class ConfigCommandTest extends Logging {
     def clientIdOpts(name: String) = Array("--entity-type", "clients", "--entity-name", name)
     for (opts <- Seq(describeOpts, alterOpts)) {
       checkEntity("users", Some(principal), userClient, opts ++ clientIdOpts(clientId))
-      checkEntity("users", Some(principal), sanitizedPrincipal + "/clients/" + ConfigEntityName.DEFAULT, opts ++ clientIdOpts(""))
-      checkEntity("users", Some(""), ConfigEntityName.DEFAULT + "/clients/" + clientId, describeOpts ++ clientIdOpts(clientId))
-      checkEntity("users", Some(""), ConfigEntityName.DEFAULT + "/clients/" + ConfigEntityName.DEFAULT, opts ++ clientIdOpts(""))
+      checkEntity("users", Some(principal), sanitizedPrincipal + "/clients/" + ZooKeeperInternals.DEFAULT_STRING, opts ++ clientIdOpts(""))
+      checkEntity("users", Some(""), ZooKeeperInternals.DEFAULT_STRING + "/clients/" + clientId, describeOpts ++ clientIdOpts(clientId))
+      checkEntity("users", Some(""), ZooKeeperInternals.DEFAULT_STRING + "/clients/" + ZooKeeperInternals.DEFAULT_STRING, opts ++ clientIdOpts(""))
     }
     checkEntity("users", Some(principal), sanitizedPrincipal + "/clients", describeOpts ++ Array("--entity-type", "clients"))
     // Both user and client-id must be provided for alter
diff --git a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
index e67203c565f38..dd7eac3b257f2 100644
--- a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
@@ -21,7 +21,7 @@ import kafka.server.QuotaType._
 import org.apache.kafka.common.metrics.Quota
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.Sanitizer
-import org.apache.kafka.server.config.{ClientQuotaManagerConfig, ConfigEntityName}
+import org.apache.kafka.server.config.{ClientQuotaManagerConfig, ZooKeeperInternals}
 import org.apache.kafka.network.Session
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
@@ -85,7 +85,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val client1 = UserClient("ANONYMOUS", "p1", None, Some("p1"))
     val client2 = UserClient("ANONYMOUS", "p2", None, Some("p2"))
     val randomClient = UserClient("ANONYMOUS", "random-client-id", None, None)
-    val defaultConfigClient = UserClient("", "", None, Some(ConfigEntityName.DEFAULT))
+    val defaultConfigClient = UserClient("", "", None, Some(ZooKeeperInternals.DEFAULT_STRING))
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
@@ -98,7 +98,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val client1 = UserClient("User1", "p1", Some("User1"), None)
     val client2 = UserClient("User2", "p2", Some("User2"), None)
     val randomClient = UserClient("RandomUser", "random-client-id", None, None)
-    val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.DEFAULT), None)
+    val defaultConfigClient = UserClient("", "", Some(ZooKeeperInternals.DEFAULT_STRING), None)
     val config = new ClientQuotaManagerConfig()
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
@@ -112,7 +112,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val client1 = UserClient("User1", "p1", Some("User1"), Some("p1"))
     val client2 = UserClient("User2", "p2", Some("User2"), Some("p2"))
     val randomClient = UserClient("RandomUser", "random-client-id", None, None)
-    val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT))
+    val defaultConfigClient = UserClient("", "", Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING))
     val config = new ClientQuotaManagerConfig()
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
@@ -125,7 +125,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val client1 = UserClient("User1", "p1", Some("User1"), None)
     val client2 = UserClient("User2", "p2", Some("User2"), None)
     val randomClient = UserClient("RandomUser", "random-client-id", None, None)
-    val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.DEFAULT), None)
+    val defaultConfigClient = UserClient("", "", Some(ZooKeeperInternals.DEFAULT_STRING), None)
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
@@ -137,7 +137,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val client1 = UserClient("User1", "p1", Some("User1"), Some("p1"))
     val client2 = UserClient("User2", "p2", Some("User2"), Some("p2"))
     val randomClient = UserClient("RandomUser", "random-client-id", None, None)
-    val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT))
+    val defaultConfigClient = UserClient("", "", Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING))
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
@@ -168,7 +168,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
       assertEquals(Double.MaxValue, clientQuotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
 
       // Set default <user> quota config
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), None, None, Some(new Quota(10, true)))
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), None, None, Some(new Quota(10, true)))
       assertEquals(10 * numFullQuotaWindows, clientQuotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
     } finally {
       clientQuotaManager.shutdown()
@@ -186,11 +186,11 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
       checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, false)
 
       // Set default <user> quota config
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), None, None, Some(new Quota(10, true)))
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), None, None, Some(new Quota(10, true)))
       checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true)
 
       // Remove default <user> quota config, back to no quotas
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), None, None, None)
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), None, None, None)
       checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, false)
     } finally {
       clientQuotaManager.shutdown()
@@ -241,14 +241,14 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
       metrics, Produce, time, "")
 
     try {
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), None, None, Some(new Quota(1000, true)))
-      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(new Quota(2000, true)))
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(new Quota(3000, true)))
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), None, None, Some(new Quota(1000, true)))
+      clientQuotaManager.updateQuota(None, Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(new Quota(2000, true)))
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(new Quota(3000, true)))
       clientQuotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(4000, true)))
       clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(5000, true)))
       clientQuotaManager.updateQuota(Some("userB"), None, None, Some(new Quota(6000, true)))
       clientQuotaManager.updateQuota(Some("userB"), Some("client1"), Some("client1"), Some(new Quota(7000, true)))
-      clientQuotaManager.updateQuota(Some("userB"), Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(new Quota(8000, true)))
+      clientQuotaManager.updateQuota(Some("userB"), Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(new Quota(8000, true)))
       clientQuotaManager.updateQuota(Some("userC"), None, None, Some(new Quota(10000, true)))
       clientQuotaManager.updateQuota(None, Some("client1"), Some("client1"), Some(new Quota(9000, true)))
 
@@ -266,14 +266,14 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
       checkQuota(clientQuotaManager, "userE", "client1", 3000, 2500, false)
 
       // Remove default <user, client> quota config, revert to <user> default
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), None)
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), None)
       checkQuota(clientQuotaManager, "userD", "client1", 1000, 0, false)    // Metrics tags changed, restart counter
       checkQuota(clientQuotaManager, "userE", "client4", 1000, 1500, true)
       checkQuota(clientQuotaManager, "userF", "client4", 1000, 800, false)  // Default <user> quota shared across clients of user
       checkQuota(clientQuotaManager, "userF", "client5", 1000, 800, true)
 
       // Remove default <user> quota config, revert to <client-id> default
-      clientQuotaManager.updateQuota(Some(ConfigEntityName.DEFAULT), None, None, None)
+      clientQuotaManager.updateQuota(Some(ZooKeeperInternals.DEFAULT_STRING), None, None, None)
       checkQuota(clientQuotaManager, "userF", "client4", 2000, 0, false)  // Default <client-id> quota shared across client-id of all users
       checkQuota(clientQuotaManager, "userF", "client5", 2000, 0, false)
       checkQuota(clientQuotaManager, "userF", "client5", 2000, 2500, true)
@@ -290,7 +290,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
       checkQuota(clientQuotaManager, "userA", "client6", 8000, 0, true)    // Throttled due to shared user quota
       clientQuotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), Some(new Quota(11000, true)))
       checkQuota(clientQuotaManager, "userA", "client6", 11000, 8500, false)
-      clientQuotaManager.updateQuota(Some("userA"), Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT), Some(new Quota(12000, true)))
+      clientQuotaManager.updateQuota(Some("userA"), Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING), Some(new Quota(12000, true)))
       clientQuotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), None)
       checkQuota(clientQuotaManager, "userA", "client6", 12000, 4000, true) // Throttled due to sum of new and earlier values
 
@@ -304,7 +304,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""))
     try {
-      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT),
+      clientQuotaManager.updateQuota(None, Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING),
         Some(new Quota(500, true)))
 
       // We have 10 second windows. Make sure that there is no quota violation
@@ -352,7 +352,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
   def testExpireThrottleTimeSensor(): Unit = {
     val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     try {
-      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT),
+      clientQuotaManager.updateQuota(None, Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING),
         Some(new Quota(500, true)))
 
       maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
@@ -374,7 +374,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
   def testExpireQuotaSensors(): Unit = {
     val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     try {
-      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT),
+      clientQuotaManager.updateQuota(None, Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING),
         Some(new Quota(500, true)))
 
       maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
@@ -401,7 +401,7 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     val clientId = "client@#$%"
     try {
-      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.DEFAULT), Some(ConfigEntityName.DEFAULT),
+      clientQuotaManager.updateQuota(None, Some(ZooKeeperInternals.DEFAULT_STRING), Some(ZooKeeperInternals.DEFAULT_STRING),
         Some(new Quota(500, true)))
 
       maybeRecord(clientQuotaManager, "ANONYMOUS", clientId, 100)
@@ -421,6 +421,6 @@ class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
     // The class under test expects only sanitized client configs. We pass both the default value (which should not be
     // sanitized to ensure it remains unique) and non-default values, so we need to take care in generating the sanitized
     // client ID
-    def sanitizedConfigClientId = configClientId.map(x => if (x == ConfigEntityName.DEFAULT) ConfigEntityName.DEFAULT else Sanitizer.sanitize(x))
+    def sanitizedConfigClientId = configClientId.map(x => if (x == ZooKeeperInternals.DEFAULT_STRING) ZooKeeperInternals.DEFAULT_STRING else Sanitizer.sanitize(x))
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala b/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
index a24c5db6a9c9a..e86e613656ede 100644
--- a/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
@@ -20,7 +20,6 @@ package kafka.server
 import java.net.InetAddress
 import java.util
 import java.util.concurrent.{ExecutionException, TimeUnit}
-
 import kafka.test.ClusterInstance
 import kafka.test.annotation.{ClusterTest, ClusterTestDefaults, Type}
 import kafka.test.junit.ClusterTestExtensions
@@ -31,7 +30,7 @@ import org.apache.kafka.common.errors.{InvalidRequestException, UnsupportedVersi
 import org.apache.kafka.common.internals.KafkaFutureImpl
 import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
 import org.apache.kafka.common.requests.{AlterClientQuotasRequest, AlterClientQuotasResponse, DescribeClientQuotasRequest, DescribeClientQuotasResponse}
-import org.apache.kafka.server.config.ConfigEntityName
+import org.apache.kafka.server.config.ZooKeeperInternals
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Tag
 import org.junit.jupiter.api.extension.ExtendWith
@@ -527,7 +526,7 @@ class ClientQuotasRequestTest(cluster: ClusterInstance) {
   def testClientQuotasWithDefaultName(): Unit = {
     // An entity using the name associated with the default entity name. The entity's name should be sanitized so
     // that it does not conflict with the default entity name.
-    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.CLIENT_ID -> ConfigEntityName.DEFAULT)).asJava)
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.CLIENT_ID -> ZooKeeperInternals.DEFAULT_STRING)).asJava)
     alterEntityQuotas(entity, Map((ProducerByteRateProp -> Some(20000.0))), validateOnly = false)
     verifyDescribeEntityQuotas(entity, Map((ProducerByteRateProp -> 20000.0)))
 
diff --git a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
index fdc1dd583bf8e..1d5c69cfadf59 100644
--- a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
@@ -44,7 +44,7 @@ import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity}
 import org.apache.kafka.common.record.{CompressionType, RecordVersion}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.server.common.MetadataVersion.IBP_3_0_IV1
-import org.apache.kafka.server.config.{ConfigEntityName, ConfigType}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
 import org.apache.kafka.storage.internals.log.LogConfig
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{Test, Timeout}
@@ -321,7 +321,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
       val ipDefaultProps = new Properties()
       ipDefaultProps.put(QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG, "20")
-      adminZkClient.changeIpConfig(ConfigEntityName.DEFAULT, ipDefaultProps)
+      adminZkClient.changeIpConfig(ZooKeeperInternals.DEFAULT_STRING, ipDefaultProps)
 
       val ipOverrideProps = new Properties()
       ipOverrideProps.put(QuotaConfigs.IP_CONNECTION_RATE_OVERRIDE_CONFIG, "10")
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
index 2aff7edf7135f..971fdecbb8333 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
@@ -142,15 +142,17 @@ class ZkConfigMigrationClientTest extends ZkMigrationTestHarness {
     RecordTestUtils.replayAllBatches(delta, batches)
     val image = delta.apply()
 
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("user" -> "").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("user" -> "user1").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("user" -> "", "client-id" -> "").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("user" -> "", "client-id" -> "clientA").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("client-id" -> "").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("client-id" -> "clientB").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("ip" -> "1.1.1.1").asJava)))
-    assertTrue(image.entities().containsKey(new ClientQuotaEntity(Map("ip" -> "").asJava)))
+    assertEquals(new util.HashSet[ClientQuotaEntity](java.util.Arrays.asList(
+      new ClientQuotaEntity(Map("user" -> null.asInstanceOf[String]).asJava),
+      new ClientQuotaEntity(Map("user" -> "user1").asJava),
+      new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava),
+      new ClientQuotaEntity(Map("user" -> null.asInstanceOf[String], "client-id" -> null.asInstanceOf[String]).asJava),
+      new ClientQuotaEntity(Map("user" -> null.asInstanceOf[String], "client-id" -> "clientA").asJava),
+      new ClientQuotaEntity(Map("client-id" -> null.asInstanceOf[String]).asJava),
+      new ClientQuotaEntity(Map("client-id" -> "clientB").asJava),
+      new ClientQuotaEntity(Map("ip" -> "1.1.1.1").asJava),
+      new ClientQuotaEntity(Map("ip" -> null.asInstanceOf[String]).asJava))),
+      image.entities().keySet())
   }
 
   @Test
@@ -186,7 +188,7 @@ class ZkConfigMigrationClientTest extends ZkMigrationTestHarness {
     assertEquals(4, migrationState.migrationZkVersion())
 
     migrationState = writeClientQuotaAndVerify(migrationClient, adminZkClient, migrationState,
-      Map(ClientQuotaEntity.USER -> ""),
+      Map(ClientQuotaEntity.USER -> null.asInstanceOf[String]),
       Map(QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG -> 200.0),
       ConfigType.USER, "<default>")
     assertEquals(5, migrationState.migrationZkVersion())
diff --git a/server-common/src/main/java/org/apache/kafka/server/config/ConfigEntityName.java b/server-common/src/main/java/org/apache/kafka/server/config/ZooKeeperInternals.java
similarity index 63%
rename from server-common/src/main/java/org/apache/kafka/server/config/ConfigEntityName.java
rename to server-common/src/main/java/org/apache/kafka/server/config/ZooKeeperInternals.java
index 7188988254078..72ac563ae52f8 100644
--- a/server-common/src/main/java/org/apache/kafka/server/config/ConfigEntityName.java
+++ b/server-common/src/main/java/org/apache/kafka/server/config/ZooKeeperInternals.java
@@ -16,6 +16,12 @@
  */
 package org.apache.kafka.server.config;
 
-public class ConfigEntityName {
-    public static final String DEFAULT = "<default>";
+public class ZooKeeperInternals {
+    /**
+     * This string is used in ZooKeeper in several places to indicate a default entity type.
+     * For example, default user quotas are stored under /config/users/&ltdefault&gt
+     * Note that AdminClient does <b>not</b> use this to indicate a default, nor do records in KRaft mode.
+     * This constant will go away in Apache Kafka 4.0 with the end of ZK mode.
+     */
+    public static final String DEFAULT_STRING = "<default>";
 }

From 3e64f8fa1b0f2ef09f1d5ff8de6539992d295799 Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Wed, 27 Mar 2024 10:19:34 +0800
Subject: [PATCH 211/258] KAFKA-16385 Enhance documentation for retention.ms
 and retention.bytes configurations (#15588)

Reviewers: Igor Soarez <soarez@apple.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../java/org/apache/kafka/common/config/TopicConfig.java  | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java b/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java
index dcb772c18a4e2..b6967503c9605 100755
--- a/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java
@@ -67,13 +67,17 @@ public class TopicConfig {
         "(which consists of log segments) can grow to before we will discard old log segments to free up space if we " +
         "are using the \"delete\" retention policy. By default there is no size limit only a time limit. " +
         "Since this limit is enforced at the partition level, multiply it by the number of partitions to compute " +
-        "the topic retention in bytes.";
+        "the topic retention in bytes. Additionally, retention.bytes configuration " +
+        "operates independently of \"segment.ms\" and \"segment.bytes\" configurations. " +
+        "Moreover, it triggers the rolling of new segment if the retention.bytes is configured to zero.";
 
     public static final String RETENTION_MS_CONFIG = "retention.ms";
     public static final String RETENTION_MS_DOC = "This configuration controls the maximum time we will retain a " +
         "log before we will discard old log segments to free up space if we are using the " +
         "\"delete\" retention policy. This represents an SLA on how soon consumers must read " +
-        "their data. If set to -1, no time limit is applied.";
+        "their data. If set to -1, no time limit is applied. Additionally, retention.ms configuration " +
+        "operates independently of \"segment.ms\" and \"segment.bytes\" configurations. " +
+        "Moreover, it triggers the rolling of new segment if the retention.ms condition is satisfied.";
 
     public static final String REMOTE_LOG_STORAGE_ENABLE_CONFIG = "remote.storage.enable";
     public static final String REMOTE_LOG_STORAGE_ENABLE_DOC = "To enable tiered storage for a topic, set this configuration as true. " +

From ae44a08051e9f0db78051b159e335f075aab1fe1 Mon Sep 17 00:00:00 2001
From: "Cheng-Kai, Zhang" <kai821104@gmail.com>
Date: Wed, 27 Mar 2024 11:11:09 +0800
Subject: [PATCH 212/258] MINOR: add docker usage documentation link in
 README.md (#15600)

We have a detail usage guide for running Kafka in docker. However, it might be easy for people to miss it if they only scan through the README.

This PR add a basic example command for quick start and link directing users to the detailed usage guide

Reviewers: Luke Chen <showuon@gmail.com>
---
 README.md | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/README.md b/README.md
index d8ee34b6188a3..27ce0dc0bce64 100644
--- a/README.md
+++ b/README.md
@@ -92,15 +92,25 @@ fail due to code changes. You can just run:
 
 ### Running a Kafka broker in KRaft mode
 
+Using compiled files:
+
     KAFKA_CLUSTER_ID="$(./bin/kafka-storage.sh random-uuid)"
     ./bin/kafka-storage.sh format -t $KAFKA_CLUSTER_ID -c config/kraft/server.properties
     ./bin/kafka-server-start.sh config/kraft/server.properties
 
+Using docker image:
+
+    docker run -p 9092:9092 apache/kafka:3.7.0
+
 ### Running a Kafka broker in ZooKeeper mode
 
+Using compiled files:
+
     ./bin/zookeeper-server-start.sh config/zookeeper.properties
     ./bin/kafka-server-start.sh config/server.properties
 
+>Since ZooKeeper mode is already deprecated and planned to be removed in Apache Kafka 4.0, the docker image only supports running in KRaft mode
+
 ### Cleaning the build ###
     ./gradlew clean
 

From 932647606504125e5c3ba0ae9470b4af335a0885 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Wed, 27 Mar 2024 11:13:54 +0800
Subject: [PATCH 213/258] KAFKA-16391: remove .lock file when FileLock#destroy
 (#15568)

Currently, server adds a .lock file to each log folder. The file is useless after server is down.

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../src/main/scala/kafka/utils/FileLock.scala |  3 +++
 .../scala/unit/kafka/log/LogManagerTest.scala | 20 +++++++++++++++++++
 .../unit/kafka/raft/RaftManagerTest.scala     |  2 +-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/kafka/utils/FileLock.scala b/core/src/main/scala/kafka/utils/FileLock.scala
index 7a5afc9fb958b..f566c1410aa00 100644
--- a/core/src/main/scala/kafka/utils/FileLock.scala
+++ b/core/src/main/scala/kafka/utils/FileLock.scala
@@ -76,6 +76,9 @@ class FileLock(val file: File) extends Logging {
   def destroy(): Unit = {
     this synchronized {
       unlock()
+      if (file.exists() && file.delete()) {
+        trace(s"Delete ${file.getAbsolutePath}")
+      }
       channel.close()
     }
   }
diff --git a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
index ce6e2a3b04ee3..95b32361cb0fa 100755
--- a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
@@ -1303,6 +1303,26 @@ class LogManagerTest {
         createLeaderAndIsrRequestForStrayDetection(present),
         onDisk.map(mockLog(_))).toSet)
   }
+
+  /**
+   * Test LogManager takes file lock by default and the lock is released after shutdown.
+   */
+  @Test
+  def testLock(): Unit = {
+    val tmpLogDir = TestUtils.tempDir()
+    val tmpLogManager = createLogManager(Seq(tmpLogDir))
+
+    try {
+      // ${tmpLogDir}.lock is acquired by tmpLogManager
+      val fileLock = new FileLock(new File(tmpLogDir, LogManager.LockFileName))
+      assertFalse(fileLock.tryLock())
+    } finally {
+      // ${tmpLogDir}.lock is removed after shutdown
+      tmpLogManager.shutdown()
+      val f = new File(tmpLogDir, LogManager.LockFileName)
+      assertFalse(f.exists())
+    }
+  }
 }
 
 object LogManagerTest {
diff --git a/core/src/test/scala/unit/kafka/raft/RaftManagerTest.scala b/core/src/test/scala/unit/kafka/raft/RaftManagerTest.scala
index 711426468c573..e716bb452d915 100644
--- a/core/src/test/scala/unit/kafka/raft/RaftManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/raft/RaftManagerTest.scala
@@ -178,7 +178,7 @@ class RaftManagerTest {
   }
 
   private def fileLocked(path: Path): Boolean = {
-    TestUtils.resource(FileChannel.open(path, StandardOpenOption.WRITE)) { channel =>
+    TestUtils.resource(FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { channel =>
       try {
         Option(channel.tryLock()).foreach(_.close())
         false

From 4cb20379c7abc07bb18075674553abc11954f286 Mon Sep 17 00:00:00 2001
From: Federico Valeri <fedevaleri@gmail.com>
Date: Wed, 27 Mar 2024 09:31:27 +0100
Subject: [PATCH 214/258] MINOR: Add retry mechanism to EOS example (#15561)

In the initial EOS example, a retry logic was implemented within the resetToLastCommittedPositions method. During refactoring, this logic was removed becasue a poison pill prevented the example from reaching the final phase of consuming from the output topic.

In this change, I suggest to add it back, but with a retry limit defined as MAX_RETRIES. Once this limit is reached, the problematic batch will be logged and skipped, allowing the processor to move on and process remaining records. If some records are skipped, the example will still hit the hard timeout (2 minutes), but after consuming all processed records.

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../main/java/kafka/examples/Consumer.java    |  2 +-
 .../examples/ExactlyOnceMessageProcessor.java | 55 +++++++++++++++++--
 .../main/java/kafka/examples/Producer.java    |  2 +-
 3 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/examples/src/main/java/kafka/examples/Consumer.java b/examples/src/main/java/kafka/examples/Consumer.java
index bd652e0a32ef1..aa971812f8758 100644
--- a/examples/src/main/java/kafka/examples/Consumer.java
+++ b/examples/src/main/java/kafka/examples/Consumer.java
@@ -109,7 +109,7 @@ public void run() {
                 }
             }
         } catch (Throwable e) {
-            Utils.printOut("Unhandled exception");
+            Utils.printErr("Unhandled exception");
             e.printStackTrace();
         }
         Utils.printOut("Fetched %d records", numRecords - remainingRecords);
diff --git a/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java b/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
index cf12cc9f15812..15488c9c47d23 100644
--- a/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
+++ b/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
@@ -36,6 +36,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -49,6 +50,8 @@
  * This class implements a read-process-write application.
  */
 public class ExactlyOnceMessageProcessor extends Thread implements ConsumerRebalanceListener, AutoCloseable {
+    private static final int MAX_RETRIES = 5;
+    
     private final String bootstrapServers;
     private final String inputTopic;
     private final String outputTopic;
@@ -103,19 +106,21 @@ public ExactlyOnceMessageProcessor(String threadName,
 
     @Override
     public void run() {
+        int retries = 0;
         int processedRecords = 0;
         long remainingRecords = Long.MAX_VALUE;
+
         // it is recommended to have a relatively short txn timeout in order to clear pending offsets faster
         int transactionTimeoutMs = 10_000;
         // consumer must be in read_committed mode, which means it won't be able to read uncommitted data
         boolean readCommitted = true;
+
         try (KafkaProducer<Integer, String> producer = new Producer("processor-producer", bootstrapServers, outputTopic,
                 true, transactionalId, true, -1, transactionTimeoutMs, null).createKafkaProducer();
              KafkaConsumer<Integer, String> consumer = new Consumer("processor-consumer", bootstrapServers, inputTopic,
                  "processor-group", Optional.of(groupInstanceId), readCommitted, -1, null).createKafkaConsumer()) {
             // called first and once to fence zombies and abort any pending transaction
             producer.initTransactions();
-
             consumer.subscribe(singleton(inputTopic), this);
 
             Utils.printOut("Processing new records");
@@ -140,6 +145,7 @@ public void run() {
                         // commit the transaction including offsets
                         producer.commitTransaction();
                         processedRecords += records.count();
+                        retries = 0;
                     }
                 } catch (AuthorizationException | UnsupportedVersionException | ProducerFencedException
                          | FencedInstanceIdException | OutOfOrderSequenceException | SerializationException e) {
@@ -151,18 +157,21 @@ public void run() {
                     Utils.printOut("Invalid or no offset found, using latest");
                     consumer.seekToEnd(emptyList());
                     consumer.commitSync();
+                    retries = 0;
                 } catch (KafkaException e) {
-                    // abort the transaction and try to continue
-                    Utils.printOut("Aborting transaction: %s", e);
+                    // abort the transaction
+                    Utils.printOut("Aborting transaction: %s", e.getMessage());
                     producer.abortTransaction();
+                    retries = maybeRetry(retries, consumer);
                 }
+
                 remainingRecords = getRemainingRecords(consumer);
                 if (remainingRecords != Long.MAX_VALUE) {
                     Utils.printOut("Remaining records: %d", remainingRecords);
                 }
             }
         } catch (Throwable e) {
-            Utils.printOut("Unhandled exception");
+            Utils.printErr("Unhandled exception");
             e.printStackTrace();
         }
         Utils.printOut("Processed %d records", processedRecords);
@@ -215,6 +224,44 @@ private long getRemainingRecords(KafkaConsumer<Integer, String> consumer) {
         }).sum();
     }
 
+    /**
+     * When we get a generic {@code KafkaException} while processing records, we retry up to {@code MAX_RETRIES} times.
+     * If we exceed this threshold, we log an error and move on to the next batch of records.
+     * In a real world application you may want to to send these records to a dead letter topic (DLT) for further processing.
+     * 
+     * @param retries Current number of retries
+     * @param consumer Consumer instance
+     * @return Updated number of retries
+     */
+    private int maybeRetry(int retries, KafkaConsumer<Integer, String> consumer) {
+        if (retries < 0) {
+            Utils.printErr("The number of retries must be greater than zero");
+            shutdown();
+        }
+        
+        if (retries < MAX_RETRIES) {
+            // retry: reset fetch offset
+            // the consumer fetch position needs to be restored to the committed offset before the transaction started
+            Map<TopicPartition, OffsetAndMetadata> committed = consumer.committed(consumer.assignment());
+            consumer.assignment().forEach(tp -> {
+                OffsetAndMetadata offsetAndMetadata = committed.get(tp);
+                if (offsetAndMetadata != null) {
+                    consumer.seek(tp, offsetAndMetadata.offset());
+                } else {
+                    consumer.seekToBeginning(Collections.singleton(tp));
+                }
+            });
+            retries++;
+        } else {
+            // continue: skip records
+            // the consumer fetch position needs to be committed as if records were processed successfully
+            Utils.printErr("Skipping records after %d retries", MAX_RETRIES);
+            consumer.commitSync();
+            retries = 0;
+        }
+        return retries;
+    }
+
     @Override
     public void close() throws Exception {
         if (producer != null) {
diff --git a/examples/src/main/java/kafka/examples/Producer.java b/examples/src/main/java/kafka/examples/Producer.java
index 36a2583954ca8..d9d454c15593a 100644
--- a/examples/src/main/java/kafka/examples/Producer.java
+++ b/examples/src/main/java/kafka/examples/Producer.java
@@ -89,7 +89,7 @@ public void run() {
                 sentRecords++;
             }
         } catch (Throwable e) {
-            Utils.printOut("Unhandled exception");
+            Utils.printErr("Unhandled exception");
             e.printStackTrace();
         }
         Utils.printOut("Sent %d records", sentRecords);

From 4a33bc8bbf7fc042de4a61c9dd15e3208dac3439 Mon Sep 17 00:00:00 2001
From: Dongnuo Lyu <139248811+dongnuo123@users.noreply.github.com>
Date: Wed, 27 Mar 2024 04:57:24 -0400
Subject: [PATCH 215/258] KAFKA-16353: Offline protocol migration integration
 tests (#15492)

This patch adds integration tests for offline protocol migration.

Reviewers: David Jacot <djacot@confluent.io>
---
 .../ConsumerProtocolMigrationTest.scala       | 250 ++++++++++++++++++
 .../server/DeleteGroupsRequestTest.scala      |   6 +-
 .../server/DescribeGroupsRequestTest.scala    |   4 +-
 .../GroupCoordinatorBaseRequestTest.scala     |  12 +-
 .../kafka/server/HeartbeatRequestTest.scala   |   4 +-
 .../kafka/server/LeaveGroupRequestTest.scala  |   4 +-
 .../kafka/server/ListGroupsRequestTest.scala  |   6 +-
 .../server/OffsetCommitRequestTest.scala      |   6 +-
 .../server/OffsetDeleteRequestTest.scala      |   6 +-
 .../kafka/server/OffsetFetchRequestTest.scala |  18 +-
 .../kafka/server/SyncGroupRequestTest.scala   |   4 +-
 11 files changed, 286 insertions(+), 34 deletions(-)
 create mode 100644 core/src/test/scala/unit/kafka/server/ConsumerProtocolMigrationTest.scala

diff --git a/core/src/test/scala/unit/kafka/server/ConsumerProtocolMigrationTest.scala b/core/src/test/scala/unit/kafka/server/ConsumerProtocolMigrationTest.scala
new file mode 100644
index 0000000000000..789a3dfb281dc
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ConsumerProtocolMigrationTest.scala
@@ -0,0 +1,250 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import kafka.test.ClusterInstance
+import kafka.test.annotation.{ClusterConfigProperty, ClusterTest, ClusterTestDefaults, Type}
+import kafka.test.junit.ClusterTestExtensions
+import org.apache.kafka.common.message.ListGroupsResponseData
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.coordinator.group.Group
+import org.apache.kafka.coordinator.group.classic.ClassicGroupState
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroup.ConsumerGroupState
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Tag
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.api.extension.ExtendWith
+
+@Timeout(120)
+@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
+@ClusterTestDefaults(clusterType = Type.KRAFT, brokers = 1)
+@Tag("integration")
+class ConsumerProtocolMigrationTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
+  @ClusterTest(serverProperties = Array(
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
+    new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
+    new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
+  ))
+  def testUpgradeFromEmptyClassicToConsumerGroup(): Unit = {
+    // Creates the __consumer_offsets topics because it won't be created automatically
+    // in this test because it does not use FindCoordinator API.
+    createOffsetsTopic()
+
+    // Create the topic.
+    createTopic(
+      topic = "foo",
+      numPartitions = 3
+    )
+
+    // Create a classic group by joining a member.
+    val groupId = "grp"
+    val (memberId, _) = joinDynamicConsumerGroupWithOldProtocol(groupId)
+
+    // The joining request from a consumer group member is rejected.
+    val responseData = consumerGroupHeartbeat(
+      groupId = groupId,
+      rebalanceTimeoutMs = 5 * 60 * 1000,
+      subscribedTopicNames = List("foo"),
+      topicPartitions = List.empty,
+      expectedError = Errors.GROUP_ID_NOT_FOUND
+    )
+    assertEquals("Group grp is not a consumer group.", responseData.errorMessage)
+
+    // The member leaves the group.
+    leaveGroup(
+      groupId = groupId,
+      memberId = memberId,
+      useNewProtocol = false,
+      version = ApiKeys.LEAVE_GROUP.latestVersion(isUnstableApiEnabled)
+    )
+
+    // Verify that the group is empty.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setProtocolType("consumer")
+          .setGroupState(ClassicGroupState.EMPTY.toString)
+          .setGroupType(Group.GroupType.CLASSIC.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CLASSIC.toString)
+      )
+    )
+
+    // The joining request with a consumer group member is accepted.
+    consumerGroupHeartbeat(
+      groupId = groupId,
+      memberId = memberId,
+      rebalanceTimeoutMs = 5 * 60 * 1000,
+      subscribedTopicNames = List("foo"),
+      topicPartitions = List.empty,
+      expectedError = Errors.NONE
+    )
+
+    // The group has become a consumer group.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setProtocolType("consumer")
+          .setGroupState(ConsumerGroupState.STABLE.toString)
+          .setGroupType(Group.GroupType.CONSUMER.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CONSUMER.toString)
+      )
+    )
+  }
+
+  @ClusterTest(serverProperties = Array(
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
+    new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
+    new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
+  ))
+  def testDowngradeFromEmptyConsumerToClassicGroup(): Unit = {
+    // Creates the __consumer_offsets topics because it won't be created automatically
+    // in this test because it does not use FindCoordinator API.
+    createOffsetsTopic()
+
+    // Create the topic.
+    createTopic(
+      topic = "foo",
+      numPartitions = 3
+    )
+
+    // Create a consumer group by joining a member.
+    val groupId = "grp"
+    val (memberId, _) = joinConsumerGroupWithNewProtocol(groupId)
+
+    // The joining request from a classic group member is rejected.
+    val joinGroupResponseData = sendJoinRequest(groupId = groupId)
+    assertEquals(Errors.GROUP_ID_NOT_FOUND.code, joinGroupResponseData.errorCode)
+
+    // The member leaves the group.
+    leaveGroup(
+      groupId = groupId,
+      memberId = memberId,
+      useNewProtocol = true,
+      version = ApiKeys.CONSUMER_GROUP_HEARTBEAT.latestVersion(isUnstableApiEnabled)
+    )
+
+    // Verify that the group is empty.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setProtocolType("consumer")
+          .setGroupState(ClassicGroupState.EMPTY.toString)
+          .setGroupType(Group.GroupType.CONSUMER.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CONSUMER.toString)
+      )
+    )
+
+    // The joining request with a classic group member is accepted.
+    joinDynamicConsumerGroupWithOldProtocol(groupId = groupId)
+
+    // The group has become a classic group.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setProtocolType("consumer")
+          .setGroupState(ClassicGroupState.STABLE.toString)
+          .setGroupType(Group.GroupType.CLASSIC.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CLASSIC.toString)
+      )
+    )
+  }
+
+  @ClusterTest(serverProperties = Array(
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
+    new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
+    new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
+  ))
+  def testUpgradeFromSimpleGroupToConsumerGroup(): Unit = {
+    // Creates the __consumer_offsets topics because it won't be created automatically
+    // in this test because it does not use FindCoordinator API.
+    createOffsetsTopic()
+
+    val topicName = "foo"
+    // Create the topic.
+    createTopic(
+      topic = topicName,
+      numPartitions = 3
+    )
+
+    // An admin client commits offsets and creates the simple group.
+    val groupId = "group-id"
+    commitOffset(
+      groupId = groupId,
+      memberId = "member-id",
+      memberEpoch = -1,
+      topic = topicName,
+      partition = 0,
+      offset = 1000L,
+      expectedError = Errors.NONE,
+      version = ApiKeys.OFFSET_COMMIT.latestVersion(isUnstableApiEnabled)
+    )
+
+    // Verify that the simple group is created.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setGroupState(ClassicGroupState.EMPTY.toString)
+          .setGroupType(Group.GroupType.CLASSIC.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CLASSIC.toString)
+      )
+    )
+
+    // The joining request with a consumer group member is accepted.
+    consumerGroupHeartbeat(
+      groupId = groupId,
+      rebalanceTimeoutMs = 5 * 60 * 1000,
+      subscribedTopicNames = List(topicName),
+      topicPartitions = List.empty,
+      expectedError = Errors.NONE
+    )
+
+    // The group has become a consumer group.
+    assertEquals(
+      List(
+        new ListGroupsResponseData.ListedGroup()
+          .setGroupId(groupId)
+          .setProtocolType("consumer")
+          .setGroupState(ConsumerGroupState.STABLE.toString)
+          .setGroupType(Group.GroupType.CONSUMER.toString)
+      ),
+      listGroups(
+        statesFilter = List.empty,
+        typesFilter = List(Group.GroupType.CONSUMER.toString)
+      )
+    )
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DeleteGroupsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DeleteGroupsRequestTest.scala
index f4bff8c7e69ba..92cea4c8d3d05 100644
--- a/core/src/test/scala/unit/kafka/server/DeleteGroupsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DeleteGroupsRequestTest.scala
@@ -32,7 +32,7 @@ import org.junit.jupiter.api.extension.ExtendWith
 @Tag("integration")
 class DeleteGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -43,7 +43,7 @@ class DeleteGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -52,7 +52,7 @@ class DeleteGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/DescribeGroupsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeGroupsRequestTest.scala
index 8305fc47599e6..45f967ffc91ca 100644
--- a/core/src/test/scala/unit/kafka/server/DescribeGroupsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DescribeGroupsRequestTest.scala
@@ -34,7 +34,7 @@ import scala.jdk.CollectionConverters._
 @Tag("integration")
 class DescribeGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -43,7 +43,7 @@ class DescribeGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinat
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/GroupCoordinatorBaseRequestTest.scala b/core/src/test/scala/unit/kafka/server/GroupCoordinatorBaseRequestTest.scala
index 6a9da0adb67c0..f2269d85b56e9 100644
--- a/core/src/test/scala/unit/kafka/server/GroupCoordinatorBaseRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/GroupCoordinatorBaseRequestTest.scala
@@ -78,7 +78,8 @@ class GroupCoordinatorBaseRequestTest(cluster: ClusterInstance) {
   }
 
   protected def isNewGroupCoordinatorEnabled: Boolean = {
-    cluster.config.serverProperties.getProperty("group.coordinator.new.enable") == "true"
+    cluster.config.serverProperties.getProperty("group.coordinator.new.enable") == "true" ||
+      cluster.config.serverProperties.getProperty("group.coordinator.rebalance.protocols").contains("consumer")
   }
 
   protected def commitOffset(
@@ -406,7 +407,7 @@ class GroupCoordinatorBaseRequestTest(cluster: ClusterInstance) {
   protected def listGroups(
     statesFilter: List[String],
     typesFilter: List[String],
-    version: Short
+    version: Short = ApiKeys.LIST_GROUPS.latestVersion(isUnstableApiEnabled)
   ): List[ListGroupsResponseData.ListedGroup] = {
     val request = new ListGroupsRequest.Builder(
       new ListGroupsRequestData()
@@ -463,7 +464,8 @@ class GroupCoordinatorBaseRequestTest(cluster: ClusterInstance) {
     rebalanceTimeoutMs: Int = -1,
     serverAssignor: String = null,
     subscribedTopicNames: List[String] = null,
-    topicPartitions: List[ConsumerGroupHeartbeatRequestData.TopicPartitions] = null
+    topicPartitions: List[ConsumerGroupHeartbeatRequestData.TopicPartitions] = null,
+    expectedError: Errors = Errors.NONE
   ): ConsumerGroupHeartbeatResponseData = {
     val consumerGroupHeartbeatRequest = new ConsumerGroupHeartbeatRequest.Builder(
       new ConsumerGroupHeartbeatRequestData()
@@ -484,7 +486,7 @@ class GroupCoordinatorBaseRequestTest(cluster: ClusterInstance) {
     var consumerGroupHeartbeatResponse: ConsumerGroupHeartbeatResponse = null
     TestUtils.waitUntilTrue(() => {
       consumerGroupHeartbeatResponse = connectAndReceive[ConsumerGroupHeartbeatResponse](consumerGroupHeartbeatRequest)
-      consumerGroupHeartbeatResponse.data.errorCode == Errors.NONE.code
+      consumerGroupHeartbeatResponse.data.errorCode == expectedError.code
     }, msg = s"Could not heartbeat successfully. Last response $consumerGroupHeartbeatResponse.")
 
     consumerGroupHeartbeatResponse.data
@@ -507,7 +509,7 @@ class GroupCoordinatorBaseRequestTest(cluster: ClusterInstance) {
     groupInstanceIds: List[String] = null,
     expectedLeaveGroupError: Errors,
     expectedMemberErrors: List[Errors],
-    version: Short
+    version: Short = ApiKeys.LEAVE_GROUP.latestVersion(isUnstableApiEnabled)
   ): Unit = {
     val leaveGroupRequest = new LeaveGroupRequest.Builder(
       groupId,
diff --git a/core/src/test/scala/unit/kafka/server/HeartbeatRequestTest.scala b/core/src/test/scala/unit/kafka/server/HeartbeatRequestTest.scala
index c3e9d6cda166d..80a308ef9b715 100644
--- a/core/src/test/scala/unit/kafka/server/HeartbeatRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/HeartbeatRequestTest.scala
@@ -38,7 +38,7 @@ import scala.concurrent.Future
 @Tag("integration")
 class HeartbeatRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -47,7 +47,7 @@ class HeartbeatRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBas
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/LeaveGroupRequestTest.scala b/core/src/test/scala/unit/kafka/server/LeaveGroupRequestTest.scala
index 1958c40a4e3cc..3eaff93d02a5c 100644
--- a/core/src/test/scala/unit/kafka/server/LeaveGroupRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LeaveGroupRequestTest.scala
@@ -32,7 +32,7 @@ import org.junit.jupiter.api.extension.ExtendWith
 @Tag("integration")
 class LeaveGroupRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -41,7 +41,7 @@ class LeaveGroupRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBa
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/ListGroupsRequestTest.scala b/core/src/test/scala/unit/kafka/server/ListGroupsRequestTest.scala
index 2c30c0ccae770..da00a765120f0 100644
--- a/core/src/test/scala/unit/kafka/server/ListGroupsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ListGroupsRequestTest.scala
@@ -34,7 +34,7 @@ import org.junit.jupiter.api.extension.ExtendWith
 @Tag("integration")
 class ListGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -45,7 +45,7 @@ class ListGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBa
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -54,7 +54,7 @@ class ListGroupsRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBa
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/OffsetCommitRequestTest.scala b/core/src/test/scala/unit/kafka/server/OffsetCommitRequestTest.scala
index 2ff776aa42a4f..346b17028fb92 100644
--- a/core/src/test/scala/unit/kafka/server/OffsetCommitRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/OffsetCommitRequestTest.scala
@@ -31,7 +31,7 @@ import org.junit.jupiter.api.extension.ExtendWith
 class OffsetCommitRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -42,7 +42,7 @@ class OffsetCommitRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -51,7 +51,7 @@ class OffsetCommitRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/OffsetDeleteRequestTest.scala b/core/src/test/scala/unit/kafka/server/OffsetDeleteRequestTest.scala
index 719697da1314f..8b4561d8a1c41 100644
--- a/core/src/test/scala/unit/kafka/server/OffsetDeleteRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/OffsetDeleteRequestTest.scala
@@ -30,7 +30,7 @@ import org.junit.jupiter.api.extension.ExtendWith
 @Tag("integration")
 class OffsetDeleteRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -41,7 +41,7 @@ class OffsetDeleteRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -50,7 +50,7 @@ class OffsetDeleteRequestTest(cluster: ClusterInstance) extends GroupCoordinator
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
diff --git a/core/src/test/scala/unit/kafka/server/OffsetFetchRequestTest.scala b/core/src/test/scala/unit/kafka/server/OffsetFetchRequestTest.scala
index 828f8ab0aaf6f..79ec4029c9a77 100644
--- a/core/src/test/scala/unit/kafka/server/OffsetFetchRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/OffsetFetchRequestTest.scala
@@ -39,7 +39,7 @@ import scala.jdk.CollectionConverters._
 class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -50,7 +50,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -61,7 +61,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -72,7 +72,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -83,7 +83,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -94,7 +94,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -105,7 +105,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -116,7 +116,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
@@ -127,7 +127,7 @@ class OffsetFetchRequestTest(cluster: ClusterInstance) extends GroupCoordinatorB
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "group.consumer.max.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "group.consumer.session.timeout.ms", value = "600000"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
diff --git a/core/src/test/scala/unit/kafka/server/SyncGroupRequestTest.scala b/core/src/test/scala/unit/kafka/server/SyncGroupRequestTest.scala
index c1c528fd8f787..ac0d068b6aa93 100644
--- a/core/src/test/scala/unit/kafka/server/SyncGroupRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/SyncGroupRequestTest.scala
@@ -39,7 +39,7 @@ import scala.concurrent.{Await, Future}
 @Tag("integration")
 class SyncGroupRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBaseRequestTest(cluster) {
   @ClusterTest(serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "true"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic,consumer"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))
@@ -48,7 +48,7 @@ class SyncGroupRequestTest(cluster: ClusterInstance) extends GroupCoordinatorBas
   }
 
   @ClusterTest(clusterType = Type.ALL, serverProperties = Array(
-    new ClusterConfigProperty(key = "group.coordinator.new.enable", value = "false"),
+    new ClusterConfigProperty(key = "group.coordinator.rebalance.protocols", value = "classic"),
     new ClusterConfigProperty(key = "offsets.topic.num.partitions", value = "1"),
     new ClusterConfigProperty(key = "offsets.topic.replication.factor", value = "1")
   ))

From eda0dc0eb0f0bb760f6b2807e2228604e2357df2 Mon Sep 17 00:00:00 2001
From: Andras Katona <41361962+akatona84@users.noreply.github.com>
Date: Wed, 27 Mar 2024 10:51:31 +0100
Subject: [PATCH 216/258] MINOR: Preventing running the :core tests twice when
 testing with coverage (#15580)

reportScoverage task (which was used previously as dependency of the registered coverage task)
creates a task for each Test task and executes them. There's unitTest, integrationTest and
the test tasks (which is just for executing both unit and integration), so reportScoverage
executes all three with their corresponding scoverage task, hence running all tests twice.
Solution is just to use the reportTestScoverage task as dependency.

Reviewers: Viktor Somogyi-Vass <viktorsomogyi@gmail.com>
---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 88d919b27c638..489473c885aef 100644
--- a/build.gradle
+++ b/build.gradle
@@ -748,7 +748,7 @@ subprojects {
   }
 
   if (userEnableTestCoverage) {
-    def coverageGen = it.path == ':core' ? 'reportScoverage' : 'jacocoTestReport'
+    def coverageGen = it.path == ':core' ? 'reportTestScoverage' : 'jacocoTestReport'
     tasks.register('reportCoverage').configure { dependsOn(coverageGen) }
   }
 

From ace2152d229da35263a7e64537ca0d0316415558 Mon Sep 17 00:00:00 2001
From: Philip Nee <pnee@confluent.io>
Date: Wed, 27 Mar 2024 04:10:29 -0700
Subject: [PATCH 217/258] =?UTF-8?q?KAFKA-16272:=20Update=20connect=5Fdistr?=
 =?UTF-8?q?ibuted=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20group=20p?=
 =?UTF-8?q?rotocol=20config=20(#15576)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Update connect_distributed_test.py to support KIP-848’s group protocol config. Not all tests are updated because only a subset of it is using the consumer directly.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>, Kirk True <ktrue@confluent.io>
---
 .../tests/connect/connect_distributed_test.py | 32 +++++++++++--------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/tests/kafkatest/tests/connect/connect_distributed_test.py b/tests/kafkatest/tests/connect/connect_distributed_test.py
index 145f29f51ebf1..e1686e809e05f 100644
--- a/tests/kafkatest/tests/connect/connect_distributed_test.py
+++ b/tests/kafkatest/tests/connect/connect_distributed_test.py
@@ -20,7 +20,7 @@
 from ducktape.cluster.remoteaccount import RemoteCommandError
 
 from kafkatest.services.zookeeper import ZookeeperService
-from kafkatest.services.kafka import KafkaService, config_property, quorum
+from kafkatest.services.kafka import KafkaService, config_property, quorum, consumer_group
 from kafkatest.services.connect import ConnectDistributedService, VerifiableSource, VerifiableSink, ConnectRestError, MockSink, MockSource
 from kafkatest.services.console_consumer import ConsoleConsumer
 from kafkatest.services.security.security_config import SecurityConfig
@@ -683,16 +683,17 @@ def test_file_source_and_sink(self, security_protocol, exactly_once_source, conn
     @matrix(
         clean=[True, False],
         connect_protocol=['sessioned', 'compatible', 'eager'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         clean=[True, False],
         connect_protocol=['sessioned', 'compatible', 'eager'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_bounce(self, clean, connect_protocol, metadata_quorum, use_new_coordinator=False):
+    def test_bounce(self, clean, connect_protocol, metadata_quorum, use_new_coordinator=False, group_protocol=None):
         """
         Validates that source and sink tasks that run continuously and produce a predictable sequence of messages
         run correctly and deliver messages exactly once when Kafka Connect workers undergo clean rolling bounces,
@@ -804,7 +805,8 @@ def test_bounce(self, clean, connect_protocol, metadata_quorum, use_new_coordina
         if not success:
             self.mark_for_collect(self.cc)
             # Also collect the data in the topic to aid in debugging
-            consumer_validator = ConsoleConsumer(self.test_context, 1, self.kafka, self.source.topic, consumer_timeout_ms=1000, print_key=True)
+            consumer_properties = consumer_group.maybe_set_group_protocol(group_protocol)
+            consumer_validator = ConsoleConsumer(self.test_context, 1, self.kafka, self.source.topic, consumer_timeout_ms=1000, print_key=True, consumer_properties=consumer_properties)
             consumer_validator.run()
             self.mark_for_collect(consumer_validator, "consumer_stdout")
 
@@ -814,16 +816,17 @@ def test_bounce(self, clean, connect_protocol, metadata_quorum, use_new_coordina
     @matrix(
         clean=[True, False],
         connect_protocol=['sessioned', 'compatible', 'eager'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         clean=[True, False],
         connect_protocol=['sessioned', 'compatible', 'eager'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_exactly_once_source(self, clean, connect_protocol, metadata_quorum, use_new_coordinator=False):
+    def test_exactly_once_source(self, clean, connect_protocol, metadata_quorum, use_new_coordinator=False, group_protocol=None):
         """
         Validates that source tasks run correctly and deliver messages exactly once
         when Kafka Connect workers undergo bounces, both clean and unclean.
@@ -878,7 +881,8 @@ def test_exactly_once_source(self, clean, connect_protocol, metadata_quorum, use
         self.source.stop()
         self.cc.stop()
 
-        consumer = ConsoleConsumer(self.test_context, 1, self.kafka, self.source.topic, message_validator=json.loads, consumer_timeout_ms=1000, isolation_level="read_committed")
+        consumer_properties = consumer_group.maybe_set_group_protocol(group_protocol)
+        consumer = ConsoleConsumer(self.test_context, 1, self.kafka, self.source.topic, message_validator=json.loads, consumer_timeout_ms=1000, isolation_level="read_committed", consumer_properties=consumer_properties)
         consumer.run()
         src_messages = consumer.messages_consumed[1]
 
@@ -921,15 +925,16 @@ def test_exactly_once_source(self, clean, connect_protocol, metadata_quorum, use
     @cluster(num_nodes=6)
     @matrix(
         connect_protocol=['sessioned', 'compatible', 'eager'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         connect_protocol=['sessioned', 'compatible', 'eager'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_transformations(self, connect_protocol, metadata_quorum, use_new_coordinator=False):
+    def test_transformations(self, connect_protocol, metadata_quorum, use_new_coordinator=False, group_protocol=None):
         self.CONNECT_PROTOCOL = connect_protocol
         self.setup_services(timestamp_type='CreateTime', include_filestream_connectors=True)
         self.cc.set_configs(lambda node: self.render("connect-distributed.properties", node=node))
@@ -959,7 +964,8 @@ def test_transformations(self, connect_protocol, metadata_quorum, use_new_coordi
         for node in self.cc.nodes:
             node.account.ssh("echo -e -n " + repr(self.FIRST_INPUTS) + " >> " + self.INPUT_FILE)
 
-        consumer = ConsoleConsumer(self.test_context, 1, self.kafka, self.TOPIC, consumer_timeout_ms=15000, print_timestamp=True)
+        consumer_properties = consumer_group.maybe_set_group_protocol(group_protocol)
+        consumer = ConsoleConsumer(self.test_context, 1, self.kafka, self.TOPIC, consumer_timeout_ms=15000, print_timestamp=True, consumer_properties=consumer_properties)
         consumer.run()
 
         assert len(consumer.messages_consumed[1]) == len(self.FIRST_INPUT_LIST)

From 6f38fe5e0a6e2fe85fec7cb9adc379061d35ce45 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Wed, 27 Mar 2024 17:37:01 +0300
Subject: [PATCH 218/258] KAFKA-14588 ZK configuration moved to ZkConfig
 (#15075)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle                                  |   5 +-
 checkstyle/import-control.xml                 |   1 +
 .../util/clusters/EmbeddedKafkaCluster.java   |   3 +-
 .../main/scala/kafka/admin/AclCommand.scala   |   8 +-
 .../scala/kafka/admin/ConfigCommand.scala     |   4 +-
 .../kafka/admin/ZkSecurityMigrator.scala      |   7 +-
 .../security/authorizer/AclAuthorizer.scala   |   7 +-
 .../main/scala/kafka/server/KafkaConfig.scala | 179 +++++-------------
 .../main/scala/kafka/server/KafkaServer.scala |  30 +--
 .../main/scala/kafka/zk/KafkaZkClient.scala   |   8 +-
 .../DescribeAuthorizedOperationsTest.scala    |   3 +-
 .../kafka/api/EndToEndAuthorizationTest.scala |   7 +-
 .../integration/kafka/api/MetricsTest.scala   |   3 +-
 .../api/PlaintextAdminIntegrationTest.scala   |   4 +-
 .../api/SaslMultiMechanismConsumerTest.scala  |   4 +-
 .../api/SaslPlainPlaintextConsumerTest.scala  |   4 +-
 .../api/SaslSslAdminIntegrationTest.scala     |   3 +-
 .../kafka/api/SaslSslConsumerTest.scala       |   4 +-
 .../kafka/api/SslAdminIntegrationTest.scala   |   5 +-
 .../DynamicBrokerReconfigurationTest.scala    |   4 +-
 .../KafkaServerKRaftRegistrationTest.scala    |   5 +-
 ...nersWithSameSecurityProtocolBaseTest.scala |   3 +-
 .../kafka/zk/ZkMigrationIntegrationTest.scala |  16 +-
 .../scala/unit/kafka/KafkaConfigTest.scala    |  29 +--
 .../ControllerChannelManagerTest.scala        |   3 +-
 .../security/authorizer/AuthorizerTest.scala  | 127 +++++++------
 .../server/DynamicBrokerConfigTest.scala      |   6 +-
 .../unit/kafka/server/KafkaConfigTest.scala   |  80 ++++----
 .../unit/kafka/server/KafkaServerTest.scala   |  51 ++---
 .../kafka/server/ServerShutdownTest.scala     |   5 +-
 .../scala/unit/kafka/server/ServerTest.scala  |   3 +-
 .../scala/unit/kafka/utils/TestUtils.scala    |   6 +-
 .../unit/kafka/zk/KafkaZkClientTest.scala     |   4 +-
 .../zk/migration/ZkMigrationTestHarness.scala |   3 +-
 .../kafka/zookeeper/ZooKeeperClientTest.scala |   3 +-
 .../metadata/MetadataRequestBenchmark.java    |   3 +-
 .../apache/kafka/server/config/Defaults.java  |  10 -
 .../apache/kafka/server/config/ZkConfigs.java | 145 ++++++++++++++
 .../utils/EmbeddedKafkaCluster.java           |   3 +-
 .../integration/utils/KafkaEmbedded.java      |   3 +-
 40 files changed, 439 insertions(+), 362 deletions(-)
 create mode 100644 server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java

diff --git a/build.gradle b/build.gradle
index 489473c885aef..0564d27a01380 100644
--- a/build.gradle
+++ b/build.gradle
@@ -2160,6 +2160,7 @@ project(':streams') {
     testImplementation project(':storage')
     testImplementation project(':server-common')
     testImplementation project(':server-common').sourceSets.test.output
+    testImplementation project(':server')
     testImplementation libs.log4j
     testImplementation libs.junitJupiter
     testImplementation libs.junitVintageEngine
@@ -2742,6 +2743,7 @@ project(':jmh-benchmarks') {
       exclude group: 'net.sf.jopt-simple', module: 'jopt-simple'
     }
     implementation project(':server-common')
+    implementation project(':server')
     implementation project(':clients')
     implementation project(':group-coordinator')
     implementation project(':metadata')
@@ -2899,7 +2901,7 @@ project(':connect:json') {
     api libs.jacksonAfterburner
 
     implementation libs.slf4jApi
-    
+
     testImplementation libs.junitJupiter
 
     testRuntimeOnly libs.slf4jlog4j
@@ -2969,6 +2971,7 @@ project(':connect:runtime') {
     testImplementation project(':metadata')
     testImplementation project(':core').sourceSets.test.output
     testImplementation project(':server-common')
+    testImplementation project(':server')
     testImplementation project(':storage')
     testImplementation project(':connect:test-plugins')
 
diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index 8bbf572821210..616e0e519c630 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -605,6 +605,7 @@
         <allow pkg="org.apache.kafka.metadata" />
         <allow pkg="org.eclipse.jetty.client"/>
         <allow class="org.apache.kafka.storage.internals.log.CleanerConfig" />
+        <allow class="org.apache.kafka.server.config.ZkConfigs" />
       </subpackage>
     </subpackage>
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
index c15aa27ae5964..8a06c8e0b0a39 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
@@ -55,6 +55,7 @@
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.metadata.BrokerState;
+import org.apache.kafka.server.config.ZkConfigs;
 import org.apache.kafka.storage.internals.log.CleanerConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -154,7 +155,7 @@ public void start() {
     }
 
     private void doStart() {
-        brokerConfig.put(KafkaConfig.ZkConnectProp(), zKConnectString());
+        brokerConfig.put(ZkConfigs.ZK_CONNECT_PROP, zKConnectString());
 
         putIfAbsent(brokerConfig, KafkaConfig.DeleteTopicEnableProp(), true);
         putIfAbsent(brokerConfig, KafkaConfig.GroupInitialRebalanceDelayMsProp(), 0);
diff --git a/core/src/main/scala/kafka/admin/AclCommand.scala b/core/src/main/scala/kafka/admin/AclCommand.scala
index 2e867d76dd2cf..7120d0d8c3439 100644
--- a/core/src/main/scala/kafka/admin/AclCommand.scala
+++ b/core/src/main/scala/kafka/admin/AclCommand.scala
@@ -21,7 +21,6 @@ import java.util.Properties
 import joptsimple._
 import joptsimple.util.EnumConverter
 import kafka.security.authorizer.{AclAuthorizer, AclEntry}
-import kafka.server.KafkaConfig
 import kafka.utils._
 import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
 import org.apache.kafka.common.acl._
@@ -33,6 +32,7 @@ import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{Utils, SecurityUtils => JSecurityUtils}
 import org.apache.kafka.security.authorizer.AuthorizerUtils
 import org.apache.kafka.server.authorizer.Authorizer
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 
 import scala.jdk.CollectionConverters._
@@ -200,7 +200,7 @@ object AclCommand extends Logging {
       // We will default the value of zookeeper.set.acl to true or false based on whether SASL is configured,
       // but if SASL is not configured and zookeeper.set.acl is supposed to be true due to mutual certificate authentication
       // then it will be up to the user to explicitly specify zookeeper.set.acl=true in the authorizer-properties.
-      val defaultProps = Map(KafkaConfig.ZkEnableSecureAclsProp -> JaasUtils.isZkSaslEnabled)
+      val defaultProps = Map(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP -> JaasUtils.isZkSaslEnabled)
       val authorizerPropertiesWithoutTls =
         if (opts.options.has(opts.authorizerPropertiesOpt)) {
           val authorizerProperties = opts.options.valuesOf(opts.authorizerPropertiesOpt)
@@ -211,7 +211,7 @@ object AclCommand extends Logging {
       val authorizerProperties =
         if (opts.options.has(opts.zkTlsConfigFile)) {
           // load in TLS configs both with and without the "authorizer." prefix
-          val validKeys = (KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList ++ KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.map("authorizer." + _).toList).asJava
+          val validKeys = (ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.toList ++ ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.map("authorizer." + _).toList).asJava
           authorizerPropertiesWithoutTls ++ Utils.loadProps(opts.options.valueOf(opts.zkTlsConfigFile), validKeys).asInstanceOf[java.util.Map[String, Any]].asScala
         }
         else
@@ -619,7 +619,7 @@ object AclCommand extends Logging {
       "DEPRECATED: Identifies the file where ZooKeeper client TLS connectivity properties are defined for" +
         " the default authorizer kafka.security.authorizer.AclAuthorizer." +
         " Any properties other than the following (with or without an \"authorizer.\" prefix) are ignored: " +
-        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.sorted.mkString(", ") +
+        ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.toList.sorted.mkString(", ") +
         ". Note that if SASL is not configured and zookeeper.set.acl is supposed to be true due to mutual certificate authentication being used" +
         " then it is necessary to explicitly specify --authorizer-properties zookeeper.set.acl=true. " +
         AclCommand.AuthorizerDeprecationMessage)
diff --git a/core/src/main/scala/kafka/admin/ConfigCommand.scala b/core/src/main/scala/kafka/admin/ConfigCommand.scala
index d03d7dfb9b52a..a10722f000f11 100644
--- a/core/src/main/scala/kafka/admin/ConfigCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConfigCommand.scala
@@ -35,7 +35,7 @@ import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity,
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter, ScramMechanism}
 import org.apache.kafka.common.utils.{Sanitizer, Time, Utils}
-import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals}
+import org.apache.kafka.server.config.{ConfigType, ZooKeeperInternals, ZkConfigs}
 import org.apache.kafka.security.{PasswordEncoder, PasswordEncoderConfigs}
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 import org.apache.kafka.storage.internals.log.LogConfig
@@ -864,7 +864,7 @@ object ConfigCommand extends Logging {
       .ofType(classOf[String])
     val zkTlsConfigFile: OptionSpec[String] = parser.accepts("zk-tls-config-file",
       "Identifies the file where ZooKeeper client TLS connectivity properties are defined.  Any properties other than " +
-        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.sorted.mkString(", ") + " are ignored.")
+        ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.toList.sorted.mkString(", ") + " are ignored.")
       .withRequiredArg().describedAs("ZooKeeper TLS configuration").ofType(classOf[String])
     options = parser.parse(args : _*)
 
diff --git a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
index 9c81bb23e2e2b..dd07f522e77dc 100644
--- a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
+++ b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
@@ -24,6 +24,7 @@ import kafka.utils.Implicits._
 import kafka.zk.{ControllerZNode, KafkaZkClient, ZkData, ZkSecurityMigratorUtils}
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
 import org.apache.zookeeper.AsyncCallback.{ChildrenCallback, StatCallback}
 import org.apache.zookeeper.KeeperException
@@ -81,7 +82,7 @@ object ZkSecurityMigrator extends Logging {
     if (jaasFile == null && !tlsClientAuthEnabled) {
       val errorMsg = s"No JAAS configuration file has been specified and no TLS client certificate has been specified. Please make sure that you set " +
         s"the system property ${JaasUtils.JAVA_LOGIN_CONFIG_PARAM} or provide a ZooKeeper client TLS configuration via --$tlsConfigFileOption <filename> " +
-        s"identifying at least ${KafkaConfig.ZkSslClientEnableProp}, ${KafkaConfig.ZkClientCnxnSocketProp}, and ${KafkaConfig.ZkSslKeyStoreLocationProp}"
+        s"identifying at least ${ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP}, and ${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP}"
       System.err.println("ERROR: %s".format(errorMsg))
       throw new IllegalArgumentException("Incorrect configuration")
     }
@@ -124,7 +125,7 @@ object ZkSecurityMigrator extends Logging {
   }
 
   def createZkClientConfigFromFile(filename: String) : ZKClientConfig = {
-    val zkTlsConfigFileProps = Utils.loadProps(filename, KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.asJava)
+    val zkTlsConfigFileProps = Utils.loadProps(filename, ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.toList.asJava)
     val zkClientConfig = new ZKClientConfig() // Initializes based on any system properties that have been set
     // Now override any set system properties with explicitly-provided values from the config file
     // Emit INFO logs due to camel-case property names encouraging mistakes -- help people see mistakes they make
@@ -156,7 +157,7 @@ object ZkSecurityMigrator extends Logging {
       "before migration. If not, exit the command.")
     val zkTlsConfigFile: OptionSpec[String] = parser.accepts(tlsConfigFileOption,
       "Identifies the file where ZooKeeper client TLS connectivity properties are defined.  Any properties other than " +
-        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.mkString(", ") + " are ignored.")
+        ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.mkString(", ") + " are ignored.")
       .withRequiredArg().describedAs("ZooKeeper TLS configuration").ofType(classOf[String])
     options = parser.parse(args : _*)
   }
diff --git a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
index 99642b33c2b90..dc307ca702fa8 100644
--- a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -37,6 +37,7 @@ import org.apache.kafka.common.utils.{SecurityUtils, Time}
 import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
 import org.apache.kafka.server.authorizer._
 import org.apache.kafka.server.common.MetadataVersion.IBP_2_0_IV1
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.zookeeper.client.ZKClientConfig
 
 import scala.annotation.nowarn
@@ -96,7 +97,7 @@ object AclAuthorizer {
   }
 
   private[authorizer] def zkClientConfigFromKafkaConfigAndMap(kafkaConfig: KafkaConfig, configMap: mutable.Map[String, _<:Any]): ZKClientConfig = {
-    val zkSslClientEnable = configMap.get(AclAuthorizer.configPrefix + KafkaConfig.ZkSslClientEnableProp).
+    val zkSslClientEnable = configMap.get(AclAuthorizer.configPrefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP).
       map(_.toString.trim).getOrElse(kafkaConfig.zkSslClientEnable.toString).toBoolean
     if (!zkSslClientEnable)
       new ZKClientConfig
@@ -105,10 +106,10 @@ object AclAuthorizer {
       // be sure to force creation since the zkSslClientEnable property in the kafkaConfig could be false
       val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(kafkaConfig, forceZkSslClientEnable = true)
       // add in any prefixed overlays
-      KafkaConfig.ZkSslConfigToSystemPropertyMap.forKeyValue { (kafkaProp, sysProp) =>
+      ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.forKeyValue { (kafkaProp, sysProp) =>
         configMap.get(AclAuthorizer.configPrefix + kafkaProp).foreach { prefixedValue =>
           zkClientConfig.setProperty(sysProp,
-            if (kafkaProp == KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp)
+            if (kafkaProp == ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP)
               (prefixedValue.toString.trim.toUpperCase == "HTTPS").toString
             else
               prefixedValue.toString.trim)
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index 4290ce1da8d45..cec5f5649f628 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -45,7 +45,7 @@ import org.apache.kafka.server.ProcessRole
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.common.{MetadataVersion, MetadataVersionValidator}
 import org.apache.kafka.server.common.MetadataVersion._
-import org.apache.kafka.server.config.{Defaults, ServerTopicConfigSynonyms}
+import org.apache.kafka.server.config.{Defaults, ServerTopicConfigSynonyms, ZkConfigs}
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.record.BrokerCompressionType
 import org.apache.kafka.server.util.Csv
@@ -67,53 +67,15 @@ object KafkaConfig {
       DynamicBrokerConfig.dynamicConfigUpdateModes))
   }
 
-  /** ********* Zookeeper Configuration ***********/
-  val ZkConnectProp = "zookeeper.connect"
-  val ZkSessionTimeoutMsProp = "zookeeper.session.timeout.ms"
-  val ZkConnectionTimeoutMsProp = "zookeeper.connection.timeout.ms"
-  val ZkEnableSecureAclsProp = "zookeeper.set.acl"
-  val ZkMaxInFlightRequestsProp = "zookeeper.max.in.flight.requests"
-  val ZkSslClientEnableProp = "zookeeper.ssl.client.enable"
-  val ZkClientCnxnSocketProp = "zookeeper.clientCnxnSocket"
-  val ZkSslKeyStoreLocationProp = "zookeeper.ssl.keystore.location"
-  val ZkSslKeyStorePasswordProp = "zookeeper.ssl.keystore.password"
-  val ZkSslKeyStoreTypeProp = "zookeeper.ssl.keystore.type"
-  val ZkSslTrustStoreLocationProp = "zookeeper.ssl.truststore.location"
-  val ZkSslTrustStorePasswordProp = "zookeeper.ssl.truststore.password"
-  val ZkSslTrustStoreTypeProp = "zookeeper.ssl.truststore.type"
-  val ZkSslProtocolProp = "zookeeper.ssl.protocol"
-  val ZkSslEnabledProtocolsProp = "zookeeper.ssl.enabled.protocols"
-  val ZkSslCipherSuitesProp = "zookeeper.ssl.cipher.suites"
-  val ZkSslEndpointIdentificationAlgorithmProp = "zookeeper.ssl.endpoint.identification.algorithm"
-  val ZkSslCrlEnableProp = "zookeeper.ssl.crl.enable"
-  val ZkSslOcspEnableProp = "zookeeper.ssl.ocsp.enable"
-
-  // a map from the Kafka config to the corresponding ZooKeeper Java system property
-  private[kafka] val ZkSslConfigToSystemPropertyMap: Map[String, String] = Map(
-    ZkSslClientEnableProp -> ZKClientConfig.SECURE_CLIENT,
-    ZkClientCnxnSocketProp -> ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET,
-    ZkSslKeyStoreLocationProp -> "zookeeper.ssl.keyStore.location",
-    ZkSslKeyStorePasswordProp -> "zookeeper.ssl.keyStore.password",
-    ZkSslKeyStoreTypeProp -> "zookeeper.ssl.keyStore.type",
-    ZkSslTrustStoreLocationProp -> "zookeeper.ssl.trustStore.location",
-    ZkSslTrustStorePasswordProp -> "zookeeper.ssl.trustStore.password",
-    ZkSslTrustStoreTypeProp -> "zookeeper.ssl.trustStore.type",
-    ZkSslProtocolProp -> "zookeeper.ssl.protocol",
-    ZkSslEnabledProtocolsProp -> "zookeeper.ssl.enabledProtocols",
-    ZkSslCipherSuitesProp -> "zookeeper.ssl.ciphersuites",
-    ZkSslEndpointIdentificationAlgorithmProp -> "zookeeper.ssl.hostnameVerification",
-    ZkSslCrlEnableProp -> "zookeeper.ssl.crl",
-    ZkSslOcspEnableProp -> "zookeeper.ssl.ocsp")
-
   private[kafka] def zooKeeperClientProperty(clientConfig: ZKClientConfig, kafkaPropName: String): Option[String] = {
-    Option(clientConfig.getProperty(ZkSslConfigToSystemPropertyMap(kafkaPropName)))
+    Option(clientConfig.getProperty(ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(kafkaPropName)))
   }
 
   private[kafka] def setZooKeeperClientProperty(clientConfig: ZKClientConfig, kafkaPropName: String, kafkaPropValue: Any): Unit = {
-    clientConfig.setProperty(ZkSslConfigToSystemPropertyMap(kafkaPropName),
+    clientConfig.setProperty(ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(kafkaPropName),
       kafkaPropName match {
-        case ZkSslEndpointIdentificationAlgorithmProp => (kafkaPropValue.toString.toUpperCase == "HTTPS").toString
-        case ZkSslEnabledProtocolsProp | ZkSslCipherSuitesProp => kafkaPropValue match {
+        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => (kafkaPropValue.toString.toUpperCase == "HTTPS").toString
+        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => kafkaPropValue match {
           case list: java.util.List[_] => list.asScala.mkString(",")
           case _ => kafkaPropValue.toString
         }
@@ -124,9 +86,9 @@ object KafkaConfig {
   // For ZooKeeper TLS client authentication to be enabled the client must (at a minimum) configure itself as using TLS
   // with both a client connection socket and a key store location explicitly set.
   private[kafka] def zkTlsClientAuthEnabled(zkClientConfig: ZKClientConfig): Boolean = {
-    zooKeeperClientProperty(zkClientConfig, ZkSslClientEnableProp).contains("true") &&
-      zooKeeperClientProperty(zkClientConfig, ZkClientCnxnSocketProp).isDefined &&
-      zooKeeperClientProperty(zkClientConfig, ZkSslKeyStoreLocationProp).isDefined
+    zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP).contains("true") &&
+      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP).isDefined &&
+      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP).isDefined
   }
 
   /** ********* General Configuration ***********/
@@ -439,51 +401,6 @@ object KafkaConfig {
   val UnstableMetadataVersionsEnableProp = "unstable.metadata.versions.enable"
 
   /* Documentation */
-  /** ********* Zookeeper Configuration ***********/
-  val ZkConnectDoc = "Specifies the ZooKeeper connection string in the form <code>hostname:port</code> where host and port are the " +
-  "host and port of a ZooKeeper server. To allow connecting through other ZooKeeper nodes when that ZooKeeper machine is " +
-  "down you can also specify multiple hosts in the form <code>hostname1:port1,hostname2:port2,hostname3:port3</code>.\n" +
-  "The server can also have a ZooKeeper chroot path as part of its ZooKeeper connection string which puts its data under some path in the global ZooKeeper namespace. " +
-  "For example to give a chroot path of <code>/chroot/path</code> you would give the connection string as <code>hostname1:port1,hostname2:port2,hostname3:port3/chroot/path</code>."
-  val ZkSessionTimeoutMsDoc = "Zookeeper session timeout"
-  val ZkConnectionTimeoutMsDoc = "The max time that the client waits to establish a connection to ZooKeeper. If not set, the value in " + ZkSessionTimeoutMsProp + " is used"
-  val ZkEnableSecureAclsDoc = "Set client to use secure ACLs"
-  val ZkMaxInFlightRequestsDoc = "The maximum number of unacknowledged requests the client will send to ZooKeeper before blocking."
-  val ZkSslClientEnableDoc = "Set client to use TLS when connecting to ZooKeeper." +
-    " An explicit value overrides any value set via the <code>zookeeper.client.secure</code> system property (note the different name)." +
-    s" Defaults to false if neither is set; when true, <code>$ZkClientCnxnSocketProp</code> must be set (typically to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code>); other values to set may include " +
-    ZkSslConfigToSystemPropertyMap.keys.toList.filter(x => x != ZkSslClientEnableProp && x != ZkClientCnxnSocketProp).sorted.mkString("<code>", "</code>, <code>", "</code>")
-  val ZkClientCnxnSocketDoc = "Typically set to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code> when using TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the same-named <code>${ZkSslConfigToSystemPropertyMap(ZkClientCnxnSocketProp)}</code> system property."
-  val ZkSslKeyStoreLocationDoc = "Keystore location when using a client-side certificate with TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslKeyStoreLocationProp)}</code> system property (note the camelCase)."
-  val ZkSslKeyStorePasswordDoc = "Keystore password when using a client-side certificate with TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslKeyStorePasswordProp)}</code> system property (note the camelCase)." +
-    " Note that ZooKeeper does not support a key password different from the keystore password, so be sure to set the key password in the keystore to be identical to the keystore password; otherwise the connection attempt to Zookeeper will fail."
-  val ZkSslKeyStoreTypeDoc = "Keystore type when using a client-side certificate with TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslKeyStoreTypeProp)}</code> system property (note the camelCase)." +
-    " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the keystore."
-  val ZkSslTrustStoreLocationDoc = "Truststore location when using TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslTrustStoreLocationProp)}</code> system property (note the camelCase)."
-  val ZkSslTrustStorePasswordDoc = "Truststore password when using TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslTrustStorePasswordProp)}</code> system property (note the camelCase)."
-  val ZkSslTrustStoreTypeDoc = "Truststore type when using TLS connectivity to ZooKeeper." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslTrustStoreTypeProp)}</code> system property (note the camelCase)." +
-    " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the truststore."
-  val ZkSslProtocolDoc = "Specifies the protocol to be used in ZooKeeper TLS negotiation." +
-    s" An explicit value overrides any value set via the same-named <code>${ZkSslConfigToSystemPropertyMap(ZkSslProtocolProp)}</code> system property."
-  val ZkSslEnabledProtocolsDoc = "Specifies the enabled protocol(s) in ZooKeeper TLS negotiation (csv)." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslEnabledProtocolsProp)}</code> system property (note the camelCase)." +
-    s" The default value of <code>null</code> means the enabled protocol will be the value of the <code>${KafkaConfig.ZkSslProtocolProp}</code> configuration property."
-  val ZkSslCipherSuitesDoc = "Specifies the enabled cipher suites to be used in ZooKeeper TLS negotiation (csv)." +
-    s""" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslCipherSuitesProp)}</code> system property (note the single word \"ciphersuites\").""" +
-    " The default value of <code>null</code> means the list of enabled cipher suites is determined by the Java runtime being used."
-  val ZkSslEndpointIdentificationAlgorithmDoc = "Specifies whether to enable hostname verification in the ZooKeeper TLS negotiation process, with (case-insensitively) \"https\" meaning ZooKeeper hostname verification is enabled and an explicit blank value meaning it is disabled (disabling it is only recommended for testing purposes)." +
-    s""" An explicit value overrides any \"true\" or \"false\" value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslEndpointIdentificationAlgorithmProp)}</code> system property (note the different name and values; true implies https and false implies blank)."""
-  val ZkSslCrlEnableDoc = "Specifies whether to enable Certificate Revocation List in the ZooKeeper TLS protocols." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslCrlEnableProp)}</code> system property (note the shorter name)."
-  val ZkSslOcspEnableDoc = "Specifies whether to enable Online Certificate Status Protocol in the ZooKeeper TLS protocols." +
-    s" Overrides any explicit value set via the <code>${ZkSslConfigToSystemPropertyMap(ZkSslOcspEnableProp)}</code> system property (note the shorter name)."
   /** ********* General Configuration ***********/
   val BrokerIdGenerationEnableDoc = s"Enable automatic broker id generation on the server. When enabled the value configured for $MaxReservedBrokerIdProp should be reviewed."
   val MaxReservedBrokerIdDoc = "Max number that can be used for a broker.id"
@@ -943,25 +860,25 @@ object KafkaConfig {
     new ConfigDef()
 
       /** ********* Zookeeper Configuration ***********/
-      .define(ZkConnectProp, STRING, null, HIGH, ZkConnectDoc)
-      .define(ZkSessionTimeoutMsProp, INT, Defaults.ZK_SESSION_TIMEOUT_MS, HIGH, ZkSessionTimeoutMsDoc)
-      .define(ZkConnectionTimeoutMsProp, INT, null, HIGH, ZkConnectionTimeoutMsDoc)
-      .define(ZkEnableSecureAclsProp, BOOLEAN, Defaults.ZK_ENABLE_SECURE_ACLS, HIGH, ZkEnableSecureAclsDoc)
-      .define(ZkMaxInFlightRequestsProp, INT, Defaults.ZK_MAX_IN_FLIGHT_REQUESTS, atLeast(1), HIGH, ZkMaxInFlightRequestsDoc)
-      .define(ZkSslClientEnableProp, BOOLEAN, Defaults.ZK_SSL_CLIENT_ENABLE, MEDIUM, ZkSslClientEnableDoc)
-      .define(ZkClientCnxnSocketProp, STRING, null, MEDIUM, ZkClientCnxnSocketDoc)
-      .define(ZkSslKeyStoreLocationProp, STRING, null, MEDIUM, ZkSslKeyStoreLocationDoc)
-      .define(ZkSslKeyStorePasswordProp, PASSWORD, null, MEDIUM, ZkSslKeyStorePasswordDoc)
-      .define(ZkSslKeyStoreTypeProp, STRING, null, MEDIUM, ZkSslKeyStoreTypeDoc)
-      .define(ZkSslTrustStoreLocationProp, STRING, null, MEDIUM, ZkSslTrustStoreLocationDoc)
-      .define(ZkSslTrustStorePasswordProp, PASSWORD, null, MEDIUM, ZkSslTrustStorePasswordDoc)
-      .define(ZkSslTrustStoreTypeProp, STRING, null, MEDIUM, ZkSslTrustStoreTypeDoc)
-      .define(ZkSslProtocolProp, STRING, Defaults.ZK_SSL_PROTOCOL, LOW, ZkSslProtocolDoc)
-      .define(ZkSslEnabledProtocolsProp, LIST, null, LOW, ZkSslEnabledProtocolsDoc)
-      .define(ZkSslCipherSuitesProp, LIST, null, LOW, ZkSslCipherSuitesDoc)
-      .define(ZkSslEndpointIdentificationAlgorithmProp, STRING, Defaults.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM, LOW, ZkSslEndpointIdentificationAlgorithmDoc)
-      .define(ZkSslCrlEnableProp, BOOLEAN, Defaults.ZK_SSL_CRL_ENABLE, LOW, ZkSslCrlEnableDoc)
-      .define(ZkSslOcspEnableProp, BOOLEAN, Defaults.ZK_SSL_OCSP_ENABLE, LOW, ZkSslOcspEnableDoc)
+      .define(ZkConfigs.ZK_CONNECT_PROP, STRING, null, HIGH, ZkConfigs.ZK_CONNECT_DOC)
+      .define(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, INT, ZkConfigs.ZK_SESSION_TIMEOUT_MS, HIGH, ZkConfigs.ZK_SESSION_TIMEOUT_MS_DOC)
+      .define(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, INT, null, HIGH, ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_DOC)
+      .define(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, BOOLEAN, ZkConfigs.ZK_ENABLE_SECURE_ACLS, HIGH, ZkConfigs.ZK_ENABLE_SECURE_ACLS_DOC)
+      .define(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP, INT, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS, atLeast(1), HIGH, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_DOC)
+      .define(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_CLIENT_ENABLE, MEDIUM, ZkConfigs.ZK_SSL_CLIENT_ENABLE_DOC)
+      .define(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_DOC)
+      .define(ZkConfigs.ZK_SSL_PROTOCOL_PROP, STRING, ZkConfigs.ZK_SSL_PROTOCOL, LOW, ZkConfigs.ZK_SSL_PROTOCOL_DOC)
+      .define(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, LIST, null, LOW, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_DOC)
+      .define(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, LIST, null, LOW, ZkConfigs.ZK_SSL_CIPHER_SUITES_DOC)
+      .define(ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, STRING, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM, LOW, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC)
+      .define(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_CRL_ENABLE, LOW, ZkConfigs.ZK_SSL_CRL_ENABLE_DOC)
+      .define(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_OCSP_ENABLE, LOW, ZkConfigs.ZK_SSL_OCSP_ENABLE_DOC)
 
       /** ********* General Configuration ***********/
       .define(BrokerIdGenerationEnableProp, BOOLEAN, Defaults.BROKER_ID_GENERATION_ENABLE, MEDIUM, BrokerIdGenerationEnableDoc)
@@ -1417,12 +1334,12 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     super.valuesWithPrefixOverride(prefix)
 
   /** ********* Zookeeper Configuration ***********/
-  val zkConnect: String = getString(KafkaConfig.ZkConnectProp)
-  val zkSessionTimeoutMs: Int = getInt(KafkaConfig.ZkSessionTimeoutMsProp)
+  val zkConnect: String = getString(ZkConfigs.ZK_CONNECT_PROP)
+  val zkSessionTimeoutMs: Int = getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP)
   val zkConnectionTimeoutMs: Int =
-    Option(getInt(KafkaConfig.ZkConnectionTimeoutMsProp)).map(_.toInt).getOrElse(getInt(KafkaConfig.ZkSessionTimeoutMsProp))
-  val zkEnableSecureAcls: Boolean = getBoolean(KafkaConfig.ZkEnableSecureAclsProp)
-  val zkMaxInFlightRequests: Int = getInt(KafkaConfig.ZkMaxInFlightRequestsProp)
+    Option(getInt(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP)).map(_.toInt).getOrElse(getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP))
+  val zkEnableSecureAcls: Boolean = getBoolean(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP)
+  val zkMaxInFlightRequests: Int = getInt(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP)
 
   private val _remoteLogManagerConfig = new RemoteLogManagerConfig(this)
   def remoteLogManagerConfig = _remoteLogManagerConfig
@@ -1470,21 +1387,21 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     }
   }
 
-  val zkSslClientEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslClientEnableProp)
-  val zkClientCnxnSocketClassName = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkClientCnxnSocketProp)
-  val zkSslKeyStoreLocation = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslKeyStoreLocationProp)
-  val zkSslKeyStorePassword = zkPasswordConfigOrSystemProperty(KafkaConfig.ZkSslKeyStorePasswordProp)
-  val zkSslKeyStoreType = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslKeyStoreTypeProp)
-  val zkSslTrustStoreLocation = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslTrustStoreLocationProp)
-  val zkSslTrustStorePassword = zkPasswordConfigOrSystemProperty(KafkaConfig.ZkSslTrustStorePasswordProp)
-  val zkSslTrustStoreType = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslTrustStoreTypeProp)
-  val ZkSslProtocol = zkStringConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslProtocolProp)
-  val ZkSslEnabledProtocols = zkListConfigOrSystemProperty(KafkaConfig.ZkSslEnabledProtocolsProp)
-  val ZkSslCipherSuites = zkListConfigOrSystemProperty(KafkaConfig.ZkSslCipherSuitesProp)
+  val zkSslClientEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP)
+  val zkClientCnxnSocketClassName = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP)
+  val zkSslKeyStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP)
+  val zkSslKeyStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP)
+  val zkSslKeyStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP)
+  val zkSslTrustStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP)
+  val zkSslTrustStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP)
+  val zkSslTrustStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP)
+  val ZkSslProtocol = zkStringConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_PROTOCOL_PROP)
+  val ZkSslEnabledProtocols = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP)
+  val ZkSslCipherSuites = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP)
   val ZkSslEndpointIdentificationAlgorithm = {
     // Use the system property if it exists and the Kafka config value was defaulted rather than actually provided
     // Need to translate any system property value from true/false to HTTPS/<blank>
-    val kafkaProp = KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp
+    val kafkaProp = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP
     val actuallyProvided = originals.containsKey(kafkaProp)
     if (actuallyProvided)
       getString(kafkaProp)
@@ -1496,8 +1413,8 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
       }
     }
   }
-  val ZkSslCrlEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslCrlEnableProp)
-  val ZkSslOcspEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslOcspEnableProp)
+  val ZkSslCrlEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP)
+  val ZkSslOcspEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP)
   /** ********* General Configuration ***********/
   val brokerIdGenerationEnable: Boolean = getBoolean(KafkaConfig.BrokerIdGenerationEnableProp)
   val maxReservedBrokerId: Int = getInt(KafkaConfig.MaxReservedBrokerIdProp)
@@ -2040,7 +1957,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     }
     if (requiresZookeeper) {
       if (zkConnect == null) {
-        throw new ConfigException(s"Missing required configuration `${KafkaConfig.ZkConnectProp}` which has no default value.")
+        throw new ConfigException(s"Missing required configuration `${ZkConfigs.ZK_CONNECT_PROP}` which has no default value.")
       }
       if (brokerIdGenerationEnable) {
         require(brokerId >= -1 && brokerId <= maxReservedBrokerId, "broker.id must be greater than or equal to -1 and not greater than reserved.broker.max.id")
@@ -2055,7 +1972,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
       }
       if (migrationEnabled) {
         if (zkConnect == null) {
-          throw new ConfigException(s"If using `${KafkaConfig.MigrationEnabledProp}` in KRaft mode, `${KafkaConfig.ZkConnectProp}` must also be set.")
+          throw new ConfigException(s"If using `${KafkaConfig.MigrationEnabledProp}` in KRaft mode, `${ZkConfigs.ZK_CONNECT_PROP}` must also be set.")
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index 55c773b0a46db..fce61e17a225a 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -57,7 +57,7 @@ import org.apache.kafka.server.NodeToControllerChannelManager
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.common.MetadataVersion._
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion}
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.server.fault.LoggingFaultHandler
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
@@ -80,20 +80,20 @@ object KafkaServer {
   def zkClientConfigFromKafkaConfig(config: KafkaConfig, forceZkSslClientEnable: Boolean = false): ZKClientConfig = {
     val clientConfig = new ZKClientConfig
     if (config.zkSslClientEnable || forceZkSslClientEnable) {
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslClientEnableProp, "true")
-      config.zkClientCnxnSocketClassName.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkClientCnxnSocketProp, _))
-      config.zkSslKeyStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStoreLocationProp, _))
-      config.zkSslKeyStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStorePasswordProp, x.value))
-      config.zkSslKeyStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStoreTypeProp, _))
-      config.zkSslTrustStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStoreLocationProp, _))
-      config.zkSslTrustStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStorePasswordProp, x.value))
-      config.zkSslTrustStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStoreTypeProp, _))
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslProtocolProp, config.ZkSslProtocol)
-      config.ZkSslEnabledProtocols.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslEnabledProtocolsProp, _))
-      config.ZkSslCipherSuites.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslCipherSuitesProp, _))
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp, config.ZkSslEndpointIdentificationAlgorithm)
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslCrlEnableProp, config.ZkSslCrlEnable.toString)
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslOcspEnableProp, config.ZkSslOcspEnable.toString)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "true")
+      config.zkClientCnxnSocketClassName.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP, _))
+      config.zkSslKeyStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, _))
+      config.zkSslKeyStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, x.value))
+      config.zkSslKeyStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, _))
+      config.zkSslTrustStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, _))
+      config.zkSslTrustStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, x.value))
+      config.zkSslTrustStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_PROTOCOL_PROP, config.ZkSslProtocol)
+      config.ZkSslEnabledProtocols.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, _))
+      config.ZkSslCipherSuites.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, config.ZkSslEndpointIdentificationAlgorithm)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, config.ZkSslCrlEnable.toString)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, config.ZkSslOcspEnable.toString)
     }
     // The zk sasl is enabled by default so it can produce false error when broker does not intend to use SASL.
     if (!JaasUtils.isZkSaslEnabled) clientConfig.setProperty(JaasUtils.ZK_SASL_CLIENT, "false")
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index 26c9453bef211..749de42853320 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -34,7 +34,7 @@ import org.apache.kafka.common.security.token.delegation.{DelegationToken, Token
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{KafkaException, TopicPartition, Uuid}
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.server.metrics.KafkaMetricsGroup
 import org.apache.kafka.storage.internals.log.LogConfig
 import org.apache.zookeeper.KeeperException.{Code, NodeExistsException}
@@ -2349,9 +2349,9 @@ object KafkaZkClient {
 
     if (secureAclsEnabled && !isZkSecurityEnabled)
       throw new java.lang.SecurityException(
-        s"${KafkaConfig.ZkEnableSecureAclsProp} is true, but ZooKeeper client TLS configuration identifying at least " +
-          s"${KafkaConfig.ZkSslClientEnableProp}, ${KafkaConfig.ZkClientCnxnSocketProp}, and " +
-          s"${KafkaConfig.ZkSslKeyStoreLocationProp} was not present and the verification of the JAAS login file failed " +
+        s"${ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP} is true, but ZooKeeper client TLS configuration identifying at least " +
+          s"${ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP}, and " +
+          s"${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP} was not present and the verification of the JAAS login file failed " +
           s"${JaasUtils.zkSecuritySysConfigString}")
 
     KafkaZkClient(config.zkConnect, secureAclsEnabled, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs,
diff --git a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
index 6ea0c2eaec040..77d0784b1a5ec 100644
--- a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
@@ -26,6 +26,7 @@ import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern,
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.server.authorizer.Authorizer
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNull}
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
 
@@ -75,7 +76,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
   import DescribeAuthorizedOperationsTest._
 
   override val brokerCount = 1
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
   this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
 
   var client: Admin = _
diff --git a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
index 46e674c00aa52..061ba3de9eaea 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
@@ -18,9 +18,9 @@
 package kafka.api
 
 import com.yammer.metrics.core.Gauge
+
 import java.util.{Collections, Properties}
 import java.util.concurrent.ExecutionException
-
 import kafka.security.authorizer.AclAuthorizer
 import kafka.security.authorizer.AclEntry.WildcardHost
 import org.apache.kafka.metadata.authorizer.StandardAuthorizer
@@ -38,11 +38,12 @@ import org.apache.kafka.common.resource._
 import org.apache.kafka.common.resource.ResourceType._
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.security.auth._
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo, Timeout}
 import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.{ValueSource, CsvSource}
+import org.junit.jupiter.params.provider.{CsvSource, ValueSource}
 
 import scala.jdk.CollectionConverters._
 
@@ -156,7 +157,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     } else {
       // The next two configuration parameters enable ZooKeeper secure ACLs
       // and sets the Kafka authorizer, both necessary to enable security.
-      this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+      this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
       this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, authorizerClass.getName)
 
       // Set the specific principal that can update ACLs.
diff --git a/core/src/test/scala/integration/kafka/api/MetricsTest.scala b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
index cb4b5b969794a..658b57cddd31f 100644
--- a/core/src/test/scala/integration/kafka/api/MetricsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
@@ -24,6 +24,7 @@ import org.apache.kafka.common.errors.{InvalidTopicException, UnknownTopicOrPart
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.security.authenticator.TestJaasConfig
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.log.remote.storage.{NoOpRemoteLogMetadataManager, NoOpRemoteStorageManager, RemoteLogManagerConfig, RemoteStorageMetrics}
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo}
@@ -43,7 +44,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   private val kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
   private val kafkaServerJaasEntryName =
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "false")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "false")
   this.serverConfig.setProperty(KafkaConfig.AutoCreateTopicsEnableProp, "false")
   this.serverConfig.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, "2.8")
   this.producerConfig.setProperty(ProducerConfig.LINGER_MS_CONFIG, "10")
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
index e084454f5ff7b..fe32106877aa3 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -44,7 +44,7 @@ import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceT
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{ConsumerGroupState, ElectionType, TopicCollection, TopicPartition, TopicPartitionInfo, TopicPartitionReplica, Uuid}
 import org.apache.kafka.controller.ControllerRequestContextUtil.ANONYMOUS_CONTEXT
-import org.apache.kafka.server.config.Defaults
+import org.apache.kafka.server.config.{Defaults, ZkConfigs}
 import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig}
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Disabled, TestInfo}
@@ -2707,7 +2707,7 @@ object PlaintextAdminIntegrationTest {
     var topicConfigEntries2 = Seq(new ConfigEntry(TopicConfig.COMPRESSION_TYPE_CONFIG, "snappy")).asJava
 
     val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, test.brokers.head.config.brokerId.toString)
-    val brokerConfigEntries = Seq(new ConfigEntry(KafkaConfig.ZkConnectProp, "localhost:2181")).asJava
+    val brokerConfigEntries = Seq(new ConfigEntry(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")).asJava
 
     // Alter configs: first and third are invalid, second is valid
     var alterResult = admin.alterConfigs(Map(
diff --git a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
index 6fc3fb9b8a1ab..4cf7a320ff290 100644
--- a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
@@ -12,10 +12,10 @@
   */
 package kafka.api
 
-import kafka.server.KafkaConfig
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo, Timeout}
 import kafka.utils.{JaasTestUtils, TestInfoUtils, TestUtils}
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 
@@ -25,7 +25,7 @@ import scala.jdk.CollectionConverters._
 class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaClientSaslMechanism = "PLAIN"
   private val kafkaServerSaslMechanisms = List("GSSAPI", "PLAIN")
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
diff --git a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
index 6a03d51aa3680..f039e16faa31e 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
@@ -13,10 +13,10 @@
 package kafka.api
 
 import java.util.Locale
-import kafka.server.KafkaConfig
 import kafka.utils.{JaasTestUtils, TestUtils}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo, Timeout}
 
 @Timeout(600)
@@ -26,7 +26,7 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
   private val kafkaServerJaasEntryName =
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "false")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "false")
   // disable secure acls of zkClient in QuorumTestHarness
   override protected def zkAclsEnabled = Some(false)
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
index 29b51713257ff..1d2f4eb8bbf4c 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
@@ -30,6 +30,7 @@ import org.apache.kafka.common.resource.ResourceType.{GROUP, TOPIC}
 import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourcePatternFilter, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.server.authorizer.Authorizer
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.storage.internals.log.LogConfig
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
@@ -46,7 +47,7 @@ class SaslSslAdminIntegrationTest extends BaseAdminIntegrationTest with SaslSetu
 
   val authorizationAdmin = new AclAuthorizationAdmin(classOf[AclAuthorizer], classOf[AclAuthorizer])
 
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
 
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
index 0ea458e25aa0a..19b6e9e46148f 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
@@ -12,14 +12,14 @@
   */
 package kafka.api
 
-import kafka.server.KafkaConfig
 import kafka.utils.{JaasTestUtils, TestUtils}
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo, Timeout}
 
 @Timeout(600)
 class SaslSslConsumerTest extends BaseConsumerTest with SaslSetup {
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
 
diff --git a/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
index 4e6ae3f81320e..ff6ea03a67745 100644
--- a/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
@@ -14,10 +14,8 @@ package kafka.api
 
 import java.util
 import java.util.concurrent._
-
 import com.yammer.metrics.core.Gauge
 import kafka.security.authorizer.AclAuthorizer
-import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, CreateAclsResult}
 import org.apache.kafka.common.acl._
@@ -25,6 +23,7 @@ import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.server.authorizer._
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotNull, assertTrue}
 import org.junit.jupiter.api.{AfterEach, Test}
@@ -80,7 +79,7 @@ object SslAdminIntegrationTest {
 class SslAdminIntegrationTest extends SaslSslAdminIntegrationTest {
   override val authorizationAdmin = new AclAuthorizationAdmin(classOf[SslAdminIntegrationTest.TestableAclAuthorizer], classOf[AclAuthorizer])
 
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
 
   override protected def securityProtocol = SecurityProtocol.SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
index c984eae0272c2..dd85bd41cb002 100644
--- a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -60,7 +60,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.security.scram.ScramCredential
 import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializer}
 import org.apache.kafka.security.PasswordEncoder
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.kafka.server.record.BrokerCompressionType
 import org.apache.kafka.server.util.ShutdownableThread
@@ -122,7 +122,7 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
         properties
       } else {
         val properties = TestUtils.createBrokerConfig(brokerId, zkConnect)
-        properties.put(KafkaConfig.ZkEnableSecureAclsProp, "true")
+        properties.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
         properties
       }
       props ++= securityProps(sslProperties1, TRUSTSTORE_PROPS)
diff --git a/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala b/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
index d4a1e8f1dd1c2..7bf86259230fd 100644
--- a/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
@@ -25,6 +25,7 @@ import kafka.testkit.{KafkaClusterTestKit, TestKitNodes}
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.raft.RaftConfig
 import org.apache.kafka.server.common.MetadataVersion
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions.{assertThrows, fail}
 import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.{Tag, Timeout}
@@ -61,7 +62,7 @@ class KafkaServerKRaftRegistrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -99,7 +100,7 @@ class KafkaServerKRaftRegistrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
index 4d7ddc15cb356..16e1a2902781f 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
@@ -30,6 +30,7 @@ import org.apache.kafka.common.config.SslConfigs
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.network.{ListenerName, Mode}
 import org.apache.kafka.coordinator.group.OffsetConfig
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
 
@@ -79,7 +80,7 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends QuorumT
       props.put(KafkaConfig.ListenerSecurityProtocolMapProp, s"$Internal:PLAINTEXT, $SecureInternal:SASL_SSL," +
         s"$External:PLAINTEXT, $SecureExternal:SASL_SSL")
       props.put(KafkaConfig.InterBrokerListenerNameProp, Internal)
-      props.put(KafkaConfig.ZkEnableSecureAclsProp, "true")
+      props.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
       props.put(KafkaConfig.SaslMechanismInterBrokerProtocolProp, kafkaClientSaslMechanism)
       props.put(s"${new ListenerName(SecureInternal).configPrefix}${KafkaConfig.SaslEnabledMechanismsProp}",
         kafkaServerSaslMechanisms(SecureInternal).mkString(","))
diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
index a3e5ecc0814d9..ff8d2b797f384 100644
--- a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -48,7 +48,7 @@ import org.apache.kafka.raft.RaftConfig
 import org.apache.kafka.security.PasswordEncoder
 import org.apache.kafka.server.ControllerRequestCompletionHandler
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotEquals, assertNotNull, assertTrue, fail}
 import org.junit.jupiter.api.{Assumptions, Timeout}
 import org.junit.jupiter.api.extension.ExtendWith
@@ -177,7 +177,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -309,7 +309,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -443,7 +443,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -508,7 +508,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -576,7 +576,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -636,7 +636,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -711,7 +711,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
diff --git a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
index 4cec0d05e5b78..3d97e7df17157 100644
--- a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
@@ -25,6 +25,7 @@ import kafka.utils.TestUtils.assertBadConfigContainingMessage
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
 import org.apache.kafka.common.config.types.Password
 import org.apache.kafka.common.internals.FatalExitError
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test}
 import org.junit.jupiter.api.Assertions._
 
@@ -155,7 +156,7 @@ class KafkaTest {
       "Missing required configuration `zookeeper.connect` which has no default value.")
 
     // Ensure that no exception is thrown once zookeeper.connect is defined (and we clear controller.listener.names)
-    propertiesFile.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    propertiesFile.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     propertiesFile.setProperty(KafkaConfig.ControllerListenerNamesProp, "")
     KafkaConfig.fromProps(propertiesFile)
   }
@@ -232,61 +233,61 @@ class KafkaTest {
 
   @Test
   def testZkSslClientEnable(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslClientEnableProp, "zookeeper.ssl.client.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "zookeeper.ssl.client.enable",
       "zookeeper.client.secure", booleanPropValueToSet, config => Some(config.zkSslClientEnable), booleanPropValueToSet, Some(false))
   }
 
   @Test
   def testZkSslKeyStoreLocation(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslKeyStoreLocationProp, "zookeeper.ssl.keystore.location",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, "zookeeper.ssl.keystore.location",
       "zookeeper.ssl.keyStore.location", stringPropValueToSet, config => config.zkSslKeyStoreLocation, stringPropValueToSet)
   }
 
   @Test
   def testZkSslTrustStoreLocation(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslTrustStoreLocationProp, "zookeeper.ssl.truststore.location",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, "zookeeper.ssl.truststore.location",
       "zookeeper.ssl.trustStore.location", stringPropValueToSet, config => config.zkSslTrustStoreLocation, stringPropValueToSet)
   }
 
   @Test
   def testZookeeperKeyStorePassword(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslKeyStorePasswordProp, "zookeeper.ssl.keystore.password",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, "zookeeper.ssl.keystore.password",
       "zookeeper.ssl.keyStore.password", passwordPropValueToSet, config => config.zkSslKeyStorePassword, new Password(passwordPropValueToSet))
   }
 
   @Test
   def testZookeeperTrustStorePassword(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslTrustStorePasswordProp, "zookeeper.ssl.truststore.password",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, "zookeeper.ssl.truststore.password",
       "zookeeper.ssl.trustStore.password", passwordPropValueToSet, config => config.zkSslTrustStorePassword, new Password(passwordPropValueToSet))
   }
 
   @Test
   def testZkSslKeyStoreType(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslKeyStoreTypeProp, "zookeeper.ssl.keystore.type",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, "zookeeper.ssl.keystore.type",
       "zookeeper.ssl.keyStore.type", stringPropValueToSet, config => config.zkSslKeyStoreType, stringPropValueToSet)
   }
 
   @Test
   def testZkSslTrustStoreType(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslTrustStoreTypeProp, "zookeeper.ssl.truststore.type",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, "zookeeper.ssl.truststore.type",
       "zookeeper.ssl.trustStore.type", stringPropValueToSet, config => config.zkSslTrustStoreType, stringPropValueToSet)
   }
 
   @Test
   def testZkSslProtocol(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslProtocolProp, "zookeeper.ssl.protocol",
+    testZkConfig(ZkConfigs.ZK_SSL_PROTOCOL_PROP, "zookeeper.ssl.protocol",
       "zookeeper.ssl.protocol", stringPropValueToSet, config => Some(config.ZkSslProtocol), stringPropValueToSet, Some("TLSv1.2"))
   }
 
   @Test
   def testZkSslEnabledProtocols(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslEnabledProtocolsProp, "zookeeper.ssl.enabled.protocols",
+    testZkConfig(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, "zookeeper.ssl.enabled.protocols",
       "zookeeper.ssl.enabledProtocols", listPropValueToSet.mkString(","), config => config.ZkSslEnabledProtocols, listPropValueToSet.asJava)
   }
 
   @Test
   def testZkSslCipherSuites(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslCipherSuitesProp, "zookeeper.ssl.cipher.suites",
+    testZkConfig(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, "zookeeper.ssl.cipher.suites",
       "zookeeper.ssl.ciphersuites", listPropValueToSet.mkString(","), config => config.ZkSslCipherSuites, listPropValueToSet.asJava)
   }
 
@@ -294,7 +295,7 @@ class KafkaTest {
   def testZkSslEndpointIdentificationAlgorithm(): Unit = {
     // this property is different than the others
     // because the system property values and the Kafka property values don't match
-    val kafkaPropName = KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp
+    val kafkaPropName = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP
     assertEquals("zookeeper.ssl.endpoint.identification.algorithm", kafkaPropName)
     val sysProp = "zookeeper.ssl.hostnameVerification"
     val expectedDefaultValue = "HTTPS"
@@ -327,13 +328,13 @@ class KafkaTest {
 
   @Test
   def testZkSslCrlEnable(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslCrlEnableProp, "zookeeper.ssl.crl.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, "zookeeper.ssl.crl.enable",
       "zookeeper.ssl.crl", booleanPropValueToSet, config => Some(config.ZkSslCrlEnable), booleanPropValueToSet, Some(false))
   }
 
   @Test
   def testZkSslOcspEnable(): Unit = {
-    testZkConfig(KafkaConfig.ZkSslOcspEnableProp, "zookeeper.ssl.ocsp.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, "zookeeper.ssl.ocsp.enable",
       "zookeeper.ssl.ocsp", booleanPropValueToSet, config => Some(config.ZkSslOcspEnable), booleanPropValueToSet, Some(false))
   }
 
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
index d6718e4203cf1..2730431565e7f 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
@@ -34,6 +34,7 @@ import org.apache.kafka.common.{TopicPartition, Uuid}
 import org.apache.kafka.metadata.LeaderRecoveryState
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.common.MetadataVersion.{IBP_0_10_0_IV1, IBP_0_10_2_IV0, IBP_0_9_0, IBP_1_0_IV0, IBP_2_2_IV0, IBP_2_4_IV0, IBP_2_4_IV1, IBP_2_6_IV0, IBP_2_8_IV1, IBP_3_2_IV0, IBP_3_4_IV0}
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
 
@@ -895,7 +896,7 @@ class ControllerChannelManagerTest {
   private def createConfig(interBrokerVersion: MetadataVersion): KafkaConfig = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, controllerId.toString)
-    props.put(KafkaConfig.ZkConnectProp, "zkConnect")
+    props.put(ZkConfigs.ZK_CONNECT_PROP, "zkConnect")
     TestUtils.setIbpAndMessageFormatVersions(props, interBrokerVersion)
     KafkaConfig.fromProps(props)
   }
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
index 80790ab905fdb..29907f220685a 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
@@ -41,6 +41,7 @@ import org.apache.kafka.metadata.authorizer.StandardAuthorizer
 import org.apache.kafka.server.authorizer._
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.common.MetadataVersion.{IBP_2_0_IV0, IBP_2_0_IV1}
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.zookeeper.client.ZKClientConfig
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
@@ -869,7 +870,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(noTlsProps),
       noTlsProps.asInstanceOf[java.util.Map[String, Any]].asScala)
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach { propName =>
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach { propName =>
       assertNull(zkClientConfig.getProperty(propName))
     }
   }
@@ -879,27 +880,27 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val props = new java.util.Properties()
     val kafkaValue = "kafkaValue"
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      KafkaConfig.ZkSslClientEnableProp -> "true",
-      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
-      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue)
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue)
     configs.foreach { case (key, value) => props.put(key, value) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
-      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
         assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
         assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case KafkaConfig.ZkSslProtocolProp =>
+      case ZkConfigs.ZK_SSL_PROTOCOL_PROP =>
         assertEquals("TLSv1.2", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       case _ => assertEquals(kafkaValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
     })
@@ -910,29 +911,29 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val props = new java.util.Properties()
     val kafkaValue = "kafkaValue"
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      KafkaConfig.ZkSslClientEnableProp -> "true",
-      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslProtocolProp -> kafkaValue,
-      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
-      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue,
-      KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "HTTPS",
-      KafkaConfig.ZkSslCrlEnableProp -> "false",
-      KafkaConfig.ZkSslOcspEnableProp -> "false")
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_PROTOCOL_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "HTTPS",
+      ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "false",
+      ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "false")
     configs.foreach{case (key, value) => props.put(key, value.toString) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
-        case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
+        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
           assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-        case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+        case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
           assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
         case _ => assertEquals(kafkaValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       })
@@ -945,43 +946,43 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val prefixedValue = "prefixedValue"
     val prefix = "authorizer."
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      KafkaConfig.ZkSslClientEnableProp -> "false",
-      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
-      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
-      KafkaConfig.ZkSslProtocolProp -> kafkaValue,
-      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
-      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue,
-      KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "HTTPS",
-      KafkaConfig.ZkSslCrlEnableProp -> "false",
-      KafkaConfig.ZkSslOcspEnableProp -> "false",
-      prefix + KafkaConfig.ZkSslClientEnableProp -> "true",
-      prefix + KafkaConfig.ZkClientCnxnSocketProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslKeyStoreLocationProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslKeyStorePasswordProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslKeyStoreTypeProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslTrustStoreLocationProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslTrustStorePasswordProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslTrustStoreTypeProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslProtocolProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslEnabledProtocolsProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslCipherSuitesProp -> prefixedValue,
-      prefix + KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "",
-      prefix + KafkaConfig.ZkSslCrlEnableProp -> "true",
-      prefix + KafkaConfig.ZkSslOcspEnableProp -> "true")
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "false",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_PROTOCOL_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "HTTPS",
+      ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "false",
+      ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "false",
+      prefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
+      prefix + ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_PROTOCOL_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "",
+      prefix + ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "true",
+      prefix + ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "true")
     configs.foreach{case (key, value) => props.put(key, value.toString) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
-      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
         assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
         assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       case _ => assertEquals(prefixedValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
     })
diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
index 4e1c8eed2769d..8e7fe06ef2eaf 100755
--- a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
@@ -34,7 +34,7 @@ import org.apache.kafka.common.config.{ConfigException, SslConfigs}
 import org.apache.kafka.common.metrics.{JmxReporter, Metrics}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.server.authorizer._
-import org.apache.kafka.server.config.Defaults
+import org.apache.kafka.server.config.{Defaults, ZkConfigs}
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.kafka.server.util.KafkaScheduler
@@ -211,7 +211,7 @@ class DynamicBrokerConfigTest {
 
     val securityPropsWithoutListenerPrefix = Map(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG -> "PKCS12")
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, securityPropsWithoutListenerPrefix)
-    val nonDynamicProps = Map(KafkaConfig.ZkConnectProp -> "somehost:2181")
+    val nonDynamicProps = Map(ZkConfigs.ZK_CONNECT_PROP -> "somehost:2181")
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, nonDynamicProps)
 
     // Test update of configs with invalid type
@@ -709,7 +709,7 @@ class DynamicBrokerConfigTest {
   @Test
   def testNonInternalValuesDoesNotExposeInternalConfigs(): Unit = {
     val props = new Properties()
-    props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     props.put(KafkaConfig.MetadataLogSegmentMinBytesProp, "1024")
     val config = new KafkaConfig(props)
     assertFalse(config.nonInternalValues.containsKey(KafkaConfig.MetadataLogSegmentMinBytesProp))
diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
index a5d4d961fe123..fd4377f658e1e 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
@@ -38,7 +38,7 @@ import org.apache.kafka.common.Node
 import org.apache.kafka.coordinator.group.Group.GroupType
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.common.MetadataVersion.{IBP_0_8_2, IBP_3_0_IV1}
-import org.apache.kafka.server.config.ServerTopicConfigSynonyms
+import org.apache.kafka.server.config.{ServerTopicConfigSynonyms, ZkConfigs}
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig}
 import org.junit.jupiter.api.function.Executable
@@ -155,7 +155,7 @@ class KafkaConfigTest {
     val hostName = "fake-host"
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, s"PLAINTEXT://$hostName:$port")
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -186,7 +186,7 @@ class KafkaConfigTest {
   def testDuplicateListeners(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     // listeners with duplicate port
     props.setProperty(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9091,SSL://localhost:9091")
@@ -212,7 +212,7 @@ class KafkaConfigTest {
   def testIPv4AndIPv6SamePortListeners(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
-    props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.put(KafkaConfig.ListenersProp, "PLAINTEXT://[::1]:9092,SSL://[::1]:9092")
     var caught = assertThrows(classOf[IllegalArgumentException], () => KafkaConfig.fromProps(props))
@@ -452,7 +452,7 @@ class KafkaConfigTest {
   def testControllerListenerNameDoesNotMapToPlaintextByDefaultForNonKRaft(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "CONTROLLER://localhost:9092")
     assertBadConfigContainingMessage(props,
       "Error creating broker listeners from 'CONTROLLER://localhost:9092': No security protocol defined for listener CONTROLLER")
@@ -465,7 +465,7 @@ class KafkaConfigTest {
   def testBadListenerProtocol(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "BAD://localhost:9091")
 
     assertFalse(isValidKafkaConfig(props))
@@ -475,7 +475,7 @@ class KafkaConfigTest {
   def testListenerNamesWithAdvertisedListenerUnset(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "CLIENT://localhost:9091,REPLICATION://localhost:9092,INTERNAL://localhost:9093")
     props.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, "CLIENT:SSL,REPLICATION:SSL,INTERNAL:PLAINTEXT")
@@ -499,7 +499,7 @@ class KafkaConfigTest {
   def testListenerAndAdvertisedListenerNames(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "EXTERNAL://localhost:9091,INTERNAL://localhost:9093")
     props.setProperty(KafkaConfig.AdvertisedListenersProp, "EXTERNAL://lb1.example.com:9000,INTERNAL://host1:9093")
@@ -530,7 +530,7 @@ class KafkaConfigTest {
   def testListenerNameMissingFromListenerSecurityProtocolMap(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091,REPLICATION://localhost:9092")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "SSL")
@@ -541,7 +541,7 @@ class KafkaConfigTest {
   def testInterBrokerListenerNameMissingFromListenerSecurityProtocolMap(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "REPLICATION")
@@ -552,7 +552,7 @@ class KafkaConfigTest {
   def testInterBrokerListenerNameAndSecurityProtocolSet(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "SSL")
@@ -564,7 +564,7 @@ class KafkaConfigTest {
   def testCaseInsensitiveListenerProtocol(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "plaintext://localhost:9091,SsL://localhost:9092")
     val config = KafkaConfig.fromProps(props)
     assertEquals(Some("SSL://localhost:9092"), config.listeners.find(_.listenerName.value == "SSL").map(_.connectionString))
@@ -579,7 +579,7 @@ class KafkaConfigTest {
   def testListenerDefaults(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
 
     // configuration with no listeners
     val conf = KafkaConfig.fromProps(props)
@@ -593,7 +593,7 @@ class KafkaConfigTest {
   def testVersionConfiguration(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     val conf = KafkaConfig.fromProps(props)
     assertEquals(MetadataVersion.latestProduction, conf.interBrokerProtocolVersion)
 
@@ -766,7 +766,7 @@ class KafkaConfigTest {
   def testFromPropsInvalid(): Unit = {
     def baseProperties: Properties = {
       val validRequiredProperties = new Properties()
-      validRequiredProperties.setProperty(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
+      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
       validRequiredProperties
     }
     // to ensure a basis is valid - bootstraps all needed validation
@@ -774,25 +774,25 @@ class KafkaConfigTest {
 
     KafkaConfig.configNames.foreach { name =>
       name match {
-        case KafkaConfig.ZkConnectProp => // ignore string
-        case KafkaConfig.ZkSessionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.ZkConnectionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case KafkaConfig.ZkEnableSecureAclsProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case KafkaConfig.ZkMaxInFlightRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
-        case KafkaConfig.ZkSslClientEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case KafkaConfig.ZkClientCnxnSocketProp =>  //ignore string
-        case KafkaConfig.ZkSslKeyStoreLocationProp =>  //ignore string
-        case KafkaConfig.ZkSslKeyStorePasswordProp =>  //ignore string
-        case KafkaConfig.ZkSslKeyStoreTypeProp =>  //ignore string
-        case KafkaConfig.ZkSslTrustStoreLocationProp =>  //ignore string
-        case KafkaConfig.ZkSslTrustStorePasswordProp =>  //ignore string
-        case KafkaConfig.ZkSslTrustStoreTypeProp =>  //ignore string
-        case KafkaConfig.ZkSslProtocolProp =>  //ignore string
-        case KafkaConfig.ZkSslEnabledProtocolsProp =>  //ignore string
-        case KafkaConfig.ZkSslCipherSuitesProp =>  //ignore string
-        case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => //ignore string
-        case KafkaConfig.ZkSslCrlEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case KafkaConfig.ZkSslOcspEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_CONNECT_PROP => // ignore string
+        case ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_PROTOCOL_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP =>  //ignore string
+        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => //ignore string
+        case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
 
         case KafkaConfig.BrokerIdProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.NumNetworkThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
@@ -1051,7 +1051,7 @@ class KafkaConfigTest {
   def testDynamicLogConfigs(): Unit = {
     def baseProperties: Properties = {
       val validRequiredProperties = new Properties()
-      validRequiredProperties.setProperty(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
+      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
       validRequiredProperties
     }
 
@@ -1141,9 +1141,9 @@ class KafkaConfigTest {
   @Test
   def testSpecificProperties(): Unit = {
     val defaults = new Properties()
-    defaults.setProperty(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
+    defaults.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
     // For ZkConnectionTimeoutMs
-    defaults.setProperty(KafkaConfig.ZkSessionTimeoutMsProp, "1234")
+    defaults.setProperty(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, "1234")
     defaults.setProperty(KafkaConfig.BrokerIdGenerationEnableProp, "false")
     defaults.setProperty(KafkaConfig.MaxReservedBrokerIdProp, "1")
     defaults.setProperty(KafkaConfig.BrokerIdProp, "1")
@@ -1187,7 +1187,7 @@ class KafkaConfigTest {
   @Test
   def testNonroutableAdvertisedListeners(): Unit = {
     val props = new Properties()
-    props.setProperty(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
     props.setProperty(KafkaConfig.ListenersProp, "PLAINTEXT://0.0.0.0:9092")
     assertFalse(isValidKafkaConfig(props))
   }
@@ -1601,7 +1601,7 @@ class KafkaConfigTest {
   @Test
   def testSaslJwksEndpointRetryDefaults(): Unit = {
     val props = new Properties()
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     val config = KafkaConfig.fromProps(props)
     assertNotNull(config.getLong(KafkaConfig.SaslOAuthBearerJwksEndpointRetryBackoffMsProp))
     assertNotNull(config.getLong(KafkaConfig.SaslOAuthBearerJwksEndpointRetryBackoffMaxMsProp))
@@ -1769,7 +1769,7 @@ class KafkaConfigTest {
       "If using `zookeeper.metadata.migration.enable` in KRaft mode, `zookeeper.connect` must also be set.",
       assertThrows(classOf[ConfigException], () => KafkaConfig.fromProps(props)).getMessage)
 
-    props.setProperty(KafkaConfig.ZkConnectProp, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
     KafkaConfig.fromProps(props)
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
index bf6312f86319a..035f60dab1db1 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
@@ -25,8 +25,11 @@ import org.junit.jupiter.api.Test
 import java.util.Properties
 import java.net.{InetAddress, ServerSocket}
 import org.apache.kafka.server.common.MetadataVersion
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig
 
+import scala.jdk.CollectionConverters._
+
 class KafkaServerTest extends QuorumTestHarness {
 
   @Test
@@ -64,7 +67,7 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkConfigWhenSaslDisabled(): Unit = {
     val props = new Properties
-    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     assertEquals("false", zkClientConfig.getProperty(JaasUtils.ZK_SASL_CLIENT))
   }
@@ -72,10 +75,10 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkTlsConfigWhenDisabled(): Unit = {
     val props = new Properties
-    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
-    props.put(KafkaConfig.ZkSslClientEnableProp, "false")
+    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "false")
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach { propName =>
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach { propName =>
       assertNull(zkClientConfig.getProperty(propName))
     }
   }
@@ -83,51 +86,51 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkTlsConfigWithTrueValues(): Unit = {
     val props = new Properties
-    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
     // should get correct config for all properties if TLS is enabled
     val someValue = "some_value"
     def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
-      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "true"
-      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "HTTPS"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "true"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "HTTPS"
       case _ => someValue
     }
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     // now check to make sure the values were set correctly
     def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
-      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "true"
-      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "true"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "true"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "true"
       case _ => someValue
     }
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp =>
-      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.getProperty(KafkaConfig.ZkSslConfigToSystemPropertyMap(kafkaProp))))
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp =>
+      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.getProperty(ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(kafkaProp))))
   }
 
   @Test
   def testCreatesProperZkTlsConfigWithFalseAndListValues(): Unit = {
     val props = new Properties
-    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
     // should get correct config for all properties if TLS is enabled
     val someValue = "some_value"
     def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
-      case KafkaConfig.ZkSslClientEnableProp => "true"
-      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "false"
-      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => ""
-      case KafkaConfig.ZkSslEnabledProtocolsProp | KafkaConfig.ZkSslCipherSuitesProp => "A,B"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => "true"
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "false"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => ""
+      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => "A,B"
       case _ => someValue
     }
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     // now check to make sure the values were set correctly
     def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
-      case KafkaConfig.ZkSslClientEnableProp => "true"
-      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "false"
-      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "false"
-      case KafkaConfig.ZkSslEnabledProtocolsProp | KafkaConfig.ZkSslCipherSuitesProp => "A,B"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => "true"
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "false"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "false"
+      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => "A,B"
       case _ => someValue
     }
-    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp =>
-      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.getProperty(KafkaConfig.ZkSslConfigToSystemPropertyMap(kafkaProp))))
+    ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp =>
+      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.getProperty(ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(kafkaProp))))
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
index 844701cc371b3..0e04bad7c8ef2 100644
--- a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
@@ -38,6 +38,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.serialization.{IntegerDeserializer, IntegerSerializer, StringDeserializer, StringSerializer}
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.metadata.BrokerState
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.{BeforeEach, Disabled, TestInfo, Timeout}
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.function.Executable
@@ -150,8 +151,8 @@ class ServerShutdownTest extends KafkaServerTestHarness {
       shutdownKRaftController()
       verifyCleanShutdownAfterFailedStartup[CancellationException]
     } else {
-      propsToChangeUponRestart.setProperty(KafkaConfig.ZkConnectionTimeoutMsProp, "50")
-      propsToChangeUponRestart.setProperty(KafkaConfig.ZkConnectProp, "some.invalid.hostname.foo.bar.local:65535")
+      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, "50")
+      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECT_PROP, "some.invalid.hostname.foo.bar.local:65535")
       verifyCleanShutdownAfterFailedStartup[ZooKeeperClientTimeoutException]
     }
   }
diff --git a/core/src/test/scala/unit/kafka/server/ServerTest.scala b/core/src/test/scala/unit/kafka/server/ServerTest.scala
index d72ad2d7bcdf4..fb80c4b890d05 100644
--- a/core/src/test/scala/unit/kafka/server/ServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerTest.scala
@@ -20,6 +20,7 @@ import java.util.Properties
 
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.metrics.MetricsContext
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
 
@@ -54,7 +55,7 @@ class ServerTest {
 
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, brokerId.toString)
-    props.put(KafkaConfig.ZkConnectProp, "127.0.0.1:0")
+    props.put(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:0")
     val config = KafkaConfig.fromProps(props)
 
     val context = Server.createKafkaMetricsContext(config, clusterId)
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 44243d39ce86a..33a4b27fd8be4 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -74,7 +74,7 @@ import org.apache.kafka.metadata.properties.MetaProperties
 import org.apache.kafka.server.ControllerRequestCompletionHandler
 import org.apache.kafka.server.authorizer.{AuthorizableRequestContext, Authorizer => JAuthorizer}
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion}
-import org.apache.kafka.server.config.Defaults
+import org.apache.kafka.server.config.{Defaults, ZkConfigs}
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.kafka.server.util.MockTime
 import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig, LogDirFailureChannel, ProducerStateManagerConfig}
@@ -359,8 +359,8 @@ object TestUtils extends Logging {
       // controllerQuorumVotersFuture instead.
       props.put(KafkaConfig.QuorumVotersProp, "1000@localhost:0")
     } else {
-      props.put(KafkaConfig.ZkConnectProp, zkConnect)
-      props.put(KafkaConfig.ZkConnectionTimeoutMsProp, "10000")
+      props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect)
+      props.put(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, "10000")
     }
     props.put(KafkaConfig.ReplicaSocketTimeoutMsProp, "1500")
     props.put(KafkaConfig.ControllerSocketTimeoutMsProp, "1500")
diff --git a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
index a1a45784ad7b7..aa97f6831c361 100644
--- a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
@@ -44,7 +44,7 @@ import org.apache.kafka.common.{TopicPartition, Uuid}
 import org.apache.kafka.metadata.LeaderRecoveryState
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.server.common.MetadataVersion
-import org.apache.kafka.server.config.ConfigType
+import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.storage.internals.log.LogConfig
 import org.apache.zookeeper.KeeperException.{Code, NoAuthException, NoNodeException, NodeExistsException}
 import org.apache.zookeeper.{CreateMode, ZooDefs}
@@ -106,7 +106,7 @@ class KafkaZkClientTest extends QuorumTestHarness {
     // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
     // to kafka.zk.EmbeddedZookeeper
     val clientConfig = new ZKClientConfig()
-    val propKey = KafkaConfig.ZkClientCnxnSocketProp
+    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP
     val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
     KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
     val client = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
index 569cb5764b4a9..b4dccaecb0ac8 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
@@ -21,6 +21,7 @@ import kafka.zk.ZkMigrationClient
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.security.PasswordEncoder
+import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.{BeforeEach, TestInfo}
 
 import java.util.Properties
@@ -40,7 +41,7 @@ class ZkMigrationTestHarness extends QuorumTestHarness {
 
   val encoder: PasswordEncoder = {
     val encoderProps = new Properties()
-    encoderProps.put(KafkaConfig.ZkConnectProp, "localhost:1234") // Get around the config validation
+    encoderProps.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:1234") // Get around the config validation
     encoderProps.put(KafkaConfig.PasswordEncoderSecretProp, SECRET) // Zk secret to encrypt the
     val encoderConfig = new KafkaConfig(encoderProps)
     PasswordEncoder.encrypting(encoderConfig.passwordEncoderSecret.get,
diff --git a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
index ebb6ccf44533a..affe6c3839c02 100644
--- a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
@@ -27,6 +27,7 @@ import kafka.utils.TestUtils
 import kafka.server.QuorumTestHarness
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.apache.zookeeper.KeeperException.{Code, NoNodeException}
 import org.apache.zookeeper.Watcher.Event.{EventType, KeeperState}
@@ -102,7 +103,7 @@ class ZooKeeperClientTest extends QuorumTestHarness {
     // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
     // to kafka.zk.EmbeddedZookeeper
     val clientConfig = new ZKClientConfig()
-    val propKey = KafkaConfig.ZkClientCnxnSocketProp
+    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP
     val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
     KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
     val client = newZooKeeperClient(clientConfig = clientConfig)
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
index 75fcf3f24b6bf..f4d17ed074020 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
@@ -62,6 +62,7 @@
 import org.apache.kafka.coordinator.group.GroupCoordinator;
 import org.apache.kafka.server.common.Features;
 import org.apache.kafka.server.common.MetadataVersion;
+import org.apache.kafka.server.config.ZkConfigs;
 import org.mockito.Mockito;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -176,7 +177,7 @@ private List<UpdateMetadataEndpoint> endpoints(final int brokerId) {
 
     private KafkaApis createKafkaApis() {
         Properties kafkaProps =  new Properties();
-        kafkaProps.put(KafkaConfig$.MODULE$.ZkConnectProp(), "zk");
+        kafkaProps.put(ZkConfigs.ZK_CONNECT_PROP, "zk");
         kafkaProps.put(KafkaConfig$.MODULE$.BrokerIdProp(), brokerId + "");
         KafkaConfig config = new KafkaConfig(kafkaProps);
         return new KafkaApisBuilder().
diff --git a/server/src/main/java/org/apache/kafka/server/config/Defaults.java b/server/src/main/java/org/apache/kafka/server/config/Defaults.java
index 0c7ca123b87c1..5b425c0dc49d3 100644
--- a/server/src/main/java/org/apache/kafka/server/config/Defaults.java
+++ b/server/src/main/java/org/apache/kafka/server/config/Defaults.java
@@ -44,16 +44,6 @@
 import java.util.stream.Collectors;
 
 public class Defaults {
-    /** ********* Zookeeper Configuration *********/
-    public static final int ZK_SESSION_TIMEOUT_MS = 18000;
-    public static final boolean ZK_ENABLE_SECURE_ACLS = false;
-    public static final int ZK_MAX_IN_FLIGHT_REQUESTS = 10;
-    public static final boolean ZK_SSL_CLIENT_ENABLE = false;
-    public static final String ZK_SSL_PROTOCOL = "TLSv1.2";
-    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM = "HTTPS";
-    public static final boolean ZK_SSL_CRL_ENABLE = false;
-    public static final boolean ZK_SSL_OCSP_ENABLE = false;
-
     /** ********* General Configuration *********/
     public static final boolean BROKER_ID_GENERATION_ENABLE = true;
     public static final int MAX_RESERVED_BROKER_ID = 1000;
diff --git a/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java b/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java
new file mode 100644
index 0000000000000..2ddd2ef2145f2
--- /dev/null
+++ b/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java
@@ -0,0 +1,145 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.server.config;
+
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public final class ZkConfigs {
+    /** ********* Zookeeper Configuration ***********/
+    public static final String ZK_CONNECT_PROP = "zookeeper.connect";
+    public static final String ZK_SESSION_TIMEOUT_MS_PROP = "zookeeper.session.timeout.ms";
+    public static final String ZK_CONNECTION_TIMEOUT_MS_PROP = "zookeeper.connection.timeout.ms";
+    public static final String ZK_ENABLE_SECURE_ACLS_PROP = "zookeeper.set.acl";
+    public static final String ZK_MAX_IN_FLIGHT_REQUESTS_PROP = "zookeeper.max.in.flight.requests";
+    public static final String ZK_SSL_CLIENT_ENABLE_PROP = "zookeeper.ssl.client.enable";
+    public static final String ZK_CLIENT_CNXN_SOCKET_PROP = "zookeeper.clientCnxnSocket";
+    public static final String ZK_SSL_KEY_STORE_LOCATION_PROP = "zookeeper.ssl.keystore.location";
+    public static final String ZK_SSL_KEY_STORE_PASSWORD_PROP = "zookeeper.ssl.keystore.password";
+    public static final String ZK_SSL_KEY_STORE_TYPE_PROP = "zookeeper.ssl.keystore.type";
+    public static final String ZK_SSL_TRUST_STORE_LOCATION_PROP = "zookeeper.ssl.truststore.location";
+    public static final String ZK_SSL_TRUST_STORE_PASSWORD_PROP = "zookeeper.ssl.truststore.password";
+    public static final String ZK_SSL_TRUST_STORE_TYPE_PROP = "zookeeper.ssl.truststore.type";
+    public static final String ZK_SSL_PROTOCOL_PROP = "zookeeper.ssl.protocol";
+    public static final String ZK_SSL_ENABLED_PROTOCOLS_PROP = "zookeeper.ssl.enabled.protocols";
+    public static final String ZK_SSL_CIPHER_SUITES_PROP = "zookeeper.ssl.cipher.suites";
+    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP = "zookeeper.ssl.endpoint.identification.algorithm";
+    public static final String ZK_SSL_CRL_ENABLE_PROP = "zookeeper.ssl.crl.enable";
+    public static final String ZK_SSL_OCSP_ENABLE_PROP = "zookeeper.ssl.ocsp.enable";
+
+    public static final String ZK_CONNECT_DOC = "Specifies the ZooKeeper connection string in the form <code>hostname:port</code> where host and port are the " +
+        "host and port of a ZooKeeper server. To allow connecting through other ZooKeeper nodes when that ZooKeeper machine is " +
+        "down you can also specify multiple hosts in the form <code>hostname1:port1,hostname2:port2,hostname3:port3</code>.\n" +
+        "The server can also have a ZooKeeper chroot path as part of its ZooKeeper connection string which puts its data under some path in the global ZooKeeper namespace. " +
+        "For example to give a chroot path of <code>/chroot/path</code> you would give the connection string as <code>hostname1:port1,hostname2:port2,hostname3:port3/chroot/path</code>.";
+    public static final String ZK_SESSION_TIMEOUT_MS_DOC = "Zookeeper session timeout";
+    public static final String ZK_CONNECTION_TIMEOUT_MS_DOC = "The max time that the client waits to establish a connection to ZooKeeper. If not set, the value in " + ZK_SESSION_TIMEOUT_MS_PROP + " is used";
+    public static final String ZK_ENABLE_SECURE_ACLS_DOC = "Set client to use secure ACLs";
+    public static final String ZK_MAX_IN_FLIGHT_REQUESTS_DOC = "The maximum number of unacknowledged requests the client will send to ZooKeeper before blocking.";
+    public static final String ZK_SSL_CLIENT_ENABLE_DOC;
+    public static final String ZK_CLIENT_CNXN_SOCKET_DOC;
+    public static final String ZK_SSL_KEY_STORE_LOCATION_DOC;
+    public static final String ZK_SSL_KEY_STORE_PASSWORD_DOC;
+    public static final String ZK_SSL_KEY_STORE_TYPE_DOC;
+    public static final String ZK_SSL_TRUST_STORE_LOCATION_DOC;
+    public static final String ZK_SSL_TRUST_STORE_PASSWORD_DOC;
+    public static final String ZK_SSL_TRUST_STORE_TYPE_DOC;
+    public static final String ZK_SSL_PROTOCOL_DOC;
+    public static final String ZK_SSL_ENABLED_PROTOCOLS_DOC;
+    public static final String ZK_SSL_CIPHER_SUITES_DOC;
+    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC;
+    public static final String ZK_SSL_CRL_ENABLE_DOC;
+    public static final String ZK_SSL_OCSP_ENABLE_DOC;
+
+    // a map from the Kafka config to the corresponding ZooKeeper Java system property
+    public static final Map<String, String> ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP;
+
+    public static final int ZK_SESSION_TIMEOUT_MS = 18000;
+    public static final boolean ZK_ENABLE_SECURE_ACLS = false;
+    public static final int ZK_MAX_IN_FLIGHT_REQUESTS = 10;
+    public static final boolean ZK_SSL_CLIENT_ENABLE = false;
+    public static final String ZK_SSL_PROTOCOL = "TLSv1.2";
+    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM = "HTTPS";
+    public static final boolean ZK_SSL_CRL_ENABLE = false;
+    public static final boolean ZK_SSL_OCSP_ENABLE = false;
+
+    // See ZKClientConfig.SECURE_CLIENT
+    private static final String SECURE_CLIENT = "zookeeper.client.secure";
+    // See ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET
+    private static final String ZOOKEEPER_CLIENT_CNXN_SOCKET = "zookeeper.clientCnxnSocket";
+
+    static {
+        Map<String, String> zkSslConfigToSystemPropertyMap = new HashMap<>();
+
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CLIENT_ENABLE_PROP, SECURE_CLIENT);
+        zkSslConfigToSystemPropertyMap.put(ZK_CLIENT_CNXN_SOCKET_PROP, ZOOKEEPER_CLIENT_CNXN_SOCKET);
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_LOCATION_PROP, "zookeeper.ssl.keyStore.location");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_PASSWORD_PROP, "zookeeper.ssl.keyStore.password");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_TYPE_PROP, "zookeeper.ssl.keyStore.type");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_LOCATION_PROP, "zookeeper.ssl.trustStore.location");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_PASSWORD_PROP, "zookeeper.ssl.trustStore.password");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_TYPE_PROP, "zookeeper.ssl.trustStore.type");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_PROTOCOL_PROP, "zookeeper.ssl.protocol");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENABLED_PROTOCOLS_PROP, "zookeeper.ssl.enabledProtocols");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CIPHER_SUITES_PROP, "zookeeper.ssl.ciphersuites");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, "zookeeper.ssl.hostnameVerification");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CRL_ENABLE_PROP, "zookeeper.ssl.crl");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_OCSP_ENABLE_PROP, "zookeeper.ssl.ocsp");
+
+        ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP = Collections.unmodifiableMap(zkSslConfigToSystemPropertyMap);
+
+        ZK_SSL_CLIENT_ENABLE_DOC = "Set client to use TLS when connecting to ZooKeeper." +
+            " An explicit value overrides any value set via the <code>zookeeper.client.secure</code> system property (note the different name)." +
+            " Defaults to false if neither is set; when true, <code>" + ZK_CLIENT_CNXN_SOCKET_PROP + "</code> must be set (typically to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code>); other values to set may include " +
+            ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.keySet().stream().filter(x -> !x.equals(ZK_SSL_CLIENT_ENABLE_PROP) && !x.equals(ZK_CLIENT_CNXN_SOCKET_PROP)).sorted().collect(Collectors.joining("<code>", "</code>, <code>", "</code>"));
+        ZK_CLIENT_CNXN_SOCKET_DOC = "Typically set to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code> when using TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_CLIENT_CNXN_SOCKET_PROP) + "</code> system property.";
+        ZK_SSL_KEY_STORE_LOCATION_DOC = "Keystore location when using a client-side certificate with TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_LOCATION_PROP) + "</code> system property (note the camelCase).";
+        ZK_SSL_KEY_STORE_PASSWORD_DOC = "Keystore password when using a client-side certificate with TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_PASSWORD_PROP) + "</code> system property (note the camelCase)." +
+            " Note that ZooKeeper does not support a key password different from the keystore password, so be sure to set the key password in the keystore to be identical to the keystore password; otherwise the connection attempt to Zookeeper will fail.";
+        ZK_SSL_KEY_STORE_TYPE_DOC = "Keystore type when using a client-side certificate with TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_TYPE_PROP) + "</code> system property (note the camelCase)." +
+            " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the keystore.";
+        ZK_SSL_TRUST_STORE_LOCATION_DOC = "Truststore location when using TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_LOCATION_PROP) + "</code> system property (note the camelCase).";
+        ZK_SSL_TRUST_STORE_PASSWORD_DOC = "Truststore password when using TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_PASSWORD_PROP) + "</code> system property (note the camelCase).";
+        ZK_SSL_TRUST_STORE_TYPE_DOC = "Truststore type when using TLS connectivity to ZooKeeper." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_TYPE_PROP) + "</code> system property (note the camelCase)." +
+            " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the truststore.";
+        ZK_SSL_PROTOCOL_DOC = "Specifies the protocol to be used in ZooKeeper TLS negotiation." +
+            " An explicit value overrides any value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_PROTOCOL_PROP) + "</code> system property.";
+        ZK_SSL_ENABLED_PROTOCOLS_DOC = "Specifies the enabled protocol(s) in ZooKeeper TLS negotiation (csv)." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENABLED_PROTOCOLS_PROP) + "</code> system property (note the camelCase)." +
+            " The default value of <code>null</code> means the enabled protocol will be the value of the <code>" + ZK_SSL_PROTOCOL_PROP + "</code> configuration property.";
+        ZK_SSL_CIPHER_SUITES_DOC = "Specifies the enabled cipher suites to be used in ZooKeeper TLS negotiation (csv)." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CIPHER_SUITES_PROP) + "</code> system property (note the single word \"ciphersuites\")." +
+            " The default value of <code>null</code> means the list of enabled cipher suites is determined by the Java runtime being used.";
+        ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC = "Specifies whether to enable hostname verification in the ZooKeeper TLS negotiation process, with (case-insensitively) \"https\" meaning ZooKeeper hostname verification is enabled and an explicit blank value meaning it is disabled (disabling it is only recommended for testing purposes)." +
+            " An explicit value overrides any \"true\" or \"false\" value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP) + "</code> system property (note the different name and values; true implies https and false implies blank).";
+        ZK_SSL_CRL_ENABLE_DOC = "Specifies whether to enable Certificate Revocation List in the ZooKeeper TLS protocols." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CRL_ENABLE_PROP) + "</code> system property (note the shorter name).";
+        ZK_SSL_OCSP_ENABLE_DOC = "Specifies whether to enable Online Certificate Status Protocol in the ZooKeeper TLS protocols." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_OCSP_ENABLE_PROP) + "</code> system property (note the shorter name).";
+    }
+
+}
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
index 4232e1d74c978..89af1d5729e84 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
@@ -23,6 +23,7 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.server.config.ConfigType;
+import org.apache.kafka.server.config.ZkConfigs;
 import org.apache.kafka.server.util.MockTime;
 import org.apache.kafka.storage.internals.log.CleanerConfig;
 import org.apache.kafka.test.TestCondition;
@@ -109,7 +110,7 @@ public void start() throws IOException {
         zookeeper = new EmbeddedZookeeper();
         log.debug("ZooKeeper instance is running at {}", zKConnectString());
 
-        brokerConfig.put(KafkaConfig.ZkConnectProp(), zKConnectString());
+        brokerConfig.put(ZkConfigs.ZK_CONNECT_PROP, zKConnectString());
         putIfAbsent(brokerConfig, KafkaConfig.ListenersProp(), "PLAINTEXT://localhost:" + DEFAULT_BROKER_PORT);
         putIfAbsent(brokerConfig, KafkaConfig.DeleteTopicEnableProp(), true);
         putIfAbsent(brokerConfig, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, 2 * 1024 * 1024L);
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
index 7f945ffe06d1b..4b57d576ae1c7 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.config.ZkConfigs;
 import org.apache.kafka.server.util.MockTime;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -90,7 +91,7 @@ private Properties effectiveConfigFrom(final Properties initialConfig) {
         effectiveConfig.put(KafkaConfig.AutoCreateTopicsEnableProp(), true);
         effectiveConfig.put(KafkaConfig.MessageMaxBytesProp(), 1000000);
         effectiveConfig.put(KafkaConfig.ControlledShutdownEnableProp(), true);
-        effectiveConfig.put(KafkaConfig.ZkSessionTimeoutMsProp(), 10000);
+        effectiveConfig.put(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, 10000);
 
         effectiveConfig.putAll(initialConfig);
         effectiveConfig.setProperty(KafkaConfig.LogDirProp(), logDir.getAbsolutePath());

From cf1ba099c0723f9cf65dda4cd334d36b7ede6327 Mon Sep 17 00:00:00 2001
From: Sanskar Jhajharia <122860866+sjhajharia@users.noreply.github.com>
Date: Thu, 28 Mar 2024 00:43:32 +0530
Subject: [PATCH 219/258] MINOR: Renaming the `Abortable_Transaction` error to
 `Transaction_Abortable` (#15609)

This is a follow-up to this PR (https://github.com/apache/kafka/pull/15486) which introduced the new ABORTABLE_TRANSACTION error as a part of KIP-890 efforts. However on further discussion, we seem to gain consensus that the error should be rather named as TRANSACTION_ABORTABLE.

This PR aims to address the same. There are no changes in the code apart from that.

Reviewers: Justine Olshan <jolshan@confluent.io>, Igor Soarez <soarez@apple.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../internals/TransactionManager.java         | 12 ++---
 ...ava => TransactionAbortableException.java} |  4 +-
 .../apache/kafka/common/protocol/Errors.java  |  4 +-
 .../message/AddOffsetsToTxnRequest.json       |  2 +-
 .../message/AddOffsetsToTxnResponse.json      |  2 +-
 .../message/AddPartitionsToTxnRequest.json    |  2 +-
 .../message/AddPartitionsToTxnResponse.json   |  2 +-
 .../common/message/EndTxnRequest.json         |  2 +-
 .../common/message/EndTxnResponse.json        |  2 +-
 .../message/FindCoordinatorRequest.json       |  2 +-
 .../message/FindCoordinatorResponse.json      |  2 +-
 .../common/message/InitProducerIdRequest.json |  2 +-
 .../message/InitProducerIdResponse.json       |  2 +-
 .../common/message/ProduceRequest.json        |  2 +-
 .../common/message/ProduceResponse.json       |  2 +-
 .../message/TxnOffsetCommitRequest.json       |  2 +-
 .../message/TxnOffsetCommitResponse.json      |  2 +-
 .../producer/internals/SenderTest.java        | 12 ++---
 .../internals/TransactionManagerTest.java     | 54 +++++++++----------
 .../transaction/TransactionCoordinator.scala  |  2 +-
 .../server/AddPartitionsToTxnManager.scala    |  4 +-
 .../TransactionCoordinatorTest.scala          |  4 +-
 .../AddPartitionsToTxnManagerTest.scala       | 28 +++++-----
 .../AddPartitionsToTxnRequestServerTest.scala |  2 +-
 24 files changed, 77 insertions(+), 77 deletions(-)
 rename clients/src/main/java/org/apache/kafka/common/errors/{AbortableTransactionException.java => TransactionAbortableException.java} (87%)

diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
index 4ff54b1759875..fab24c5e5e270 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java
@@ -1330,7 +1330,7 @@ public void handleResponse(AbstractResponse response) {
                 // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
                 // just treat it the same as PRODUCE_FENCED.
                 fatalError(Errors.PRODUCER_FENCED.exception());
-            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+            } else if (error == Errors.TRANSACTION_ABORTABLE) {
                 abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unexpected error in InitProducerIdResponse; " + error.message()));
@@ -1401,7 +1401,7 @@ public void handleResponse(AbstractResponse response) {
                 } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
                     abortableErrorIfPossible(error.exception());
                     return;
-                } else if (error == Errors.ABORTABLE_TRANSACTION) {
+                } else if (error == Errors.TRANSACTION_ABORTABLE) {
                     abortableError(error.exception());
                     return;
                 } else {
@@ -1507,7 +1507,7 @@ public void handleResponse(AbstractResponse response) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                 abortableError(GroupAuthorizationException.forGroupId(key));
-            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+            } else if (error == Errors.TRANSACTION_ABORTABLE) {
                 abortableError(error.exception());
             } else {
                 fatalError(new KafkaException(String.format("Could not find a coordinator with type %s with key %s due to " +
@@ -1562,7 +1562,7 @@ public void handleResponse(AbstractResponse response) {
                 fatalError(error.exception());
             } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
                 abortableErrorIfPossible(error.exception());
-            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+            } else if (error == Errors.TRANSACTION_ABORTABLE) {
                 abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unhandled error in EndTxnResponse: " + error.message()));
@@ -1622,7 +1622,7 @@ public void handleResponse(AbstractResponse response) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                 abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
-            } else if (error == Errors.ABORTABLE_TRANSACTION) {
+            } else if (error == Errors.TRANSACTION_ABORTABLE) {
                 abortableError(error.exception());
             } else {
                 fatalError(new KafkaException("Unexpected error in AddOffsetsToTxnResponse: " + error.message()));
@@ -1687,7 +1687,7 @@ public void handleResponse(AbstractResponse response) {
                     abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
                     break;
                 } else if (error == Errors.FENCED_INSTANCE_ID ||
-                        error == Errors.ABORTABLE_TRANSACTION) {
+                        error == Errors.TRANSACTION_ABORTABLE) {
                     abortableError(error.exception());
                     break;
                 } else if (error == Errors.UNKNOWN_MEMBER_ID
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java b/clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortableException.java
similarity index 87%
rename from clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java
rename to clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortableException.java
index cfea649e750c6..aa592d552bf00 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/AbortableTransactionException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortableException.java
@@ -16,8 +16,8 @@
  */
 package org.apache.kafka.common.errors;
 
-public class AbortableTransactionException extends ApiException {
-    public AbortableTransactionException(String message) {
+public class TransactionAbortableException extends ApiException {
+    public TransactionAbortableException(String message) {
         super(message);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
index 30bc250c6dec2..900d191c8f9d4 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java
@@ -138,7 +138,7 @@
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedSaslMechanismException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.errors.AbortableTransactionException;
+import org.apache.kafka.common.errors.TransactionAbortableException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -394,7 +394,7 @@ public enum Errors {
     UNKNOWN_SUBSCRIPTION_ID(117, "Client sent a push telemetry request with an invalid or outdated subscription ID.", UnknownSubscriptionIdException::new),
     TELEMETRY_TOO_LARGE(118, "Client sent a push telemetry request larger than the maximum size the broker will accept.", TelemetryTooLargeException::new),
     INVALID_REGISTRATION(119, "The controller has considered the broker registration to be invalid.", InvalidRegistrationException::new),
-    ABORTABLE_TRANSACTION(120, "The server encountered an error with the transaction. The client can abort the transaction to continue using this transactional ID.", AbortableTransactionException::new);
+    TRANSACTION_ABORTABLE(120, "The server encountered an error with the transaction. The client can abort the transaction to continue using this transactional ID.", TransactionAbortableException::new);
 
     private static final Logger log = LoggerFactory.getLogger(Errors.class);
 
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
index 9d7b63c313332..157ae20c0a41d 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
@@ -24,7 +24,7 @@
   //
   // Version 3 enables flexible versions.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
index 6b3b1c481d663..6a713fea1af62 100644
--- a/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.json
@@ -23,7 +23,7 @@
   //
   // Version 3 enables flexible versions.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
index 139d1436a6520..2270f7a8469f5 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
@@ -26,7 +26,7 @@
   //
   // Version 4 adds VerifyOnly field to check if partitions are already in transaction and adds support to batch multiple transactions.
   //
-  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   // Versions 3 and below will be exclusively used by clients and versions 4 and above will be used by brokers.
   "latestVersionUnstable": false,
   "validVersions": "0-5",
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
index a2af388dba55f..6c4eefd2cc001 100644
--- a/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.json
@@ -25,7 +25,7 @@
   //
   // Version 4 adds support to batch multiple transactions and a top level error code.
   //
-  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-5",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/EndTxnRequest.json b/clients/src/main/resources/common/message/EndTxnRequest.json
index 5bf6c577343d6..bc66adcf50ade 100644
--- a/clients/src/main/resources/common/message/EndTxnRequest.json
+++ b/clients/src/main/resources/common/message/EndTxnRequest.json
@@ -24,7 +24,7 @@
   //
   // Version 3 enables flexible versions.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/EndTxnResponse.json b/clients/src/main/resources/common/message/EndTxnResponse.json
index 53b0250f808b9..08ac6cddd3895 100644
--- a/clients/src/main/resources/common/message/EndTxnResponse.json
+++ b/clients/src/main/resources/common/message/EndTxnResponse.json
@@ -23,7 +23,7 @@
   //
   // Version 3 enables flexible versions.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/FindCoordinatorRequest.json b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
index e6786f5b10e3e..42b2f4c891ad5 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorRequest.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorRequest.json
@@ -26,7 +26,7 @@
   //
   // Version 4 adds support for batching via CoordinatorKeys (KIP-699)
   //
-  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-5",
   "deprecatedVersions": "0",
   "flexibleVersions": "3+",
diff --git a/clients/src/main/resources/common/message/FindCoordinatorResponse.json b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
index a744a1928d6c7..860d655a252b2 100644
--- a/clients/src/main/resources/common/message/FindCoordinatorResponse.json
+++ b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
@@ -25,7 +25,7 @@
   //
   // Version 4 adds support for batching via Coordinators (KIP-699)
   //
-  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-5",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/InitProducerIdRequest.json b/clients/src/main/resources/common/message/InitProducerIdRequest.json
index 92e6a6b25376d..39f546dbc0439 100644
--- a/clients/src/main/resources/common/message/InitProducerIdRequest.json
+++ b/clients/src/main/resources/common/message/InitProducerIdRequest.json
@@ -26,7 +26,7 @@
   //
   // Version 4 adds the support for new error code PRODUCER_FENCED.
   //
-  // Verison 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Verison 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-5",
   "flexibleVersions": "2+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/InitProducerIdResponse.json b/clients/src/main/resources/common/message/InitProducerIdResponse.json
index c0f10b2e851a6..c5dfec6e321cb 100644
--- a/clients/src/main/resources/common/message/InitProducerIdResponse.json
+++ b/clients/src/main/resources/common/message/InitProducerIdResponse.json
@@ -25,7 +25,7 @@
   //
   // Version 4 adds the support for new error code PRODUCER_FENCED.
   //
-  // Version 5 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 5 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-5",
   "flexibleVersions": "2+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/ProduceRequest.json b/clients/src/main/resources/common/message/ProduceRequest.json
index d396d66070da9..ae01fe5c8c08a 100644
--- a/clients/src/main/resources/common/message/ProduceRequest.json
+++ b/clients/src/main/resources/common/message/ProduceRequest.json
@@ -36,7 +36,7 @@
   //
   // Version 10 is the same as version 9 (KIP-951).
   //
-  // Version 11 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 11 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-11",
   "deprecatedVersions": "0-6",
   "flexibleVersions": "9+",
diff --git a/clients/src/main/resources/common/message/ProduceResponse.json b/clients/src/main/resources/common/message/ProduceResponse.json
index 1c097d3bc3b5d..92c7a2223da09 100644
--- a/clients/src/main/resources/common/message/ProduceResponse.json
+++ b/clients/src/main/resources/common/message/ProduceResponse.json
@@ -35,7 +35,7 @@
   //
   // Version 10 adds 'CurrentLeader' and 'NodeEndpoints' as tagged fields (KIP-951)
   //
-  // Version 11 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 11 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-11",
   "flexibleVersions": "9+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
index 1df18c64b52d8..3cb63aa8fb8e3 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
@@ -24,7 +24,7 @@
   //
   // Version 3 adds the member.id, group.instance.id and generation.id.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
index 0f6c1f272418a..1a04cef9d5e4a 100644
--- a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
@@ -23,7 +23,7 @@
   //
   // Version 3 adds illegal generation, fenced instance id, and unknown member id errors.
   //
-  // Version 4 adds support for new error code ABORTABLE_TRANSACTION (KIP-890).
+  // Version 4 adds support for new error code TRANSACTION_ABORTABLE (KIP-890).
   "validVersions": "0-4",
   "flexibleVersions": "3+",
   "fields": [
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
index b9c734a2cb1bf..4ef9dab4d096c 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
@@ -42,7 +42,7 @@
 import org.apache.kafka.common.errors.TransactionAbortedException;
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.errors.AbortableTransactionException;
+import org.apache.kafka.common.errors.TransactionAbortableException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.AddPartitionsToTxnResponseData;
 import org.apache.kafka.common.message.ApiMessageType;
@@ -3149,10 +3149,10 @@ public void testInvalidTxnStateIsAnAbortableError() throws Exception {
     }
 
     @Test
-    public void testAbortableTxnExceptionIsAnAbortableError() throws Exception {
+    public void testTransactionAbortablenExceptionIsAnAbortableError() throws Exception {
         ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
         apiVersions.update("0", NodeApiVersions.create(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 3));
-        TransactionManager txnManager = new TransactionManager(logContext, "textAbortableTxnException", 60000, 100, apiVersions);
+        TransactionManager txnManager = new TransactionManager(logContext, "textTransactionAbortableException", 60000, 100, apiVersions);
 
         setupWithTransactionState(txnManager);
         doInitTransactions(txnManager, producerIdAndEpoch);
@@ -3164,11 +3164,11 @@ public void testAbortableTxnExceptionIsAnAbortableError() throws Exception {
 
         Future<RecordMetadata> request = appendToAccumulator(tp0);
         sender.runOnce();  // send request
-        sendIdempotentProducerResponse(0, tp0, Errors.ABORTABLE_TRANSACTION, -1);
+        sendIdempotentProducerResponse(0, tp0, Errors.TRANSACTION_ABORTABLE, -1);
 
-        // Return AbortableTransactionException error. It should be abortable.
+        // Return TransactionAbortableException error. It should be abortable.
         sender.runOnce();
-        assertFutureFailure(request, AbortableTransactionException.class);
+        assertFutureFailure(request, TransactionAbortableException.class);
         assertTrue(txnManager.hasAbortableError());
         TransactionalRequestResult result = txnManager.beginAbort();
         sender.runOnce();
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
index 4d8c445be0ac5..51299ad337ea1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java
@@ -40,7 +40,7 @@
 import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.errors.AbortableTransactionException;
+import org.apache.kafka.common.errors.TransactionAbortableException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
@@ -3516,22 +3516,22 @@ public void testForegroundInvalidStateTransitionIsRecoverable() {
     }
 
     @Test
-    public void testAbortableTransactionExceptionInInitProducerId() {
+    public void testTransactionAbortableExceptionInInitProducerId() {
         TransactionalRequestResult initPidResult = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
         runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
-        prepareInitPidResponse(Errors.ABORTABLE_TRANSACTION, false, producerId, RecordBatch.NO_PRODUCER_EPOCH);
+        prepareInitPidResponse(Errors.TRANSACTION_ABORTABLE, false, producerId, RecordBatch.NO_PRODUCER_EPOCH);
         runUntil(transactionManager::hasError);
         assertTrue(initPidResult.isCompleted());
         assertFalse(initPidResult.isSuccessful());
-        assertThrows(AbortableTransactionException.class, initPidResult::await);
-        assertAbortableError(AbortableTransactionException.class);
+        assertThrows(TransactionAbortableException.class, initPidResult::await);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     @Test
-    public void testAbortableTransactionExceptionInAddPartitions() {
+    public void testTransactionAbortableExceptionInAddPartitions() {
         final TopicPartition tp = new TopicPartition("foo", 0);
 
         doInitTransactions();
@@ -3539,15 +3539,15 @@ public void testAbortableTransactionExceptionInAddPartitions() {
         transactionManager.beginTransaction();
         transactionManager.maybeAddPartition(tp);
 
-        prepareAddPartitionsToTxn(tp, Errors.ABORTABLE_TRANSACTION);
+        prepareAddPartitionsToTxn(tp, Errors.TRANSACTION_ABORTABLE);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(transactionManager.lastError() instanceof TransactionAbortableException);
 
-        assertAbortableError(AbortableTransactionException.class);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     @Test
-    public void testAbortableTransactionExceptionInFindCoordinator() {
+    public void testTransactionAbortableExceptionInFindCoordinator() {
         doInitTransactions();
 
         transactionManager.beginTransaction();
@@ -3557,19 +3557,19 @@ public void testAbortableTransactionExceptionInFindCoordinator() {
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
         runUntil(() -> !transactionManager.hasPartitionsToAdd());
 
-        prepareFindCoordinatorResponse(Errors.ABORTABLE_TRANSACTION, false, CoordinatorType.GROUP, consumerGroupId);
+        prepareFindCoordinatorResponse(Errors.TRANSACTION_ABORTABLE, false, CoordinatorType.GROUP, consumerGroupId);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(transactionManager.lastError() instanceof TransactionAbortableException);
 
         runUntil(sendOffsetsResult::isCompleted);
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
+        assertTrue(sendOffsetsResult.error() instanceof TransactionAbortableException);
 
-        assertAbortableError(AbortableTransactionException.class);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     @Test
-    public void testAbortableTransactionExceptionInEndTxn() throws InterruptedException {
+    public void testTransactionAbortableExceptionInEndTxn() throws InterruptedException {
         doInitTransactions();
 
         transactionManager.beginTransaction();
@@ -3581,7 +3581,7 @@ public void testAbortableTransactionExceptionInEndTxn() throws InterruptedExcept
         assertFalse(responseFuture.isDone());
         prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
         prepareProduceResponse(Errors.NONE, producerId, epoch);
-        prepareEndTxnResponse(Errors.ABORTABLE_TRANSACTION, TransactionResult.COMMIT, producerId, epoch);
+        prepareEndTxnResponse(Errors.TRANSACTION_ABORTABLE, TransactionResult.COMMIT, producerId, epoch);
 
         runUntil(commitResult::isCompleted);
         runUntil(responseFuture::isDone);
@@ -3590,11 +3590,11 @@ public void testAbortableTransactionExceptionInEndTxn() throws InterruptedExcept
         assertFalse(commitResult.isSuccessful());
         assertTrue(commitResult.isAcked());
 
-        assertAbortableError(AbortableTransactionException.class);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     @Test
-    public void testAbortableTransactionExceptionInAddOffsetsToTxn() {
+    public void testTransactionAbortableExceptionInAddOffsetsToTxn() {
         final TopicPartition tp = new TopicPartition("foo", 0);
 
         doInitTransactions();
@@ -3603,18 +3603,18 @@ public void testAbortableTransactionExceptionInAddOffsetsToTxn() {
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
                 singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
-        prepareAddOffsetsToTxnResponse(Errors.ABORTABLE_TRANSACTION, consumerGroupId, producerId, epoch);
+        prepareAddOffsetsToTxnResponse(Errors.TRANSACTION_ABORTABLE, consumerGroupId, producerId, epoch);
         runUntil(transactionManager::hasError);
-        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(transactionManager.lastError() instanceof TransactionAbortableException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
+        assertTrue(sendOffsetsResult.error() instanceof TransactionAbortableException);
 
-        assertAbortableError(AbortableTransactionException.class);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     @Test
-    public void testAbortableTransactionExceptionInTxnOffsetCommit() {
+    public void testTransactionAbortableExceptionInTxnOffsetCommit() {
         final TopicPartition tp = new TopicPartition("foo", 0);
 
         doInitTransactions();
@@ -3625,14 +3625,14 @@ public void testAbortableTransactionExceptionInTxnOffsetCommit() {
 
         prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.ABORTABLE_TRANSACTION));
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.TRANSACTION_ABORTABLE));
         runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.lastError() instanceof AbortableTransactionException);
+        assertTrue(transactionManager.lastError() instanceof TransactionAbortableException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
-        assertTrue(sendOffsetsResult.error() instanceof AbortableTransactionException);
-        assertAbortableError(AbortableTransactionException.class);
+        assertTrue(sendOffsetsResult.error() instanceof TransactionAbortableException);
+        assertAbortableError(TransactionAbortableException.class);
     }
 
     private FutureRecordMetadata appendToAccumulator(TopicPartition tp) throws InterruptedException {
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
index 4d9ba7fa55893..aca2e601accf8 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
@@ -369,7 +369,7 @@ class TransactionCoordinator(txnConfig: TransactionConfig,
                   if (txnMetadata.topicPartitions.contains(part))
                     (part, Errors.NONE)
                   else
-                    (part, Errors.ABORTABLE_TRANSACTION)
+                    (part, Errors.TRANSACTION_ABORTABLE)
                 }.toMap)
               }
             }
diff --git a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
index ec909580a8400..3aad9b74acc62 100644
--- a/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
+++ b/core/src/main/scala/kafka/server/AddPartitionsToTxnManager.scala
@@ -45,7 +45,7 @@ object AddPartitionsToTxnManager {
 /**
  * This is an enum which handles the Partition Response based on the Request Version and the exact operation
  *    defaultError:       This is the default workflow which maps to cases when the Produce Request Version or the Txn_offset_commit request was lower than the first version supporting the new Error Class
- *    genericError:       This maps to the case when the clients are updated to handle the AbortableTxnException
+ *    genericError:       This maps to the case when the clients are updated to handle the TransactionAbortableException
  *    addPartition:       This is a WIP. To be updated as a part of KIP-890 Part 2
  */
 sealed trait SupportedOperation
@@ -226,7 +226,7 @@ class AddPartitionsToTxnManager(
                   val code =
                     if (partitionResult.partitionErrorCode == Errors.PRODUCER_FENCED.code)
                       Errors.INVALID_PRODUCER_EPOCH.code
-                    else if (partitionResult.partitionErrorCode() == Errors.ABORTABLE_TRANSACTION.code && transactionDataAndCallbacks.supportedOperation != genericError) // For backward compatibility with clients.
+                    else if (partitionResult.partitionErrorCode() == Errors.TRANSACTION_ABORTABLE.code && transactionDataAndCallbacks.supportedOperation != genericError) // For backward compatibility with clients.
                       Errors.INVALID_TXN_STATE.code
                     else
                       partitionResult.partitionErrorCode
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
index 1da8d664801e0..ddcaa4ca8bbfa 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
@@ -283,7 +283,7 @@ class TransactionCoordinatorTest {
 
     coordinator.handleVerifyPartitionsInTransaction(transactionalId, 0L, 0, partitions, verifyPartitionsInTxnCallback)
     errors.foreach { case (_, error) =>
-      assertEquals(Errors.ABORTABLE_TRANSACTION, error)
+      assertEquals(Errors.TRANSACTION_ABORTABLE, error)
     }
   }
 
@@ -399,7 +399,7 @@ class TransactionCoordinatorTest {
     val extraPartitions = partitions ++ Set(new TopicPartition("topic2", 0))
     
     coordinator.handleVerifyPartitionsInTransaction(transactionalId, 0L, 0, extraPartitions, verifyPartitionsInTxnCallback)
-    assertEquals(Errors.ABORTABLE_TRANSACTION, errors(new TopicPartition("topic2", 0)))
+    assertEquals(Errors.TRANSACTION_ABORTABLE, errors(new TopicPartition("topic2", 0)))
     assertEquals(Errors.NONE, errors(new TopicPartition("topic1", 0)))
     verify(transactionManager).getTransactionState(ArgumentMatchers.eq(transactionalId))
   }
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
index 5ee9cf8302698..a2569647092b3 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnManagerTest.scala
@@ -299,31 +299,31 @@ class AddPartitionsToTxnManagerTest {
     assertEquals(expectedTransaction1Errors, transaction1Errors)
     assertEquals(expectedTransaction2Errors, transaction2Errors)
 
-    val preConvertedAbortableTransaction1Errors = topicPartitions.map(_ -> Errors.ABORTABLE_TRANSACTION).toMap
-    val preConvertedAbortableTransaction2Errors = Map(new TopicPartition("foo", 1) -> Errors.NONE,
-      new TopicPartition("foo", 2) -> Errors.ABORTABLE_TRANSACTION,
+    val preConvertedTransactionAbortableErrorsTxn1 = topicPartitions.map(_ -> Errors.TRANSACTION_ABORTABLE).toMap
+    val preConvertedTransactionAbortableErrorsTxn2 = Map(new TopicPartition("foo", 1) -> Errors.NONE,
+      new TopicPartition("foo", 2) -> Errors.TRANSACTION_ABORTABLE,
       new TopicPartition("foo", 3) -> Errors.NONE)
-    val abortableTransaction1ErrorResponse = AddPartitionsToTxnResponse.resultForTransaction(transactionalId1, preConvertedAbortableTransaction1Errors.asJava)
-    val abortableTransaction2ErrorResponse = AddPartitionsToTxnResponse.resultForTransaction(transactionalId2, preConvertedAbortableTransaction2Errors.asJava)
+    val transactionAbortableErrorResponseTxn1 = AddPartitionsToTxnResponse.resultForTransaction(transactionalId1, preConvertedTransactionAbortableErrorsTxn1.asJava)
+    val transactionAbortableErrorResponseTxn2 = AddPartitionsToTxnResponse.resultForTransaction(transactionalId2, preConvertedTransactionAbortableErrorsTxn2.asJava)
     val mixedErrorsAddPartitionsResponseAbortableError = new AddPartitionsToTxnResponse(new AddPartitionsToTxnResponseData()
-      .setResultsByTransaction(new AddPartitionsToTxnResultCollection(Seq(abortableTransaction1ErrorResponse, abortableTransaction2ErrorResponse).iterator.asJava)))
+      .setResultsByTransaction(new AddPartitionsToTxnResultCollection(Seq(transactionAbortableErrorResponseTxn1, transactionAbortableErrorResponseTxn2).iterator.asJava)))
     val mixedAbortableErrorsResponse = clientResponse(mixedErrorsAddPartitionsResponseAbortableError)
 
-    val expectedAbortableTransaction1ErrorsLowerVersion = topicPartitions.map(_ -> Errors.INVALID_TXN_STATE).toMap
-    val expectedAbortableTransaction2ErrorsLowerVersion = Map(new TopicPartition("foo", 2) -> Errors.INVALID_TXN_STATE)
+    val expectedTransactionAbortableErrorsTxn1LowerVersion = topicPartitions.map(_ -> Errors.INVALID_TXN_STATE).toMap
+    val expectedTransactionAbortableErrorsTxn2LowerVersion = Map(new TopicPartition("foo", 2) -> Errors.INVALID_TXN_STATE)
 
-    val expectedAbortableTransaction1ErrorsHigherVersion = topicPartitions.map(_ -> Errors.ABORTABLE_TRANSACTION).toMap
-    val expectedAbortableTransaction2ErrorsHigherVersion = Map(new TopicPartition("foo", 2) -> Errors.ABORTABLE_TRANSACTION)
+    val expectedTransactionAbortableErrorsTxn1HigherVersion = topicPartitions.map(_ -> Errors.TRANSACTION_ABORTABLE).toMap
+    val expectedTransactionAbortableErrorsTxn2HigherVersion = Map(new TopicPartition("foo", 2) -> Errors.TRANSACTION_ABORTABLE)
 
     addTransactionsToVerifyRequestVersion(defaultError)
     receiveResponse(mixedAbortableErrorsResponse)
-    assertEquals(expectedAbortableTransaction1ErrorsLowerVersion, transaction1Errors)
-    assertEquals(expectedAbortableTransaction2ErrorsLowerVersion, transaction2Errors)
+    assertEquals(expectedTransactionAbortableErrorsTxn1LowerVersion, transaction1Errors)
+    assertEquals(expectedTransactionAbortableErrorsTxn2LowerVersion, transaction2Errors)
 
     addTransactionsToVerifyRequestVersion(genericError)
     receiveResponse(mixedAbortableErrorsResponse)
-    assertEquals(expectedAbortableTransaction1ErrorsHigherVersion, transaction1Errors)
-    assertEquals(expectedAbortableTransaction2ErrorsHigherVersion, transaction2Errors)
+    assertEquals(expectedTransactionAbortableErrorsTxn1HigherVersion, transaction1Errors)
+    assertEquals(expectedTransactionAbortableErrorsTxn2HigherVersion, transaction2Errors)
   }
 
   @Test
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
index bc3da1852abdc..6daa5e80b3c74 100644
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
@@ -163,7 +163,7 @@ class AddPartitionsToTxnRequestServerTest extends BaseRequestTest {
 
     val verifyErrors = verifyResponse.errors()
 
-    assertEquals(Collections.singletonMap(transactionalId, Collections.singletonMap(tp0, Errors.ABORTABLE_TRANSACTION)), verifyErrors)
+    assertEquals(Collections.singletonMap(transactionalId, Collections.singletonMap(tp0, Errors.TRANSACTION_ABORTABLE)), verifyErrors)
   }
   
   private def setUpTransactions(transactionalId: String, verifyOnly: Boolean, partitions: Set[TopicPartition]): (Int, AddPartitionsToTxnTransaction) = {

From f40c06690bd3d982ed704411b952e5f9645d5f76 Mon Sep 17 00:00:00 2001
From: Colin Patrick McCabe <cmccabe@apache.org>
Date: Wed, 27 Mar 2024 12:26:36 -0700
Subject: [PATCH 220/258] KAFKA-16428: Fix bug where config change notification
 znode may not get created during migration (#15608)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>, Luke Chen <showuon@gmail.com>
---
 .../main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
index fda6e93c2492d..a1f6e1a112f1c 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkConfigMigrationClient.scala
@@ -244,7 +244,7 @@ class ZkConfigMigrationClient(
         state
       } else if (responses.head.resultCode.equals(Code.OK)) {
         // Write the notification znode if our update was successful
-        zkClient.createConfigChangeNotification(s"$configType/$configName")
+        zkClient.createConfigChangeNotification(s"${configType.get}/$configName")
         state.withMigrationZkVersion(migrationZkVersion)
       } else {
         throw KeeperException.create(responses.head.resultCode, path)

From 4cb6806cb8f1cdbf9f47cb6521127fd3f49fa712 Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Thu, 28 Mar 2024 10:31:22 +0800
Subject: [PATCH 221/258] KAFKA-16232: kafka hangs forever in the starting
 process if the authorizer future is not returned (#15549)

add logs before and after future waiting, to allow admin to know we're waiting for the authorizer future.

Reviewers: Luke Chen <showuon@gmail.com>
---
 core/src/main/scala/kafka/server/KafkaServer.scala | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index fce61e17a225a..3d41f01b70497 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -617,17 +617,22 @@ class KafkaServer(
             }
           }
         }
+
         val enableRequestProcessingFuture = socketServer.enableRequestProcessing(authorizerFutures)
         // Block here until all the authorizer futures are complete
         try {
+          info("Start processing authorizer futures")
           CompletableFuture.allOf(authorizerFutures.values.toSeq: _*).join()
+          info("End processing authorizer futures")
         } catch {
           case t: Throwable => throw new RuntimeException("Received a fatal error while " +
             "waiting for all of the authorizer futures to be completed.", t)
         }
         // Wait for all the SocketServer ports to be open, and the Acceptors to be started.
         try {
+          info("Start processing enable request processing future")
           enableRequestProcessingFuture.join()
+          info("End processing enable request processing future")
         } catch {
           case t: Throwable => throw new RuntimeException("Received a fatal error while " +
             "waiting for the SocketServer Acceptors to be started.", t)

From 4ccbf1634afb063615616b5995ef279a063fbeab Mon Sep 17 00:00:00 2001
From: Alyssa Huang <ahuang@confluent.io>
Date: Thu, 28 Mar 2024 03:22:02 -0700
Subject: [PATCH 222/258] MINOR: Metadata image test improvements (#15373)

Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 .../ProducerIdControlManagerTest.java         |   9 +-
 .../kafka/image/ClientQuotasImageTest.java    |  23 ++--
 .../apache/kafka/image/ClusterImageTest.java  | 104 +++++++++++++++++-
 .../apache/kafka/image/FeaturesImageTest.java |  43 +++++++-
 .../kafka/image/ImageDowngradeTest.java       |  33 ++++++
 .../kafka/image/ProducerIdsImageTest.java     |  12 +-
 .../apache/kafka/image/ScramImageTest.java    |  16 ++-
 7 files changed, 219 insertions(+), 21 deletions(-)

diff --git a/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
index ffe6a7b38caa3..d8c3770f85af0 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
@@ -105,7 +105,14 @@ public void testMonotonic() {
                     .setBrokerEpoch(100)
                     .setNextProducerId(40));
         }, "Producer ID range must only increase");
-        range = producerIdControlManager.generateNextProducerId(1, 100).response();
+        assertThrows(RuntimeException.class, () -> {
+            producerIdControlManager.replay(
+                new ProducerIdsRecord()
+                    .setBrokerId(2)
+                    .setBrokerEpoch(100)
+                    .setNextProducerId(42));
+        }, "Producer ID range must only increase");
+        range = producerIdControlManager.generateNextProducerId(3, 100).response();
         assertEquals(42, range.firstProducerId());
 
         // Gaps in the ID range are okay.
diff --git a/metadata/src/test/java/org/apache/kafka/image/ClientQuotasImageTest.java b/metadata/src/test/java/org/apache/kafka/image/ClientQuotasImageTest.java
index 8d1a5883cc41a..9b6a8fda9d06e 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ClientQuotasImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ClientQuotasImageTest.java
@@ -30,6 +30,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -51,27 +52,31 @@ public class ClientQuotasImageTest {
 
     static {
         Map<ClientQuotaEntity, ClientQuotaImage> entities1 = new HashMap<>();
-        Map<String, String> fooUser = new HashMap<>();
-        fooUser.put(ClientQuotaEntity.USER, "foo");
-        Map<String, Double> fooUserQuotas = new HashMap<>();
-        fooUserQuotas.put(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 123.0);
+        Map<String, String> fooUser = Collections.singletonMap(ClientQuotaEntity.USER, "foo");
+        Map<String, Double> fooUserQuotas = Collections.singletonMap(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 123.0);
         entities1.put(new ClientQuotaEntity(fooUser), new ClientQuotaImage(fooUserQuotas));
         Map<String, String> barUserAndIp = new HashMap<>();
         barUserAndIp.put(ClientQuotaEntity.USER, "bar");
         barUserAndIp.put(ClientQuotaEntity.IP, "127.0.0.1");
-        Map<String, Double> barUserAndIpQuotas = new HashMap<>();
-        barUserAndIpQuotas.put(QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG, 456.0);
-        entities1.put(new ClientQuotaEntity(barUserAndIp),
-            new ClientQuotaImage(barUserAndIpQuotas));
+        Map<String, Double> barUserAndIpQuotas = Collections.singletonMap(QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG, 456.0);
+        entities1.put(new ClientQuotaEntity(barUserAndIp), new ClientQuotaImage(barUserAndIpQuotas));
         IMAGE1 = new ClientQuotasImage(entities1);
 
         DELTA1_RECORDS = new ArrayList<>();
+        // remove quota
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new ClientQuotaRecord().
                 setEntity(Arrays.asList(
                     new EntityData().setEntityType(ClientQuotaEntity.USER).setEntityName("bar"),
                     new EntityData().setEntityType(ClientQuotaEntity.IP).setEntityName("127.0.0.1"))).
                 setKey(QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG).
                 setRemove(true), CLIENT_QUOTA_RECORD.highestSupportedVersion()));
+        // alter quota
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ClientQuotaRecord().
+            setEntity(Arrays.asList(
+                new EntityData().setEntityType(ClientQuotaEntity.USER).setEntityName("foo"))).
+            setKey(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG).
+            setValue(234.0), CLIENT_QUOTA_RECORD.highestSupportedVersion()));
+        // add quota to entity with existing quota
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new ClientQuotaRecord().
             setEntity(Arrays.asList(
                 new EntityData().setEntityType(ClientQuotaEntity.USER).setEntityName("foo"))).
@@ -83,7 +88,7 @@ public class ClientQuotasImageTest {
 
         Map<ClientQuotaEntity, ClientQuotaImage> entities2 = new HashMap<>();
         Map<String, Double> fooUserQuotas2 = new HashMap<>();
-        fooUserQuotas2.put(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 123.0);
+        fooUserQuotas2.put(QuotaConfigs.PRODUCER_BYTE_RATE_OVERRIDE_CONFIG, 234.0);
         fooUserQuotas2.put(QuotaConfigs.CONSUMER_BYTE_RATE_OVERRIDE_CONFIG, 999.0);
         entities2.put(new ClientQuotaEntity(fooUser), new ClientQuotaImage(fooUserQuotas2));
         IMAGE2 = new ClientQuotasImage(entities2);
diff --git a/metadata/src/test/java/org/apache/kafka/image/ClusterImageTest.java b/metadata/src/test/java/org/apache/kafka/image/ClusterImageTest.java
index 7d121aaf46606..820d5a83fa82f 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ClusterImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ClusterImageTest.java
@@ -21,6 +21,11 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord;
 import org.apache.kafka.common.metadata.FenceBrokerRecord;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpoint;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpointCollection;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerFeature;
+import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerFeatureCollection;
 import org.apache.kafka.common.metadata.RegisterControllerRecord;
 import org.apache.kafka.common.metadata.RegisterControllerRecord.ControllerEndpoint;
 import org.apache.kafka.common.metadata.RegisterControllerRecord.ControllerEndpointCollection;
@@ -48,7 +53,9 @@
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static org.apache.kafka.common.metadata.MetadataRecordType.BROKER_REGISTRATION_CHANGE_RECORD;
 import static org.apache.kafka.common.metadata.MetadataRecordType.FENCE_BROKER_RECORD;
+import static org.apache.kafka.common.metadata.MetadataRecordType.REGISTER_BROKER_RECORD;
 import static org.apache.kafka.common.metadata.MetadataRecordType.UNFENCE_BROKER_RECORD;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
@@ -56,13 +63,19 @@
 @Timeout(value = 40)
 public class ClusterImageTest {
 
-    public final static ClusterImage IMAGE1;
+    public static final ClusterImage IMAGE1;
 
     static final List<ApiMessageAndVersion> DELTA1_RECORDS;
 
-    final static ClusterDelta DELTA1;
+    static final ClusterDelta DELTA1;
 
-    final static ClusterImage IMAGE2;
+    static final ClusterImage IMAGE2;
+
+    static final List<ApiMessageAndVersion> DELTA2_RECORDS;
+
+    static final ClusterDelta DELTA2;
+
+    static final ClusterImage IMAGE3;
 
     static {
         Map<Integer, BrokerRegistration> map1 = new HashMap<>();
@@ -88,7 +101,7 @@ public class ClusterImageTest {
             setId(2).
             setEpoch(123).
             setIncarnationId(Uuid.fromString("hr4TVh3YQiu3p16Awkka6w")).
-            setListeners(Arrays.asList(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9093))).
+            setListeners(Arrays.asList(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9094))).
             setSupportedFeatures(Collections.emptyMap()).
             setRack(Optional.of("arack")).
             setFenced(false).
@@ -104,14 +117,18 @@ public class ClusterImageTest {
         IMAGE1 = new ClusterImage(map1, cmap1);
 
         DELTA1_RECORDS = new ArrayList<>();
+        // unfence b0
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new UnfenceBrokerRecord().
             setId(0).setEpoch(1000), UNFENCE_BROKER_RECORD.highestSupportedVersion()));
+        // fence b1
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new FenceBrokerRecord().
             setId(1).setEpoch(1001), FENCE_BROKER_RECORD.highestSupportedVersion()));
+        // mark b0 in controlled shutdown
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
             setBrokerId(0).setBrokerEpoch(1000).setInControlledShutdown(
-            BrokerRegistrationInControlledShutdownChange.IN_CONTROLLED_SHUTDOWN.value()),
-            (short) 0));
+                BrokerRegistrationInControlledShutdownChange.IN_CONTROLLED_SHUTDOWN.value()),
+            BROKER_REGISTRATION_CHANGE_RECORD.highestSupportedVersion()));
+        // unregister b2
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new UnregisterBrokerRecord().
             setBrokerId(2).setBrokerEpoch(123),
             (short) 0));
@@ -160,6 +177,67 @@ public class ClusterImageTest {
                 new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 19093))).
             setSupportedFeatures(Collections.emptyMap()).build());
         IMAGE2 = new ClusterImage(map2, cmap2);
+
+        DELTA2_RECORDS = new ArrayList<>(DELTA1_RECORDS);
+        // fence b0
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new FenceBrokerRecord().
+            setId(0).setEpoch(1000), FENCE_BROKER_RECORD.highestSupportedVersion()));
+        // unfence b1
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new UnfenceBrokerRecord().
+            setId(1).setEpoch(1001), UNFENCE_BROKER_RECORD.highestSupportedVersion()));
+        // mark b0 as not in controlled shutdown
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
+            setBrokerId(0).setBrokerEpoch(1000).setInControlledShutdown(
+                BrokerRegistrationInControlledShutdownChange.NONE.value()),
+            BROKER_REGISTRATION_CHANGE_RECORD.highestSupportedVersion()));
+        // re-register b2
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new RegisterBrokerRecord().
+            setBrokerId(2).setIsMigratingZkBroker(true).setIncarnationId(Uuid.fromString("Am5Yse7GQxaw0b2alM74bP")).
+            setBrokerEpoch(1002).setEndPoints(new BrokerEndpointCollection(
+                Arrays.asList(new BrokerEndpoint().setName("PLAINTEXT").setHost("localhost").
+                    setPort(9094).setSecurityProtocol((short) 0)).iterator())).
+            setFeatures(new BrokerFeatureCollection(
+                Collections.singleton(new BrokerFeature().
+                    setName(MetadataVersion.FEATURE_NAME).
+                    setMinSupportedVersion(MetadataVersion.IBP_3_3_IV3.featureLevel()).
+                    setMaxSupportedVersion(MetadataVersion.IBP_3_6_IV0.featureLevel())).iterator())).
+            setRack("rack3"),
+            REGISTER_BROKER_RECORD.highestSupportedVersion()));
+
+        DELTA2 = new ClusterDelta(IMAGE2);
+        RecordTestUtils.replayAll(DELTA2, DELTA2_RECORDS);
+
+        Map<Integer, BrokerRegistration> map3 = new HashMap<>();
+        map3.put(0, new BrokerRegistration.Builder().
+            setId(0).
+            setEpoch(1000).
+            setIncarnationId(Uuid.fromString("vZKYST0pSA2HO5x_6hoO2Q")).
+            setListeners(Arrays.asList(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9092))).
+            setSupportedFeatures(Collections.singletonMap("foo", VersionRange.of((short) 1, (short) 3))).
+            setRack(Optional.empty()).
+            setFenced(true).
+            setInControlledShutdown(true).build());
+        map3.put(1, new BrokerRegistration.Builder().
+            setId(1).
+            setEpoch(1001).
+            setIncarnationId(Uuid.fromString("U52uRe20RsGI0RvpcTx33Q")).
+            setListeners(Arrays.asList(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9093))).
+            setSupportedFeatures(Collections.singletonMap("foo", VersionRange.of((short) 1, (short) 3))).
+            setRack(Optional.empty()).
+            setFenced(false).
+            setInControlledShutdown(false).build());
+        map3.put(2, new BrokerRegistration.Builder().
+            setId(2).
+            setEpoch(1002).
+            setIncarnationId(Uuid.fromString("Am5Yse7GQxaw0b2alM74bP")).
+            setListeners(Arrays.asList(new Endpoint("PLAINTEXT", SecurityProtocol.PLAINTEXT, "localhost", 9094))).
+            setSupportedFeatures(Collections.singletonMap("metadata.version",
+                VersionRange.of(MetadataVersion.IBP_3_3_IV3.featureLevel(), MetadataVersion.IBP_3_6_IV0.featureLevel()))).
+            setRack(Optional.of("rack3")).
+            setFenced(true).
+            setIsMigratingZkBroker(true).build());
+
+        IMAGE3 = new ClusterImage(map3, cmap2);
     }
 
     @Test
@@ -186,6 +264,20 @@ public void testImage2RoundTrip() {
         testToImage(IMAGE2);
     }
 
+    @Test
+    public void testApplyDelta2() {
+        assertEquals(IMAGE3, DELTA2.apply());
+        // check image2 + delta2 = image3, since records for image2 + delta2 might differ from records from image3
+        List<ApiMessageAndVersion> records = getImageRecords(IMAGE2);
+        records.addAll(DELTA2_RECORDS);
+        testToImage(IMAGE3, records);
+    }
+
+    @Test
+    public void testImage3RoundTrip() {
+        testToImage(IMAGE3);
+    }
+
     private static void testToImage(ClusterImage image) {
         testToImage(image, Optional.empty());
     }
diff --git a/metadata/src/test/java/org/apache/kafka/image/FeaturesImageTest.java b/metadata/src/test/java/org/apache/kafka/image/FeaturesImageTest.java
index bc5149a217286..99e97d64757a7 100644
--- a/metadata/src/test/java/org/apache/kafka/image/FeaturesImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/FeaturesImageTest.java
@@ -45,23 +45,28 @@ public class FeaturesImageTest {
     public final static List<ApiMessageAndVersion> DELTA1_RECORDS;
     final static FeaturesDelta DELTA1;
     final static FeaturesImage IMAGE2;
+    final static List<ApiMessageAndVersion> DELTA2_RECORDS;
+    final static FeaturesDelta DELTA2;
+    final static FeaturesImage IMAGE3;
 
     static {
         Map<String, Short> map1 = new HashMap<>();
         map1.put("foo", (short) 2);
         map1.put("bar", (short) 1);
-        map1.put("baz", (short) 8);
         IMAGE1 = new FeaturesImage(map1, MetadataVersion.latestTesting(), ZkMigrationState.NONE);
 
         DELTA1_RECORDS = new ArrayList<>();
+        // change feature level
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
             setName("foo").setFeatureLevel((short) 3),
             (short) 0));
+        // remove feature
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
             setName("bar").setFeatureLevel((short) 0),
             (short) 0));
+        // add feature
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
-            setName("baz").setFeatureLevel((short) 0),
+            setName("baz").setFeatureLevel((short) 8),
             (short) 0));
 
         DELTA1 = new FeaturesDelta(IMAGE1);
@@ -69,7 +74,27 @@ public class FeaturesImageTest {
 
         Map<String, Short> map2 = new HashMap<>();
         map2.put("foo", (short) 3);
+        map2.put("baz", (short) 8);
         IMAGE2 = new FeaturesImage(map2, MetadataVersion.latestTesting(), ZkMigrationState.NONE);
+
+        DELTA2_RECORDS = new ArrayList<>();
+        // remove all features
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
+            setName("foo").setFeatureLevel((short) 0),
+            (short) 0));
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
+            setName("baz").setFeatureLevel((short) 0),
+            (short) 0));
+        // add feature back with different feature level
+        DELTA2_RECORDS.add(new ApiMessageAndVersion(new FeatureLevelRecord().
+            setName("bar").setFeatureLevel((short) 1),
+            (short) 0));
+
+        DELTA2 = new FeaturesDelta(IMAGE2);
+        RecordTestUtils.replayAll(DELTA2, DELTA2_RECORDS);
+
+        Map<String, Short> map3 = Collections.singletonMap("bar", (short) 1);
+        IMAGE3 = new FeaturesImage(map3, MetadataVersion.latestTesting(), ZkMigrationState.NONE);
     }
 
     @Test
@@ -96,6 +121,20 @@ public void testImage2RoundTrip() {
         testToImage(IMAGE2);
     }
 
+    @Test
+    public void testImage3RoundTrip() {
+        testToImage(IMAGE3);
+    }
+
+    @Test
+    public void testApplyDelta2() {
+        assertEquals(IMAGE3, DELTA2.apply());
+        // check image2 + delta2 = image3, since records for image2 + delta2 might differ from records from image3
+        List<ApiMessageAndVersion> records = getImageRecords(IMAGE2);
+        records.addAll(DELTA2_RECORDS);
+        testToImage(IMAGE3, records);
+    }
+
     private static void testToImage(FeaturesImage image) {
         testToImage(image, Optional.empty());
     }
diff --git a/metadata/src/test/java/org/apache/kafka/image/ImageDowngradeTest.java b/metadata/src/test/java/org/apache/kafka/image/ImageDowngradeTest.java
index 55323c67713d2..4a29779e0acff 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ImageDowngradeTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ImageDowngradeTest.java
@@ -128,6 +128,39 @@ public void testPreControlledShutdownStateVersion() {
                         TEST_RECORDS.get(1)));
     }
 
+    /**
+     * Test downgrading to a MetadataVersion that doesn't support ZK migration.
+     */
+    @Test
+    public void testPreZkMigrationSupportVersion() {
+        writeWithExpectedLosses(MetadataVersion.IBP_3_3_IV3,
+            Arrays.asList(
+                "the isMigratingZkBroker state of one or more brokers"),
+            Arrays.asList(
+                metadataVersionRecord(MetadataVersion.IBP_3_4_IV0),
+                new ApiMessageAndVersion(new RegisterBrokerRecord().
+                    setBrokerId(123).
+                    setIncarnationId(Uuid.fromString("XgjKo16hRWeWrTui0iR5Nw")).
+                    setBrokerEpoch(456).
+                    setRack(null).
+                    setFenced(false).
+                    setInControlledShutdown(true).
+                    setIsMigratingZkBroker(true), (short) 2),
+                TEST_RECORDS.get(0),
+                TEST_RECORDS.get(1)),
+            Arrays.asList(
+                metadataVersionRecord(MetadataVersion.IBP_3_3_IV3),
+                new ApiMessageAndVersion(new RegisterBrokerRecord().
+                    setBrokerId(123).
+                    setIncarnationId(Uuid.fromString("XgjKo16hRWeWrTui0iR5Nw")).
+                    setBrokerEpoch(456).
+                    setRack(null).
+                    setFenced(false).
+                    setInControlledShutdown(true), (short) 1),
+                TEST_RECORDS.get(0),
+                TEST_RECORDS.get(1)));
+    }
+
     @Test
     void testDirectoryAssignmentState() {
         MetadataVersion outputMetadataVersion = MetadataVersion.IBP_3_7_IV0;
diff --git a/metadata/src/test/java/org/apache/kafka/image/ProducerIdsImageTest.java b/metadata/src/test/java/org/apache/kafka/image/ProducerIdsImageTest.java
index 738582fc108cd..32d42efda7a8c 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ProducerIdsImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ProducerIdsImageTest.java
@@ -53,12 +53,20 @@ public class ProducerIdsImageTest {
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new ProducerIdsRecord().
             setBrokerId(3).
             setBrokerEpoch(100).
-            setNextProducerId(789), (short) 0));
+            setNextProducerId(780), (short) 0));
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ProducerIdsRecord().
+            setBrokerId(3).
+            setBrokerEpoch(100).
+            setNextProducerId(785), (short) 0));
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new ProducerIdsRecord().
+            setBrokerId(2).
+            setBrokerEpoch(100).
+            setNextProducerId(800), (short) 0));
 
         DELTA1 = new ProducerIdsDelta(IMAGE1);
         RecordTestUtils.replayAll(DELTA1, DELTA1_RECORDS);
 
-        IMAGE2 = new ProducerIdsImage(789);
+        IMAGE2 = new ProducerIdsImage(800);
     }
 
     @Test
diff --git a/metadata/src/test/java/org/apache/kafka/image/ScramImageTest.java b/metadata/src/test/java/org/apache/kafka/image/ScramImageTest.java
index 038a5c956c34b..26f860579e07d 100644
--- a/metadata/src/test/java/org/apache/kafka/image/ScramImageTest.java
+++ b/metadata/src/test/java/org/apache/kafka/image/ScramImageTest.java
@@ -85,10 +85,15 @@ static ScramCredentialData randomScramCredentialData(Random random) {
         IMAGE1 = new ScramImage(image1mechanisms);
 
         DELTA1_RECORDS = new ArrayList<>();
+        // remove all sha512 credentials
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new RemoveUserScramCredentialRecord().
+            setName("alpha").
+            setMechanism(SCRAM_SHA_512.type()), (short) 0));
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new RemoveUserScramCredentialRecord().
             setName("gamma").
             setMechanism(SCRAM_SHA_512.type()), (short) 0));
         ScramCredentialData secondAlpha256Credential = randomScramCredentialData(random);
+        // add sha256 credential
         DELTA1_RECORDS.add(new ApiMessageAndVersion(new UserScramCredentialRecord().
                 setName("alpha").
                 setMechanism(SCRAM_SHA_256.type()).
@@ -96,6 +101,15 @@ static ScramCredentialData randomScramCredentialData(Random random) {
                 setStoredKey(secondAlpha256Credential.storedKey()).
                 setServerKey(secondAlpha256Credential.serverKey()).
                 setIterations(secondAlpha256Credential.iterations()), (short) 0));
+        // add sha512 credential re-using name
+        ScramCredentialData secondAlpha512Credential = randomScramCredentialData(random);
+        DELTA1_RECORDS.add(new ApiMessageAndVersion(new UserScramCredentialRecord().
+            setName("alpha").
+            setMechanism(SCRAM_SHA_512.type()).
+            setSalt(secondAlpha512Credential.salt()).
+            setStoredKey(secondAlpha512Credential.storedKey()).
+            setServerKey(secondAlpha512Credential.serverKey()).
+            setIterations(secondAlpha512Credential.iterations()), (short) 0));
         DELTA1 = new ScramDelta(IMAGE1);
         RecordTestUtils.replayAll(DELTA1, DELTA1_RECORDS);
 
@@ -107,7 +121,7 @@ static ScramCredentialData randomScramCredentialData(Random random) {
         image2mechanisms.put(SCRAM_SHA_256, image2sha256);
 
         Map<String, ScramCredentialData> image2sha512 = new HashMap<>();
-        image2sha512.put("alpha", image1sha512.get("alpha"));
+        image2sha512.put("alpha", secondAlpha512Credential);
         image2mechanisms.put(SCRAM_SHA_512, image2sha512);
 
         IMAGE2 = new ScramImage(image2mechanisms);

From b411ac0c700fcb74639502cf38181de6096b6404 Mon Sep 17 00:00:00 2001
From: Lianet Magrans <98415067+lianetm@users.noreply.github.com>
Date: Thu, 28 Mar 2024 14:15:56 +0100
Subject: [PATCH 223/258] KAFKA-16406 [2] : Split consumer commit tests
 (#15612)

Follow-up to #15535, splitting consumer integration tests defined in the long-running PlainTextConsumerTest. This PR extracts the tests that directly relate to committing offsets. No changes in logic.

Reviewers: Lucas Brutschy <lbrutschy@confluent.io>
---
 .../api/PlaintextConsumerCommitTest.scala     | 320 ++++++++++++++++++
 .../kafka/api/PlaintextConsumerTest.scala     | 279 +--------------
 2 files changed, 321 insertions(+), 278 deletions(-)
 create mode 100644 core/src/test/scala/integration/kafka/api/PlaintextConsumerCommitTest.scala

diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerCommitTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerCommitTest.scala
new file mode 100644
index 0000000000000..db5ef480d1493
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerCommitTest.scala
@@ -0,0 +1,320 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+  * specific language governing permissions and limitations under the License.
+  */
+package kafka.api
+
+import kafka.utils.{TestInfoUtils, TestUtils}
+import org.apache.kafka.clients.consumer._
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializer}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.test.MockConsumerInterceptor
+import org.junit.jupiter.api.Assertions._
+import org.junit.jupiter.api.Timeout
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.{Arguments, MethodSource}
+
+import java.time.Duration
+import java.util
+import java.util.Optional
+import java.util.stream.Stream
+import scala.jdk.CollectionConverters._
+
+/**
+ * Integration tests for the consumer that covers the logic related to committing offsets.
+ */
+@Timeout(600)
+class PlaintextConsumerCommitTest extends AbstractConsumerTest {
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAutoCommitOnClose(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
+    val consumer = createConsumer()
+
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, Set(tp, tp2))
+
+    // should auto-commit sought positions before closing
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
+    consumer.close()
+
+    // now we should see the committed positions from another consumer
+    val anotherConsumer = createConsumer()
+    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAutoCommitOnCloseAfterWakeup(quorum: String, groupProtocol: String): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
+    val consumer = createConsumer()
+
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, Set(tp, tp2))
+
+    // should auto-commit sought positions before closing
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
+
+    // wakeup the consumer before closing to simulate trying to break a poll
+    // loop from another thread
+    consumer.wakeup()
+    consumer.close()
+
+    // now we should see the committed positions from another consumer
+    val anotherConsumer = createConsumer()
+    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testCommitMetadata(quorum: String, groupProtocol: String): Unit = {
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+
+    // sync commit
+    val syncMetadata = new OffsetAndMetadata(5, Optional.of(15), "foo")
+    consumer.commitSync(Map((tp, syncMetadata)).asJava)
+    assertEquals(syncMetadata, consumer.committed(Set(tp).asJava).get(tp))
+
+    // async commit
+    val asyncMetadata = new OffsetAndMetadata(10, "bar")
+    sendAndAwaitAsyncCommit(consumer, Some(Map(tp -> asyncMetadata)))
+    assertEquals(asyncMetadata, consumer.committed(Set(tp).asJava).get(tp))
+
+    // handle null metadata
+    val nullMetadata = new OffsetAndMetadata(5, null)
+    consumer.commitSync(Map(tp -> nullMetadata).asJava)
+    assertEquals(nullMetadata, consumer.committed(Set(tp).asJava).get(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAsyncCommit(quorum: String, groupProtocol: String): Unit = {
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+
+    val callback = new CountConsumerCommitCallback
+    val count = 5
+
+    for (i <- 1 to count)
+      consumer.commitAsync(Map(tp -> new OffsetAndMetadata(i)).asJava, callback)
+
+    TestUtils.pollUntilTrue(consumer, () => callback.successCount >= count || callback.lastError.isDefined,
+      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
+
+    assertEquals(None, callback.lastError)
+    assertEquals(count, callback.successCount)
+    assertEquals(new OffsetAndMetadata(count), consumer.committed(Set(tp).asJava).get(tp))
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAutoCommitIntercept(quorum: String, groupProtocol: String): Unit = {
+    val topic2 = "topic2"
+    createTopic(topic2, 2, brokerCount)
+
+    // produce records
+    val numRecords = 100
+    val testProducer = createProducer(keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
+    (0 until numRecords).map { i =>
+      testProducer.send(new ProducerRecord(tp.topic(), tp.partition(), s"key $i", s"value $i"))
+    }.foreach(_.get)
+
+    // create consumer with interceptor
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
+    this.consumerConfig.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockConsumerInterceptor")
+    val testConsumer = createConsumer(keyDeserializer = new StringDeserializer, valueDeserializer = new StringDeserializer)
+    val rebalanceListener = new ConsumerRebalanceListener {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
+        // keep partitions paused in this test so that we can verify the commits based on specific seeks
+        testConsumer.pause(partitions)
+      }
+
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {}
+    }
+    changeConsumerSubscriptionAndValidateAssignment(testConsumer, List(topic), Set(tp, tp2), rebalanceListener)
+    testConsumer.seek(tp, 10)
+    testConsumer.seek(tp2, 20)
+
+    // change subscription to trigger rebalance
+    val commitCountBeforeRebalance = MockConsumerInterceptor.ON_COMMIT_COUNT.intValue()
+    changeConsumerSubscriptionAndValidateAssignment(testConsumer,
+      List(topic, topic2),
+      Set(tp, tp2, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)),
+      rebalanceListener)
+
+    // after rebalancing, we should have reset to the committed positions
+    assertEquals(10, testConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(20, testConsumer.committed(Set(tp2).asJava).get(tp2).offset)
+
+    // In both CLASSIC and CONSUMER protocols, interceptors are executed in poll and close.
+    // However, in the CONSUMER protocol, the assignment may be changed outside of a poll, so
+    // we need to poll once to ensure the interceptor is called.
+    if (groupProtocol.toUpperCase == GroupProtocol.CONSUMER.name) {
+      testConsumer.poll(Duration.ZERO);
+    }
+
+    assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeRebalance)
+
+    // verify commits are intercepted on close
+    val commitCountBeforeClose = MockConsumerInterceptor.ON_COMMIT_COUNT.intValue()
+    testConsumer.close()
+    assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeClose)
+    testProducer.close()
+
+    // cleanup
+    MockConsumerInterceptor.resetCounters()
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testCommitSpecifiedOffsets(quorum: String, groupProtocol: String): Unit = {
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 5, tp)
+    sendRecords(producer, numRecords = 7, tp2)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp, tp2).asJava)
+
+    val pos1 = consumer.position(tp)
+    val pos2 = consumer.position(tp2)
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava)
+    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertNull(consumer.committed(Set(tp2).asJava).get(tp2))
+
+    // Positions should not change
+    assertEquals(pos1, consumer.position(tp))
+    assertEquals(pos2, consumer.position(tp2))
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(5L))).asJava)
+    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(5, consumer.committed(Set(tp2).asJava).get(tp2).offset)
+
+    // Using async should pick up the committed changes after commit completes
+    sendAndAwaitAsyncCommit(consumer, Some(Map(tp2 -> new OffsetAndMetadata(7L))))
+    assertEquals(7, consumer.committed(Set(tp2).asJava).get(tp2).offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testAutoCommitOnRebalance(quorum: String, groupProtocol: String): Unit = {
+    val topic2 = "topic2"
+    createTopic(topic2, 2, brokerCount)
+
+    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
+    val consumer = createConsumer()
+
+    val numRecords = 10000
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    val rebalanceListener = new ConsumerRebalanceListener {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
+        // keep partitions paused in this test so that we can verify the commits based on specific seeks
+        consumer.pause(partitions)
+      }
+
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {}
+    }
+
+    consumer.subscribe(List(topic).asJava, rebalanceListener)
+
+    awaitAssignment(consumer, Set(tp, tp2))
+
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
+
+    // change subscription to trigger rebalance
+    consumer.subscribe(List(topic, topic2).asJava, rebalanceListener)
+
+    val newAssignment = Set(tp, tp2, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1))
+    awaitAssignment(consumer, newAssignment)
+
+    // after rebalancing, we should have reset to the committed positions
+    assertEquals(300, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, consumer.committed(Set(tp2).asJava).get(tp2).offset)
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testSubscribeAndCommitSync(quorum: String, groupProtocol: String): Unit = {
+    // This test ensure that the member ID is propagated from the group coordinator when the
+    // assignment is received into a subsequent offset commit
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment.size)
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, Set(tp, tp2))
+
+    consumer.seek(tp, 0)
+
+    consumer.commitSync()
+  }
+
+  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
+  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
+  def testPositionAndCommit(quorum: String, groupProtocol: String): Unit = {
+    val producer = createProducer()
+    var startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords = 5, tp, startingTimestamp = startingTimestamp)
+
+    val topicPartition = new TopicPartition(topic, 15)
+    val consumer = createConsumer()
+    assertNull(consumer.committed(Set(topicPartition).asJava).get(topicPartition))
+
+    // position() on a partition that we aren't subscribed to throws an exception
+    assertThrows(classOf[IllegalStateException], () => consumer.position(topicPartition))
+
+    consumer.assign(List(tp).asJava)
+
+    assertEquals(0L, consumer.position(tp), "position() on a partition that we are subscribed to should reset the offset")
+    consumer.commitSync()
+    assertEquals(0L, consumer.committed(Set(tp).asJava).get(tp).offset)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 0, startingTimestamp = startingTimestamp)
+    assertEquals(5L, consumer.position(tp), "After consuming 5 records, position should be 5")
+    consumer.commitSync()
+    assertEquals(5L, consumer.committed(Set(tp).asJava).get(tp).offset, "Committed offset should be returned")
+
+    startingTimestamp = System.currentTimeMillis()
+    sendRecords(producer, numRecords = 1, tp, startingTimestamp = startingTimestamp)
+
+    // another consumer in the same group should get the same position
+    val otherConsumer = createConsumer()
+    otherConsumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = otherConsumer, numRecords = 1, startingOffset = 5, startingTimestamp = startingTimestamp)
+  }
+
+  def changeConsumerSubscriptionAndValidateAssignment[K, V](consumer: Consumer[K, V],
+                                                            topicsToSubscribe: List[String],
+                                                            expectedAssignment: Set[TopicPartition],
+                                                            rebalanceListener: ConsumerRebalanceListener): Unit = {
+    consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
+    awaitAssignment(consumer, expectedAssignment)
+  }
+}
+
+object PlaintextConsumerCommitTest {
+
+  def getTestQuorumAndGroupProtocolParametersAll: Stream[Arguments] =
+    BaseConsumerTest.getTestQuorumAndGroupProtocolParametersAll()
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 8784e5b43fc23..77ba47fa08fdf 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -15,7 +15,7 @@ package kafka.api
 import java.time.Duration
 import java.util
 import java.util.Arrays.asList
-import java.util.{Locale, Optional, Properties}
+import java.util.{Locale, Properties}
 import kafka.server.{KafkaBroker, QuotaType}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin.{NewPartitions, NewTopic}
@@ -135,58 +135,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     testHeadersSerializeDeserialize(extendedSerializer, extendedDeserializer)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAutoCommitOnClose(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer = createConsumer()
-
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp)
-
-    consumer.subscribe(List(topic).asJava)
-    awaitAssignment(consumer, Set(tp, tp2))
-
-    // should auto-commit sought positions before closing
-    consumer.seek(tp, 300)
-    consumer.seek(tp2, 500)
-    consumer.close()
-
-    // now we should see the committed positions from another consumer
-    val anotherConsumer = createConsumer()
-    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAutoCommitOnCloseAfterWakeup(quorum: String, groupProtocol: String): Unit = {
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer = createConsumer()
-
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp)
-
-    consumer.subscribe(List(topic).asJava)
-    awaitAssignment(consumer, Set(tp, tp2))
-
-    // should auto-commit sought positions before closing
-    consumer.seek(tp, 300)
-    consumer.seek(tp2, 500)
-
-    // wakeup the consumer before closing to simulate trying to break a poll
-    // loop from another thread
-    consumer.wakeup()
-    consumer.close()
-
-    // now we should see the committed positions from another consumer
-    val anotherConsumer = createConsumer()
-    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testAutoOffsetReset(quorum: String, groupProtocol: String): Unit = {
@@ -211,48 +159,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0, startingTimestamp = startingTimestamp)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testCommitMetadata(quorum: String, groupProtocol: String): Unit = {
-    val consumer = createConsumer()
-    consumer.assign(List(tp).asJava)
-
-    // sync commit
-    val syncMetadata = new OffsetAndMetadata(5, Optional.of(15), "foo")
-    consumer.commitSync(Map((tp, syncMetadata)).asJava)
-    assertEquals(syncMetadata, consumer.committed(Set(tp).asJava).get(tp))
-
-    // async commit
-    val asyncMetadata = new OffsetAndMetadata(10, "bar")
-    sendAndAwaitAsyncCommit(consumer, Some(Map(tp -> asyncMetadata)))
-    assertEquals(asyncMetadata, consumer.committed(Set(tp).asJava).get(tp))
-
-    // handle null metadata
-    val nullMetadata = new OffsetAndMetadata(5, null)
-    consumer.commitSync(Map(tp -> nullMetadata).asJava)
-    assertEquals(nullMetadata, consumer.committed(Set(tp).asJava).get(tp))
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAsyncCommit(quorum: String, groupProtocol: String): Unit = {
-    val consumer = createConsumer()
-    consumer.assign(List(tp).asJava)
-
-    val callback = new CountConsumerCommitCallback
-    val count = 5
-
-    for (i <- 1 to count)
-      consumer.commitAsync(Map(tp -> new OffsetAndMetadata(i)).asJava, callback)
-
-    TestUtils.pollUntilTrue(consumer, () => callback.successCount >= count || callback.lastError.isDefined,
-      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
-
-    assertEquals(None, callback.lastError)
-    assertEquals(count, callback.successCount)
-    assertEquals(new OffsetAndMetadata(count), consumer.committed(Set(tp).asJava).get(tp))
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testPartitionsFor(quorum: String, groupProtocol: String): Unit = {
@@ -338,39 +244,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     producer.close()
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testPositionAndCommit(quorum: String, groupProtocol: String): Unit = {
-    val producer = createProducer()
-    var startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords = 5, tp, startingTimestamp = startingTimestamp)
-
-    val topicPartition = new TopicPartition(topic, 15)
-    val consumer = createConsumer()
-    assertNull(consumer.committed(Set(topicPartition).asJava).get(topicPartition))
-
-    // position() on a partition that we aren't subscribed to throws an exception
-    assertThrows(classOf[IllegalStateException], () => consumer.position(topicPartition))
-
-    consumer.assign(List(tp).asJava)
-
-    assertEquals(0L, consumer.position(tp), "position() on a partition that we are subscribed to should reset the offset")
-    consumer.commitSync()
-    assertEquals(0L, consumer.committed(Set(tp).asJava).get(tp).offset)
-    consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 0, startingTimestamp = startingTimestamp)
-    assertEquals(5L, consumer.position(tp), "After consuming 5 records, position should be 5")
-    consumer.commitSync()
-    assertEquals(5L, consumer.committed(Set(tp).asJava).get(tp).offset, "Committed offset should be returned")
-
-    startingTimestamp = System.currentTimeMillis()
-    sendRecords(producer, numRecords = 1, tp, startingTimestamp = startingTimestamp)
-
-    // another consumer in the same group should get the same position
-    val otherConsumer = createConsumer()
-    otherConsumer.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = otherConsumer, numRecords = 1, startingOffset = 5, startingTimestamp = startingTimestamp)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testPartitionPauseAndResume(quorum: String, groupProtocol: String): Unit = {
@@ -450,65 +323,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     MockProducerInterceptor.resetCounters()
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAutoCommitIntercept(quorum: String, groupProtocol: String): Unit = {
-    val topic2 = "topic2"
-    createTopic(topic2, 2, brokerCount)
-
-    // produce records
-    val numRecords = 100
-    val testProducer = createProducer(keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
-    (0 until numRecords).map { i =>
-      testProducer.send(new ProducerRecord(tp.topic(), tp.partition(), s"key $i", s"value $i"))
-    }.foreach(_.get)
-
-    // create consumer with interceptor
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    this.consumerConfig.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockConsumerInterceptor")
-    val testConsumer = createConsumer(keyDeserializer = new StringDeserializer, valueDeserializer = new StringDeserializer)
-    val rebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
-        // keep partitions paused in this test so that we can verify the commits based on specific seeks
-        testConsumer.pause(partitions)
-      }
-
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {}
-    }
-    changeConsumerSubscriptionAndValidateAssignment(testConsumer, List(topic), Set(tp, tp2), rebalanceListener)
-    testConsumer.seek(tp, 10)
-    testConsumer.seek(tp2, 20)
-
-    // change subscription to trigger rebalance
-    val commitCountBeforeRebalance = MockConsumerInterceptor.ON_COMMIT_COUNT.intValue()
-    changeConsumerSubscriptionAndValidateAssignment(testConsumer,
-                                                    List(topic, topic2),
-                                                    Set(tp, tp2, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)),
-                                                    rebalanceListener)
-
-    // after rebalancing, we should have reset to the committed positions
-    assertEquals(10, testConsumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(20, testConsumer.committed(Set(tp2).asJava).get(tp2).offset)
-
-    // In both CLASSIC and CONSUMER protocols, interceptors are executed in poll and close.
-    // However, in the CONSUMER protocol, the assignment may be changed outside of a poll, so
-    // we need to poll once to ensure the interceptor is called.
-    if (groupProtocol.toUpperCase == GroupProtocol.CONSUMER.name) {
-      testConsumer.poll(Duration.ZERO);
-    }
-
-    assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeRebalance)
-
-    // verify commits are intercepted on close
-    val commitCountBeforeClose = MockConsumerInterceptor.ON_COMMIT_COUNT.intValue()
-    testConsumer.close()
-    assertTrue(MockConsumerInterceptor.ON_COMMIT_COUNT.intValue() > commitCountBeforeClose)
-    testProducer.close()
-
-    // cleanup
-    MockConsumerInterceptor.resetCounters()
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testInterceptorsWithWrongKeyValue(quorum: String, groupProtocol: String): Unit = {
@@ -626,74 +440,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     consumeAndVerifyRecords(consumer = consumer, numRecords = 0, startingOffset = 5, startingTimestamp = startingTimestamp)
   }
 
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testCommitSpecifiedOffsets(quorum: String, groupProtocol: String): Unit = {
-    val producer = createProducer()
-    sendRecords(producer, numRecords = 5, tp)
-    sendRecords(producer, numRecords = 7, tp2)
-
-    val consumer = createConsumer()
-    consumer.assign(List(tp, tp2).asJava)
-
-    val pos1 = consumer.position(tp)
-    val pos2 = consumer.position(tp2)
-    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava)
-    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
-    assertNull(consumer.committed(Set(tp2).asJava).get(tp2))
-
-    // Positions should not change
-    assertEquals(pos1, consumer.position(tp))
-    assertEquals(pos2, consumer.position(tp2))
-    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(5L))).asJava)
-    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(5, consumer.committed(Set(tp2).asJava).get(tp2).offset)
-
-    // Using async should pick up the committed changes after commit completes
-    sendAndAwaitAsyncCommit(consumer, Some(Map(tp2 -> new OffsetAndMetadata(7L))))
-    assertEquals(7, consumer.committed(Set(tp2).asJava).get(tp2).offset)
-  }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testAutoCommitOnRebalance(quorum: String, groupProtocol: String): Unit = {
-    val topic2 = "topic2"
-    createTopic(topic2, 2, brokerCount)
-
-    this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer = createConsumer()
-
-    val numRecords = 10000
-    val producer = createProducer()
-    sendRecords(producer, numRecords, tp)
-
-    val rebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
-        // keep partitions paused in this test so that we can verify the commits based on specific seeks
-        consumer.pause(partitions)
-      }
-
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {}
-    }
-
-    consumer.subscribe(List(topic).asJava, rebalanceListener)
-
-    awaitAssignment(consumer, Set(tp, tp2))
-
-    consumer.seek(tp, 300)
-    consumer.seek(tp2, 500)
-
-    // change subscription to trigger rebalance
-    consumer.subscribe(List(topic, topic2).asJava, rebalanceListener)
-
-    val newAssignment = Set(tp, tp2, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1))
-    awaitAssignment(consumer, newAssignment)
-
-    // after rebalancing, we should have reset to the committed positions
-    assertEquals(300, consumer.committed(Set(tp).asJava).get(tp).offset)
-    assertEquals(500, consumer.committed(Set(tp2).asJava).get(tp2).offset)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testPerPartitionLeadMetricsCleanUpWithSubscribe(quorum: String, groupProtocol: String): Unit = {
@@ -893,14 +639,6 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     brokers.foreach(assertNoMetric(_, "throttle-time", QuotaType.Request, consumerClientId))
   }
 
-  def changeConsumerSubscriptionAndValidateAssignment[K, V](consumer: Consumer[K, V],
-                                                            topicsToSubscribe: List[String],
-                                                            expectedAssignment: Set[TopicPartition],
-                                                            rebalanceListener: ConsumerRebalanceListener): Unit = {
-    consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
-    awaitAssignment(consumer, expectedAssignment)
-  }
-
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
   @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
   def testConsumingWithNullGroupId(quorum: String, groupProtocol: String): Unit = {
@@ -1067,19 +805,4 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     consumer2.close()
   }
-
-  @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumAndGroupProtocolNames)
-  @MethodSource(Array("getTestQuorumAndGroupProtocolParametersAll"))
-  def testSubscribeAndCommitSync(quorum: String, groupProtocol: String): Unit = {
-    // This test ensure that the member ID is propagated from the group coordinator when the
-    // assignment is received into a subsequent offset commit
-    val consumer = createConsumer()
-    assertEquals(0, consumer.assignment.size)
-    consumer.subscribe(List(topic).asJava)
-    awaitAssignment(consumer, Set(tp, tp2))
-
-    consumer.seek(tp, 0)
-
-    consumer.commitSync()
-  }
 }

From 355873aa54439790d37a6b7c1c9d0c072511e47a Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Thu, 28 Mar 2024 17:52:34 +0300
Subject: [PATCH 224/258] MINOR: Use CONFIG suffix in ZkConfigs (#15614)

Reviewers: Mickael Maison <mickael.maison@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>, Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
Co-authored-by: n.izhikov <n.izhikov@vk.team>
---
 .../util/clusters/EmbeddedKafkaCluster.java   |   2 +-
 .../main/scala/kafka/admin/AclCommand.scala   |   2 +-
 .../kafka/admin/ZkSecurityMigrator.scala      |   2 +-
 .../security/authorizer/AclAuthorizer.scala   |   4 +-
 .../main/scala/kafka/server/KafkaConfig.scala |  90 ++++++-------
 .../main/scala/kafka/server/KafkaServer.scala |  28 ++---
 .../main/scala/kafka/zk/KafkaZkClient.scala   |   6 +-
 .../DescribeAuthorizedOperationsTest.scala    |   2 +-
 .../kafka/api/EndToEndAuthorizationTest.scala |   2 +-
 .../integration/kafka/api/MetricsTest.scala   |   2 +-
 .../api/PlaintextAdminIntegrationTest.scala   |   2 +-
 .../api/SaslMultiMechanismConsumerTest.scala  |   2 +-
 .../api/SaslPlainPlaintextConsumerTest.scala  |   2 +-
 .../api/SaslSslAdminIntegrationTest.scala     |   2 +-
 .../kafka/api/SaslSslConsumerTest.scala       |   2 +-
 .../kafka/api/SslAdminIntegrationTest.scala   |   2 +-
 .../DynamicBrokerReconfigurationTest.scala    |   2 +-
 .../KafkaServerKRaftRegistrationTest.scala    |   4 +-
 ...nersWithSameSecurityProtocolBaseTest.scala |   2 +-
 .../kafka/zk/ZkMigrationIntegrationTest.scala |  14 +--
 .../scala/unit/kafka/KafkaConfigTest.scala    |  28 ++---
 .../ControllerChannelManagerTest.scala        |   2 +-
 .../security/authorizer/AuthorizerTest.scala  | 118 +++++++++---------
 .../server/DynamicBrokerConfigTest.scala      |   4 +-
 .../unit/kafka/server/KafkaConfigTest.scala   |  78 ++++++------
 .../unit/kafka/server/KafkaServerTest.scala   |  34 ++---
 .../kafka/server/ServerShutdownTest.scala     |   4 +-
 .../scala/unit/kafka/server/ServerTest.scala  |   2 +-
 .../scala/unit/kafka/utils/TestUtils.scala    |   4 +-
 .../unit/kafka/zk/KafkaZkClientTest.scala     |   2 +-
 .../zk/migration/ZkMigrationTestHarness.scala |   2 +-
 .../kafka/zookeeper/ZooKeeperClientTest.scala |   2 +-
 .../metadata/MetadataRequestBenchmark.java    |   2 +-
 .../apache/kafka/server/config/ZkConfigs.java | 100 +++++++--------
 .../utils/EmbeddedKafkaCluster.java           |   2 +-
 .../integration/utils/KafkaEmbedded.java      |   2 +-
 36 files changed, 280 insertions(+), 280 deletions(-)

diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
index 8a06c8e0b0a39..9f1f85ee2933f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
@@ -155,7 +155,7 @@ public void start() {
     }
 
     private void doStart() {
-        brokerConfig.put(ZkConfigs.ZK_CONNECT_PROP, zKConnectString());
+        brokerConfig.put(ZkConfigs.ZK_CONNECT_CONFIG, zKConnectString());
 
         putIfAbsent(brokerConfig, KafkaConfig.DeleteTopicEnableProp(), true);
         putIfAbsent(brokerConfig, KafkaConfig.GroupInitialRebalanceDelayMsProp(), 0);
diff --git a/core/src/main/scala/kafka/admin/AclCommand.scala b/core/src/main/scala/kafka/admin/AclCommand.scala
index 7120d0d8c3439..0415598b4ed5b 100644
--- a/core/src/main/scala/kafka/admin/AclCommand.scala
+++ b/core/src/main/scala/kafka/admin/AclCommand.scala
@@ -200,7 +200,7 @@ object AclCommand extends Logging {
       // We will default the value of zookeeper.set.acl to true or false based on whether SASL is configured,
       // but if SASL is not configured and zookeeper.set.acl is supposed to be true due to mutual certificate authentication
       // then it will be up to the user to explicitly specify zookeeper.set.acl=true in the authorizer-properties.
-      val defaultProps = Map(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP -> JaasUtils.isZkSaslEnabled)
+      val defaultProps = Map(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG -> JaasUtils.isZkSaslEnabled)
       val authorizerPropertiesWithoutTls =
         if (opts.options.has(opts.authorizerPropertiesOpt)) {
           val authorizerProperties = opts.options.valuesOf(opts.authorizerPropertiesOpt)
diff --git a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
index dd07f522e77dc..4dc334728df96 100644
--- a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
+++ b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
@@ -82,7 +82,7 @@ object ZkSecurityMigrator extends Logging {
     if (jaasFile == null && !tlsClientAuthEnabled) {
       val errorMsg = s"No JAAS configuration file has been specified and no TLS client certificate has been specified. Please make sure that you set " +
         s"the system property ${JaasUtils.JAVA_LOGIN_CONFIG_PARAM} or provide a ZooKeeper client TLS configuration via --$tlsConfigFileOption <filename> " +
-        s"identifying at least ${ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP}, and ${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP}"
+        s"identifying at least ${ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG}, and ${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG}"
       System.err.println("ERROR: %s".format(errorMsg))
       throw new IllegalArgumentException("Incorrect configuration")
     }
diff --git a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
index dc307ca702fa8..fc2d30f8e8cb1 100644
--- a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -97,7 +97,7 @@ object AclAuthorizer {
   }
 
   private[authorizer] def zkClientConfigFromKafkaConfigAndMap(kafkaConfig: KafkaConfig, configMap: mutable.Map[String, _<:Any]): ZKClientConfig = {
-    val zkSslClientEnable = configMap.get(AclAuthorizer.configPrefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP).
+    val zkSslClientEnable = configMap.get(AclAuthorizer.configPrefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG).
       map(_.toString.trim).getOrElse(kafkaConfig.zkSslClientEnable.toString).toBoolean
     if (!zkSslClientEnable)
       new ZKClientConfig
@@ -109,7 +109,7 @@ object AclAuthorizer {
       ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.forKeyValue { (kafkaProp, sysProp) =>
         configMap.get(AclAuthorizer.configPrefix + kafkaProp).foreach { prefixedValue =>
           zkClientConfig.setProperty(sysProp,
-            if (kafkaProp == ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP)
+            if (kafkaProp == ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG)
               (prefixedValue.toString.trim.toUpperCase == "HTTPS").toString
             else
               prefixedValue.toString.trim)
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index cec5f5649f628..c6d6ad83379c9 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -74,8 +74,8 @@ object KafkaConfig {
   private[kafka] def setZooKeeperClientProperty(clientConfig: ZKClientConfig, kafkaPropName: String, kafkaPropValue: Any): Unit = {
     clientConfig.setProperty(ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(kafkaPropName),
       kafkaPropName match {
-        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => (kafkaPropValue.toString.toUpperCase == "HTTPS").toString
-        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => kafkaPropValue match {
+        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => (kafkaPropValue.toString.toUpperCase == "HTTPS").toString
+        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG | ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG => kafkaPropValue match {
           case list: java.util.List[_] => list.asScala.mkString(",")
           case _ => kafkaPropValue.toString
         }
@@ -86,9 +86,9 @@ object KafkaConfig {
   // For ZooKeeper TLS client authentication to be enabled the client must (at a minimum) configure itself as using TLS
   // with both a client connection socket and a key store location explicitly set.
   private[kafka] def zkTlsClientAuthEnabled(zkClientConfig: ZKClientConfig): Boolean = {
-    zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP).contains("true") &&
-      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP).isDefined &&
-      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP).isDefined
+    zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG).contains("true") &&
+      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG).isDefined &&
+      zooKeeperClientProperty(zkClientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG).isDefined
   }
 
   /** ********* General Configuration ***********/
@@ -860,25 +860,25 @@ object KafkaConfig {
     new ConfigDef()
 
       /** ********* Zookeeper Configuration ***********/
-      .define(ZkConfigs.ZK_CONNECT_PROP, STRING, null, HIGH, ZkConfigs.ZK_CONNECT_DOC)
-      .define(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, INT, ZkConfigs.ZK_SESSION_TIMEOUT_MS, HIGH, ZkConfigs.ZK_SESSION_TIMEOUT_MS_DOC)
-      .define(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, INT, null, HIGH, ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_DOC)
-      .define(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, BOOLEAN, ZkConfigs.ZK_ENABLE_SECURE_ACLS, HIGH, ZkConfigs.ZK_ENABLE_SECURE_ACLS_DOC)
-      .define(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP, INT, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS, atLeast(1), HIGH, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_DOC)
-      .define(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_CLIENT_ENABLE, MEDIUM, ZkConfigs.ZK_SSL_CLIENT_ENABLE_DOC)
-      .define(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_DOC)
-      .define(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_DOC)
-      .define(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_DOC)
-      .define(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_DOC)
-      .define(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_DOC)
-      .define(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_DOC)
-      .define(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_DOC)
-      .define(ZkConfigs.ZK_SSL_PROTOCOL_PROP, STRING, ZkConfigs.ZK_SSL_PROTOCOL, LOW, ZkConfigs.ZK_SSL_PROTOCOL_DOC)
-      .define(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, LIST, null, LOW, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_DOC)
-      .define(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, LIST, null, LOW, ZkConfigs.ZK_SSL_CIPHER_SUITES_DOC)
-      .define(ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, STRING, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM, LOW, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC)
-      .define(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_CRL_ENABLE, LOW, ZkConfigs.ZK_SSL_CRL_ENABLE_DOC)
-      .define(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, BOOLEAN, ZkConfigs.ZK_SSL_OCSP_ENABLE, LOW, ZkConfigs.ZK_SSL_OCSP_ENABLE_DOC)
+      .define(ZkConfigs.ZK_CONNECT_CONFIG, STRING, null, HIGH, ZkConfigs.ZK_CONNECT_DOC)
+      .define(ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG, INT, ZkConfigs.ZK_SESSION_TIMEOUT_MS, HIGH, ZkConfigs.ZK_SESSION_TIMEOUT_MS_DOC)
+      .define(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_CONFIG, INT, null, HIGH, ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_DOC)
+      .define(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, BOOLEAN, ZkConfigs.ZK_ENABLE_SECURE_ACLS, HIGH, ZkConfigs.ZK_ENABLE_SECURE_ACLS_DOC)
+      .define(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_CONFIG, INT, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS, atLeast(1), HIGH, ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_DOC)
+      .define(ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG, BOOLEAN, ZkConfigs.ZK_SSL_CLIENT_ENABLE, MEDIUM, ZkConfigs.ZK_SSL_CLIENT_ENABLE_DOC)
+      .define(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG, STRING, null, MEDIUM, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_DOC)
+      .define(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG, PASSWORD, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_DOC)
+      .define(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG, STRING, null, MEDIUM, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_DOC)
+      .define(ZkConfigs.ZK_SSL_PROTOCOL_CONFIG, STRING, ZkConfigs.ZK_SSL_PROTOCOL, LOW, ZkConfigs.ZK_SSL_PROTOCOL_DOC)
+      .define(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG, LIST, null, LOW, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_DOC)
+      .define(ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG, LIST, null, LOW, ZkConfigs.ZK_SSL_CIPHER_SUITES_DOC)
+      .define(ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, STRING, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM, LOW, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC)
+      .define(ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG, BOOLEAN, ZkConfigs.ZK_SSL_CRL_ENABLE, LOW, ZkConfigs.ZK_SSL_CRL_ENABLE_DOC)
+      .define(ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG, BOOLEAN, ZkConfigs.ZK_SSL_OCSP_ENABLE, LOW, ZkConfigs.ZK_SSL_OCSP_ENABLE_DOC)
 
       /** ********* General Configuration ***********/
       .define(BrokerIdGenerationEnableProp, BOOLEAN, Defaults.BROKER_ID_GENERATION_ENABLE, MEDIUM, BrokerIdGenerationEnableDoc)
@@ -1334,12 +1334,12 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     super.valuesWithPrefixOverride(prefix)
 
   /** ********* Zookeeper Configuration ***********/
-  val zkConnect: String = getString(ZkConfigs.ZK_CONNECT_PROP)
-  val zkSessionTimeoutMs: Int = getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP)
+  val zkConnect: String = getString(ZkConfigs.ZK_CONNECT_CONFIG)
+  val zkSessionTimeoutMs: Int = getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG)
   val zkConnectionTimeoutMs: Int =
-    Option(getInt(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP)).map(_.toInt).getOrElse(getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP))
-  val zkEnableSecureAcls: Boolean = getBoolean(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP)
-  val zkMaxInFlightRequests: Int = getInt(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP)
+    Option(getInt(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_CONFIG)).map(_.toInt).getOrElse(getInt(ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG))
+  val zkEnableSecureAcls: Boolean = getBoolean(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG)
+  val zkMaxInFlightRequests: Int = getInt(ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_CONFIG)
 
   private val _remoteLogManagerConfig = new RemoteLogManagerConfig(this)
   def remoteLogManagerConfig = _remoteLogManagerConfig
@@ -1387,21 +1387,21 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     }
   }
 
-  val zkSslClientEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP)
-  val zkClientCnxnSocketClassName = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP)
-  val zkSslKeyStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP)
-  val zkSslKeyStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP)
-  val zkSslKeyStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP)
-  val zkSslTrustStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP)
-  val zkSslTrustStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP)
-  val zkSslTrustStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP)
-  val ZkSslProtocol = zkStringConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_PROTOCOL_PROP)
-  val ZkSslEnabledProtocols = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP)
-  val ZkSslCipherSuites = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP)
+  val zkSslClientEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG)
+  val zkClientCnxnSocketClassName = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG)
+  val zkSslKeyStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG)
+  val zkSslKeyStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG)
+  val zkSslKeyStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG)
+  val zkSslTrustStoreLocation = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG)
+  val zkSslTrustStorePassword = zkPasswordConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG)
+  val zkSslTrustStoreType = zkOptionalStringConfigOrSystemProperty(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG)
+  val ZkSslProtocol = zkStringConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_PROTOCOL_CONFIG)
+  val ZkSslEnabledProtocols = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG)
+  val ZkSslCipherSuites = zkListConfigOrSystemProperty(ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG)
   val ZkSslEndpointIdentificationAlgorithm = {
     // Use the system property if it exists and the Kafka config value was defaulted rather than actually provided
     // Need to translate any system property value from true/false to HTTPS/<blank>
-    val kafkaProp = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP
+    val kafkaProp = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG
     val actuallyProvided = originals.containsKey(kafkaProp)
     if (actuallyProvided)
       getString(kafkaProp)
@@ -1413,8 +1413,8 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
       }
     }
   }
-  val ZkSslCrlEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP)
-  val ZkSslOcspEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP)
+  val ZkSslCrlEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG)
+  val ZkSslOcspEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG)
   /** ********* General Configuration ***********/
   val brokerIdGenerationEnable: Boolean = getBoolean(KafkaConfig.BrokerIdGenerationEnableProp)
   val maxReservedBrokerId: Int = getInt(KafkaConfig.MaxReservedBrokerIdProp)
@@ -1957,7 +1957,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
     }
     if (requiresZookeeper) {
       if (zkConnect == null) {
-        throw new ConfigException(s"Missing required configuration `${ZkConfigs.ZK_CONNECT_PROP}` which has no default value.")
+        throw new ConfigException(s"Missing required configuration `${ZkConfigs.ZK_CONNECT_CONFIG}` which has no default value.")
       }
       if (brokerIdGenerationEnable) {
         require(brokerId >= -1 && brokerId <= maxReservedBrokerId, "broker.id must be greater than or equal to -1 and not greater than reserved.broker.max.id")
@@ -1972,7 +1972,7 @@ class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynami
       }
       if (migrationEnabled) {
         if (zkConnect == null) {
-          throw new ConfigException(s"If using `${KafkaConfig.MigrationEnabledProp}` in KRaft mode, `${ZkConfigs.ZK_CONNECT_PROP}` must also be set.")
+          throw new ConfigException(s"If using `${KafkaConfig.MigrationEnabledProp}` in KRaft mode, `${ZkConfigs.ZK_CONNECT_CONFIG}` must also be set.")
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index 3d41f01b70497..03f4fe9363267 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -80,20 +80,20 @@ object KafkaServer {
   def zkClientConfigFromKafkaConfig(config: KafkaConfig, forceZkSslClientEnable: Boolean = false): ZKClientConfig = {
     val clientConfig = new ZKClientConfig
     if (config.zkSslClientEnable || forceZkSslClientEnable) {
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "true")
-      config.zkClientCnxnSocketClassName.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP, _))
-      config.zkSslKeyStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, _))
-      config.zkSslKeyStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, x.value))
-      config.zkSslKeyStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, _))
-      config.zkSslTrustStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, _))
-      config.zkSslTrustStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, x.value))
-      config.zkSslTrustStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, _))
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_PROTOCOL_PROP, config.ZkSslProtocol)
-      config.ZkSslEnabledProtocols.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, _))
-      config.ZkSslCipherSuites.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, _))
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, config.ZkSslEndpointIdentificationAlgorithm)
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, config.ZkSslCrlEnable.toString)
-      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, config.ZkSslOcspEnable.toString)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG, "true")
+      config.zkClientCnxnSocketClassName.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG, _))
+      config.zkSslKeyStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG, _))
+      config.zkSslKeyStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG, x.value))
+      config.zkSslKeyStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG, _))
+      config.zkSslTrustStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG, _))
+      config.zkSslTrustStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG, x.value))
+      config.zkSslTrustStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_PROTOCOL_CONFIG, config.ZkSslProtocol)
+      config.ZkSslEnabledProtocols.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG, _))
+      config.ZkSslCipherSuites.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, config.ZkSslEndpointIdentificationAlgorithm)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG, config.ZkSslCrlEnable.toString)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG, config.ZkSslOcspEnable.toString)
     }
     // The zk sasl is enabled by default so it can produce false error when broker does not intend to use SASL.
     if (!JaasUtils.isZkSaslEnabled) clientConfig.setProperty(JaasUtils.ZK_SASL_CLIENT, "false")
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index 749de42853320..e2303efa90824 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -2349,9 +2349,9 @@ object KafkaZkClient {
 
     if (secureAclsEnabled && !isZkSecurityEnabled)
       throw new java.lang.SecurityException(
-        s"${ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP} is true, but ZooKeeper client TLS configuration identifying at least " +
-          s"${ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP}, and " +
-          s"${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP} was not present and the verification of the JAAS login file failed " +
+        s"${ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG} is true, but ZooKeeper client TLS configuration identifying at least " +
+          s"${ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG}, ${ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG}, and " +
+          s"${ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG} was not present and the verification of the JAAS login file failed " +
           s"${JaasUtils.zkSecuritySysConfigString}")
 
     KafkaZkClient(config.zkConnect, secureAclsEnabled, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs,
diff --git a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
index 77d0784b1a5ec..9469a2d4e4d18 100644
--- a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
@@ -76,7 +76,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
   import DescribeAuthorizedOperationsTest._
 
   override val brokerCount = 1
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
   this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
 
   var client: Admin = _
diff --git a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
index 061ba3de9eaea..0ac83d640a3d4 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
@@ -157,7 +157,7 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     } else {
       // The next two configuration parameters enable ZooKeeper secure ACLs
       // and sets the Kafka authorizer, both necessary to enable security.
-      this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+      this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
       this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, authorizerClass.getName)
 
       // Set the specific principal that can update ACLs.
diff --git a/core/src/test/scala/integration/kafka/api/MetricsTest.scala b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
index 658b57cddd31f..e48302675869f 100644
--- a/core/src/test/scala/integration/kafka/api/MetricsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
@@ -44,7 +44,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   private val kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
   private val kafkaServerJaasEntryName =
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "false")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "false")
   this.serverConfig.setProperty(KafkaConfig.AutoCreateTopicsEnableProp, "false")
   this.serverConfig.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, "2.8")
   this.producerConfig.setProperty(ProducerConfig.LINGER_MS_CONFIG, "10")
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
index fe32106877aa3..b8690d238dc05 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -2707,7 +2707,7 @@ object PlaintextAdminIntegrationTest {
     var topicConfigEntries2 = Seq(new ConfigEntry(TopicConfig.COMPRESSION_TYPE_CONFIG, "snappy")).asJava
 
     val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, test.brokers.head.config.brokerId.toString)
-    val brokerConfigEntries = Seq(new ConfigEntry(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")).asJava
+    val brokerConfigEntries = Seq(new ConfigEntry(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")).asJava
 
     // Alter configs: first and third are invalid, second is valid
     var alterResult = admin.alterConfigs(Map(
diff --git a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
index 4cf7a320ff290..4cd2145319094 100644
--- a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
@@ -25,7 +25,7 @@ import scala.jdk.CollectionConverters._
 class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaClientSaslMechanism = "PLAIN"
   private val kafkaServerSaslMechanisms = List("GSSAPI", "PLAIN")
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
diff --git a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
index f039e16faa31e..83a54f3bdc610 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
@@ -26,7 +26,7 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
   private val kafkaServerJaasEntryName =
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "false")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "false")
   // disable secure acls of zkClient in QuorumTestHarness
   override protected def zkAclsEnabled = Some(false)
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
index 1d2f4eb8bbf4c..22579d92eafca 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
@@ -47,7 +47,7 @@ class SaslSslAdminIntegrationTest extends BaseAdminIntegrationTest with SaslSetu
 
   val authorizationAdmin = new AclAuthorizationAdmin(classOf[AclAuthorizer], classOf[AclAuthorizer])
 
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
 
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
index 19b6e9e46148f..fe276fdd719f7 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslConsumerTest.scala
@@ -19,7 +19,7 @@ import org.junit.jupiter.api.{AfterEach, BeforeEach, TestInfo, Timeout}
 
 @Timeout(600)
 class SaslSslConsumerTest extends BaseConsumerTest with SaslSetup {
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
 
diff --git a/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
index ff6ea03a67745..e1cf49b5af7bf 100644
--- a/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
@@ -79,7 +79,7 @@ object SslAdminIntegrationTest {
 class SslAdminIntegrationTest extends SaslSslAdminIntegrationTest {
   override val authorizationAdmin = new AclAuthorizationAdmin(classOf[SslAdminIntegrationTest.TestableAclAuthorizer], classOf[AclAuthorizer])
 
-  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+  this.serverConfig.setProperty(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
 
   override protected def securityProtocol = SecurityProtocol.SSL
   override protected lazy val trustStoreFile = Some(TestUtils.tempFile("truststore", ".jks"))
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
index dd85bd41cb002..9eeb79c74e29c 100644
--- a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -122,7 +122,7 @@ class DynamicBrokerReconfigurationTest extends QuorumTestHarness with SaslSetup
         properties
       } else {
         val properties = TestUtils.createBrokerConfig(brokerId, zkConnect)
-        properties.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+        properties.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
         properties
       }
       props ++= securityProps(sslProperties1, TRUSTSTORE_PROPS)
diff --git a/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala b/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
index 7bf86259230fd..87abb572f4604 100644
--- a/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
+++ b/core/src/test/scala/integration/kafka/server/KafkaServerKRaftRegistrationTest.scala
@@ -62,7 +62,7 @@ class KafkaServerKRaftRegistrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -100,7 +100,7 @@ class KafkaServerKRaftRegistrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
index 16e1a2902781f..aeee7ddced30a 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
@@ -80,7 +80,7 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends QuorumT
       props.put(KafkaConfig.ListenerSecurityProtocolMapProp, s"$Internal:PLAINTEXT, $SecureInternal:SASL_SSL," +
         s"$External:PLAINTEXT, $SecureExternal:SASL_SSL")
       props.put(KafkaConfig.InterBrokerListenerNameProp, Internal)
-      props.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP, "true")
+      props.put(ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG, "true")
       props.put(KafkaConfig.SaslMechanismInterBrokerProtocolProp, kafkaClientSaslMechanism)
       props.put(s"${new ListenerName(SecureInternal).configPrefix}${KafkaConfig.SaslEnabledMechanismsProp}",
         kafkaServerSaslMechanisms(SecureInternal).mkString(","))
diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
index ff8d2b797f384..e905e3a8a1098 100644
--- a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -177,7 +177,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -309,7 +309,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -443,7 +443,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -508,7 +508,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -576,7 +576,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -636,7 +636,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
@@ -711,7 +711,7 @@ class ZkMigrationIntegrationTest {
         setNumBrokerNodes(0).
         setNumControllerNodes(1).build())
       .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
-      .setConfigProp(ZkConfigs.ZK_CONNECT_PROP, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
+      .setConfigProp(ZkConfigs.ZK_CONNECT_CONFIG, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
       .build()
     try {
       kraftCluster.format()
diff --git a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
index 3d97e7df17157..3c24df93a91ed 100644
--- a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
@@ -156,7 +156,7 @@ class KafkaTest {
       "Missing required configuration `zookeeper.connect` which has no default value.")
 
     // Ensure that no exception is thrown once zookeeper.connect is defined (and we clear controller.listener.names)
-    propertiesFile.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    propertiesFile.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     propertiesFile.setProperty(KafkaConfig.ControllerListenerNamesProp, "")
     KafkaConfig.fromProps(propertiesFile)
   }
@@ -233,61 +233,61 @@ class KafkaTest {
 
   @Test
   def testZkSslClientEnable(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "zookeeper.ssl.client.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG, "zookeeper.ssl.client.enable",
       "zookeeper.client.secure", booleanPropValueToSet, config => Some(config.zkSslClientEnable), booleanPropValueToSet, Some(false))
   }
 
   @Test
   def testZkSslKeyStoreLocation(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP, "zookeeper.ssl.keystore.location",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG, "zookeeper.ssl.keystore.location",
       "zookeeper.ssl.keyStore.location", stringPropValueToSet, config => config.zkSslKeyStoreLocation, stringPropValueToSet)
   }
 
   @Test
   def testZkSslTrustStoreLocation(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP, "zookeeper.ssl.truststore.location",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG, "zookeeper.ssl.truststore.location",
       "zookeeper.ssl.trustStore.location", stringPropValueToSet, config => config.zkSslTrustStoreLocation, stringPropValueToSet)
   }
 
   @Test
   def testZookeeperKeyStorePassword(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP, "zookeeper.ssl.keystore.password",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG, "zookeeper.ssl.keystore.password",
       "zookeeper.ssl.keyStore.password", passwordPropValueToSet, config => config.zkSslKeyStorePassword, new Password(passwordPropValueToSet))
   }
 
   @Test
   def testZookeeperTrustStorePassword(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP, "zookeeper.ssl.truststore.password",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG, "zookeeper.ssl.truststore.password",
       "zookeeper.ssl.trustStore.password", passwordPropValueToSet, config => config.zkSslTrustStorePassword, new Password(passwordPropValueToSet))
   }
 
   @Test
   def testZkSslKeyStoreType(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP, "zookeeper.ssl.keystore.type",
+    testZkConfig(ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG, "zookeeper.ssl.keystore.type",
       "zookeeper.ssl.keyStore.type", stringPropValueToSet, config => config.zkSslKeyStoreType, stringPropValueToSet)
   }
 
   @Test
   def testZkSslTrustStoreType(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP, "zookeeper.ssl.truststore.type",
+    testZkConfig(ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG, "zookeeper.ssl.truststore.type",
       "zookeeper.ssl.trustStore.type", stringPropValueToSet, config => config.zkSslTrustStoreType, stringPropValueToSet)
   }
 
   @Test
   def testZkSslProtocol(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_PROTOCOL_PROP, "zookeeper.ssl.protocol",
+    testZkConfig(ZkConfigs.ZK_SSL_PROTOCOL_CONFIG, "zookeeper.ssl.protocol",
       "zookeeper.ssl.protocol", stringPropValueToSet, config => Some(config.ZkSslProtocol), stringPropValueToSet, Some("TLSv1.2"))
   }
 
   @Test
   def testZkSslEnabledProtocols(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP, "zookeeper.ssl.enabled.protocols",
+    testZkConfig(ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG, "zookeeper.ssl.enabled.protocols",
       "zookeeper.ssl.enabledProtocols", listPropValueToSet.mkString(","), config => config.ZkSslEnabledProtocols, listPropValueToSet.asJava)
   }
 
   @Test
   def testZkSslCipherSuites(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP, "zookeeper.ssl.cipher.suites",
+    testZkConfig(ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG, "zookeeper.ssl.cipher.suites",
       "zookeeper.ssl.ciphersuites", listPropValueToSet.mkString(","), config => config.ZkSslCipherSuites, listPropValueToSet.asJava)
   }
 
@@ -295,7 +295,7 @@ class KafkaTest {
   def testZkSslEndpointIdentificationAlgorithm(): Unit = {
     // this property is different than the others
     // because the system property values and the Kafka property values don't match
-    val kafkaPropName = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP
+    val kafkaPropName = ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG
     assertEquals("zookeeper.ssl.endpoint.identification.algorithm", kafkaPropName)
     val sysProp = "zookeeper.ssl.hostnameVerification"
     val expectedDefaultValue = "HTTPS"
@@ -328,13 +328,13 @@ class KafkaTest {
 
   @Test
   def testZkSslCrlEnable(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_CRL_ENABLE_PROP, "zookeeper.ssl.crl.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG, "zookeeper.ssl.crl.enable",
       "zookeeper.ssl.crl", booleanPropValueToSet, config => Some(config.ZkSslCrlEnable), booleanPropValueToSet, Some(false))
   }
 
   @Test
   def testZkSslOcspEnable(): Unit = {
-    testZkConfig(ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP, "zookeeper.ssl.ocsp.enable",
+    testZkConfig(ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG, "zookeeper.ssl.ocsp.enable",
       "zookeeper.ssl.ocsp", booleanPropValueToSet, config => Some(config.ZkSslOcspEnable), booleanPropValueToSet, Some(false))
   }
 
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
index 2730431565e7f..1d5e6d53d37e3 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
@@ -896,7 +896,7 @@ class ControllerChannelManagerTest {
   private def createConfig(interBrokerVersion: MetadataVersion): KafkaConfig = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, controllerId.toString)
-    props.put(ZkConfigs.ZK_CONNECT_PROP, "zkConnect")
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, "zkConnect")
     TestUtils.setIbpAndMessageFormatVersions(props, interBrokerVersion)
     KafkaConfig.fromProps(props)
   }
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
index 29907f220685a..8b424d44944be 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
@@ -880,27 +880,27 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val props = new java.util.Properties()
     val kafkaValue = "kafkaValue"
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
-      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue)
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG -> "true",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG -> kafkaValue)
     configs.foreach { case (key, value) => props.put(key, value) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG =>
         assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG =>
         assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case ZkConfigs.ZK_SSL_PROTOCOL_PROP =>
+      case ZkConfigs.ZK_SSL_PROTOCOL_CONFIG =>
         assertEquals("TLSv1.2", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       case _ => assertEquals(kafkaValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
     })
@@ -911,29 +911,29 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val props = new java.util.Properties()
     val kafkaValue = "kafkaValue"
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
-      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_PROTOCOL_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "HTTPS",
-      ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "false",
-      ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "false")
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG -> "true",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_PROTOCOL_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG -> "HTTPS",
+      ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG -> "false",
+      ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG -> "false")
     configs.foreach{case (key, value) => props.put(key, value.toString) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
-        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
+        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG | ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG =>
           assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-        case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
+        case ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG =>
           assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
         case _ => assertEquals(kafkaValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       })
@@ -946,43 +946,43 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val prefixedValue = "prefixedValue"
     val prefix = "authorizer."
     val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
-      ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "false",
-      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_PROTOCOL_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> kafkaValue,
-      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "HTTPS",
-      ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "false",
-      ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "false",
-      prefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP -> "true",
-      prefix + ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_PROTOCOL_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP -> prefixedValue,
-      prefix + ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP -> "",
-      prefix + ZkConfigs.ZK_SSL_CRL_ENABLE_PROP -> "true",
-      prefix + ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP -> "true")
+      ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG -> "false",
+      ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_PROTOCOL_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG -> kafkaValue,
+      ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG -> "HTTPS",
+      ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG -> "false",
+      ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG -> "false",
+      prefix + ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG -> "true",
+      prefix + ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_PROTOCOL_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG -> prefixedValue,
+      prefix + ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG -> "",
+      prefix + ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG -> "true",
+      prefix + ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG -> "true")
     configs.foreach{case (key, value) => props.put(key, value.toString) }
 
     val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
       KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
     // confirm we get all the values we expect
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(prop => prop match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP =>
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG | ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG =>
         assertEquals("true", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
-      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP =>
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG =>
         assertEquals("false", KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
       case _ => assertEquals(prefixedValue, KafkaConfig.zooKeeperClientProperty(zkClientConfig, prop).getOrElse("<None>"))
     })
diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
index 8e7fe06ef2eaf..973737b4bab7f 100755
--- a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
@@ -211,7 +211,7 @@ class DynamicBrokerConfigTest {
 
     val securityPropsWithoutListenerPrefix = Map(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG -> "PKCS12")
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, securityPropsWithoutListenerPrefix)
-    val nonDynamicProps = Map(ZkConfigs.ZK_CONNECT_PROP -> "somehost:2181")
+    val nonDynamicProps = Map(ZkConfigs.ZK_CONNECT_CONFIG -> "somehost:2181")
     verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, nonDynamicProps)
 
     // Test update of configs with invalid type
@@ -709,7 +709,7 @@ class DynamicBrokerConfigTest {
   @Test
   def testNonInternalValuesDoesNotExposeInternalConfigs(): Unit = {
     val props = new Properties()
-    props.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     props.put(KafkaConfig.MetadataLogSegmentMinBytesProp, "1024")
     val config = new KafkaConfig(props)
     assertFalse(config.nonInternalValues.containsKey(KafkaConfig.MetadataLogSegmentMinBytesProp))
diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
index fd4377f658e1e..391a93ab00008 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
@@ -155,7 +155,7 @@ class KafkaConfigTest {
     val hostName = "fake-host"
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, s"PLAINTEXT://$hostName:$port")
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -186,7 +186,7 @@ class KafkaConfigTest {
   def testDuplicateListeners(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     // listeners with duplicate port
     props.setProperty(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9091,SSL://localhost:9091")
@@ -212,7 +212,7 @@ class KafkaConfigTest {
   def testIPv4AndIPv6SamePortListeners(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
-    props.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.put(KafkaConfig.ListenersProp, "PLAINTEXT://[::1]:9092,SSL://[::1]:9092")
     var caught = assertThrows(classOf[IllegalArgumentException], () => KafkaConfig.fromProps(props))
@@ -452,7 +452,7 @@ class KafkaConfigTest {
   def testControllerListenerNameDoesNotMapToPlaintextByDefaultForNonKRaft(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "CONTROLLER://localhost:9092")
     assertBadConfigContainingMessage(props,
       "Error creating broker listeners from 'CONTROLLER://localhost:9092': No security protocol defined for listener CONTROLLER")
@@ -465,7 +465,7 @@ class KafkaConfigTest {
   def testBadListenerProtocol(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "BAD://localhost:9091")
 
     assertFalse(isValidKafkaConfig(props))
@@ -475,7 +475,7 @@ class KafkaConfigTest {
   def testListenerNamesWithAdvertisedListenerUnset(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "CLIENT://localhost:9091,REPLICATION://localhost:9092,INTERNAL://localhost:9093")
     props.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, "CLIENT:SSL,REPLICATION:SSL,INTERNAL:PLAINTEXT")
@@ -499,7 +499,7 @@ class KafkaConfigTest {
   def testListenerAndAdvertisedListenerNames(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "EXTERNAL://localhost:9091,INTERNAL://localhost:9093")
     props.setProperty(KafkaConfig.AdvertisedListenersProp, "EXTERNAL://lb1.example.com:9000,INTERNAL://host1:9093")
@@ -530,7 +530,7 @@ class KafkaConfigTest {
   def testListenerNameMissingFromListenerSecurityProtocolMap(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091,REPLICATION://localhost:9092")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "SSL")
@@ -541,7 +541,7 @@ class KafkaConfigTest {
   def testInterBrokerListenerNameMissingFromListenerSecurityProtocolMap(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "REPLICATION")
@@ -552,7 +552,7 @@ class KafkaConfigTest {
   def testInterBrokerListenerNameAndSecurityProtocolSet(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     props.setProperty(KafkaConfig.ListenersProp, "SSL://localhost:9091")
     props.setProperty(KafkaConfig.InterBrokerListenerNameProp, "SSL")
@@ -564,7 +564,7 @@ class KafkaConfigTest {
   def testCaseInsensitiveListenerProtocol(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     props.setProperty(KafkaConfig.ListenersProp, "plaintext://localhost:9091,SsL://localhost:9092")
     val config = KafkaConfig.fromProps(props)
     assertEquals(Some("SSL://localhost:9092"), config.listeners.find(_.listenerName.value == "SSL").map(_.connectionString))
@@ -579,7 +579,7 @@ class KafkaConfigTest {
   def testListenerDefaults(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
 
     // configuration with no listeners
     val conf = KafkaConfig.fromProps(props)
@@ -593,7 +593,7 @@ class KafkaConfigTest {
   def testVersionConfiguration(): Unit = {
     val props = new Properties()
     props.setProperty(KafkaConfig.BrokerIdProp, "1")
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     val conf = KafkaConfig.fromProps(props)
     assertEquals(MetadataVersion.latestProduction, conf.interBrokerProtocolVersion)
 
@@ -766,7 +766,7 @@ class KafkaConfigTest {
   def testFromPropsInvalid(): Unit = {
     def baseProperties: Properties = {
       val validRequiredProperties = new Properties()
-      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
+      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "127.0.0.1:2181")
       validRequiredProperties
     }
     // to ensure a basis is valid - bootstraps all needed validation
@@ -774,25 +774,25 @@ class KafkaConfigTest {
 
     KafkaConfig.configNames.foreach { name =>
       name match {
-        case ZkConfigs.ZK_CONNECT_PROP => // ignore string
-        case ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number")
-        case ZkConfigs.ZK_ENABLE_SECURE_ACLS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_PROP => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
-        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_KEY_STORE_TYPE_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_PROTOCOL_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP =>  //ignore string
-        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => //ignore string
-        case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
-        case ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_CONNECT_CONFIG => // ignore string
+        case ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case ZkConfigs.ZK_ENABLE_SECURE_ACLS_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_MAX_IN_FLIGHT_REQUESTS_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_LOCATION_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_PASSWORD_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_KEY_STORE_TYPE_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_LOCATION_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_PASSWORD_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_TRUST_STORE_TYPE_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_PROTOCOL_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG =>  //ignore string
+        case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => //ignore string
+        case ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
 
         case KafkaConfig.BrokerIdProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.NumNetworkThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
@@ -1051,7 +1051,7 @@ class KafkaConfigTest {
   def testDynamicLogConfigs(): Unit = {
     def baseProperties: Properties = {
       val validRequiredProperties = new Properties()
-      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
+      validRequiredProperties.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "127.0.0.1:2181")
       validRequiredProperties
     }
 
@@ -1141,9 +1141,9 @@ class KafkaConfigTest {
   @Test
   def testSpecificProperties(): Unit = {
     val defaults = new Properties()
-    defaults.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
+    defaults.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "127.0.0.1:2181")
     // For ZkConnectionTimeoutMs
-    defaults.setProperty(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, "1234")
+    defaults.setProperty(ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG, "1234")
     defaults.setProperty(KafkaConfig.BrokerIdGenerationEnableProp, "false")
     defaults.setProperty(KafkaConfig.MaxReservedBrokerIdProp, "1")
     defaults.setProperty(KafkaConfig.BrokerIdProp, "1")
@@ -1187,7 +1187,7 @@ class KafkaConfigTest {
   @Test
   def testNonroutableAdvertisedListeners(): Unit = {
     val props = new Properties()
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "127.0.0.1:2181")
     props.setProperty(KafkaConfig.ListenersProp, "PLAINTEXT://0.0.0.0:9092")
     assertFalse(isValidKafkaConfig(props))
   }
@@ -1601,7 +1601,7 @@ class KafkaConfigTest {
   @Test
   def testSaslJwksEndpointRetryDefaults(): Unit = {
     val props = new Properties()
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     val config = KafkaConfig.fromProps(props)
     assertNotNull(config.getLong(KafkaConfig.SaslOAuthBearerJwksEndpointRetryBackoffMsProp))
     assertNotNull(config.getLong(KafkaConfig.SaslOAuthBearerJwksEndpointRetryBackoffMaxMsProp))
@@ -1769,7 +1769,7 @@ class KafkaConfigTest {
       "If using `zookeeper.metadata.migration.enable` in KRaft mode, `zookeeper.connect` must also be set.",
       assertThrows(classOf[ConfigException], () => KafkaConfig.fromProps(props)).getMessage)
 
-    props.setProperty(ZkConfigs.ZK_CONNECT_PROP, "localhost:2181")
+    props.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:2181")
     KafkaConfig.fromProps(props)
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
index 035f60dab1db1..6cf19458dabc8 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
@@ -67,7 +67,7 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkConfigWhenSaslDisabled(): Unit = {
     val props = new Properties
-    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, zkConnect) // required, otherwise we would leave it out
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     assertEquals("false", zkClientConfig.getProperty(JaasUtils.ZK_SASL_CLIENT))
   }
@@ -75,8 +75,8 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkTlsConfigWhenDisabled(): Unit = {
     val props = new Properties
-    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
-    props.put(ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP, "false")
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG, "false")
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach { propName =>
       assertNull(zkClientConfig.getProperty(propName))
@@ -86,20 +86,20 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkTlsConfigWithTrueValues(): Unit = {
     val props = new Properties
-    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, zkConnect) // required, otherwise we would leave it out
     // should get correct config for all properties if TLS is enabled
     val someValue = "some_value"
     def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "true"
-      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "HTTPS"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG | ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG => "true"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => "HTTPS"
       case _ => someValue
     }
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     // now check to make sure the values were set correctly
     def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP | ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "true"
-      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "true"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG | ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG => "true"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => "true"
       case _ => someValue
     }
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp =>
@@ -109,24 +109,24 @@ class KafkaServerTest extends QuorumTestHarness {
   @Test
   def testCreatesProperZkTlsConfigWithFalseAndListValues(): Unit = {
     val props = new Properties
-    props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect) // required, otherwise we would leave it out
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, zkConnect) // required, otherwise we would leave it out
     // should get correct config for all properties if TLS is enabled
     val someValue = "some_value"
     def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => "true"
-      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "false"
-      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => ""
-      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => "A,B"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG => "true"
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG => "false"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => ""
+      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG | ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG => "A,B"
       case _ => someValue
     }
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
     val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
     // now check to make sure the values were set correctly
     def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
-      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_PROP => "true"
-      case ZkConfigs.ZK_SSL_CRL_ENABLE_PROP | ZkConfigs.ZK_SSL_OCSP_ENABLE_PROP => "false"
-      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP => "false"
-      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_PROP | ZkConfigs.ZK_SSL_CIPHER_SUITES_PROP => "A,B"
+      case ZkConfigs.ZK_SSL_CLIENT_ENABLE_CONFIG => "true"
+      case ZkConfigs.ZK_SSL_CRL_ENABLE_CONFIG | ZkConfigs.ZK_SSL_OCSP_ENABLE_CONFIG => "false"
+      case ZkConfigs.ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG => "false"
+      case ZkConfigs.ZK_SSL_ENABLED_PROTOCOLS_CONFIG | ZkConfigs.ZK_SSL_CIPHER_SUITES_CONFIG => "A,B"
       case _ => someValue
     }
     ZkConfigs.ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.asScala.keys.foreach(kafkaProp =>
diff --git a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
index 0e04bad7c8ef2..f6d0853d56f26 100644
--- a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
@@ -151,8 +151,8 @@ class ServerShutdownTest extends KafkaServerTestHarness {
       shutdownKRaftController()
       verifyCleanShutdownAfterFailedStartup[CancellationException]
     } else {
-      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, "50")
-      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECT_PROP, "some.invalid.hostname.foo.bar.local:65535")
+      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_CONFIG, "50")
+      propsToChangeUponRestart.setProperty(ZkConfigs.ZK_CONNECT_CONFIG, "some.invalid.hostname.foo.bar.local:65535")
       verifyCleanShutdownAfterFailedStartup[ZooKeeperClientTimeoutException]
     }
   }
diff --git a/core/src/test/scala/unit/kafka/server/ServerTest.scala b/core/src/test/scala/unit/kafka/server/ServerTest.scala
index fb80c4b890d05..af6233c23068a 100644
--- a/core/src/test/scala/unit/kafka/server/ServerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerTest.scala
@@ -55,7 +55,7 @@ class ServerTest {
 
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, brokerId.toString)
-    props.put(ZkConfigs.ZK_CONNECT_PROP, "127.0.0.1:0")
+    props.put(ZkConfigs.ZK_CONNECT_CONFIG, "127.0.0.1:0")
     val config = KafkaConfig.fromProps(props)
 
     val context = Server.createKafkaMetricsContext(config, clusterId)
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 33a4b27fd8be4..a08dc797af491 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -359,8 +359,8 @@ object TestUtils extends Logging {
       // controllerQuorumVotersFuture instead.
       props.put(KafkaConfig.QuorumVotersProp, "1000@localhost:0")
     } else {
-      props.put(ZkConfigs.ZK_CONNECT_PROP, zkConnect)
-      props.put(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_PROP, "10000")
+      props.put(ZkConfigs.ZK_CONNECT_CONFIG, zkConnect)
+      props.put(ZkConfigs.ZK_CONNECTION_TIMEOUT_MS_CONFIG, "10000")
     }
     props.put(KafkaConfig.ReplicaSocketTimeoutMsProp, "1500")
     props.put(KafkaConfig.ControllerSocketTimeoutMsProp, "1500")
diff --git a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
index aa97f6831c361..27720860ee411 100644
--- a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
@@ -106,7 +106,7 @@ class KafkaZkClientTest extends QuorumTestHarness {
     // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
     // to kafka.zk.EmbeddedZookeeper
     val clientConfig = new ZKClientConfig()
-    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP
+    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG
     val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
     KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
     val client = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
index b4dccaecb0ac8..3e4652c6180c4 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkMigrationTestHarness.scala
@@ -41,7 +41,7 @@ class ZkMigrationTestHarness extends QuorumTestHarness {
 
   val encoder: PasswordEncoder = {
     val encoderProps = new Properties()
-    encoderProps.put(ZkConfigs.ZK_CONNECT_PROP, "localhost:1234") // Get around the config validation
+    encoderProps.put(ZkConfigs.ZK_CONNECT_CONFIG, "localhost:1234") // Get around the config validation
     encoderProps.put(KafkaConfig.PasswordEncoderSecretProp, SECRET) // Zk secret to encrypt the
     val encoderConfig = new KafkaConfig(encoderProps)
     PasswordEncoder.encrypting(encoderConfig.passwordEncoderSecret.get,
diff --git a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
index affe6c3839c02..650da171290a3 100644
--- a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
@@ -103,7 +103,7 @@ class ZooKeeperClientTest extends QuorumTestHarness {
     // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
     // to kafka.zk.EmbeddedZookeeper
     val clientConfig = new ZKClientConfig()
-    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_PROP
+    val propKey = ZkConfigs.ZK_CLIENT_CNXN_SOCKET_CONFIG
     val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
     KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
     val client = newZooKeeperClient(clientConfig = clientConfig)
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
index f4d17ed074020..1c5e499145fcf 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/MetadataRequestBenchmark.java
@@ -177,7 +177,7 @@ private List<UpdateMetadataEndpoint> endpoints(final int brokerId) {
 
     private KafkaApis createKafkaApis() {
         Properties kafkaProps =  new Properties();
-        kafkaProps.put(ZkConfigs.ZK_CONNECT_PROP, "zk");
+        kafkaProps.put(ZkConfigs.ZK_CONNECT_CONFIG, "zk");
         kafkaProps.put(KafkaConfig$.MODULE$.BrokerIdProp(), brokerId + "");
         KafkaConfig config = new KafkaConfig(kafkaProps);
         return new KafkaApisBuilder().
diff --git a/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java b/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java
index 2ddd2ef2145f2..57a443ed56a02 100644
--- a/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java
+++ b/server/src/main/java/org/apache/kafka/server/config/ZkConfigs.java
@@ -24,25 +24,25 @@
 
 public final class ZkConfigs {
     /** ********* Zookeeper Configuration ***********/
-    public static final String ZK_CONNECT_PROP = "zookeeper.connect";
-    public static final String ZK_SESSION_TIMEOUT_MS_PROP = "zookeeper.session.timeout.ms";
-    public static final String ZK_CONNECTION_TIMEOUT_MS_PROP = "zookeeper.connection.timeout.ms";
-    public static final String ZK_ENABLE_SECURE_ACLS_PROP = "zookeeper.set.acl";
-    public static final String ZK_MAX_IN_FLIGHT_REQUESTS_PROP = "zookeeper.max.in.flight.requests";
-    public static final String ZK_SSL_CLIENT_ENABLE_PROP = "zookeeper.ssl.client.enable";
-    public static final String ZK_CLIENT_CNXN_SOCKET_PROP = "zookeeper.clientCnxnSocket";
-    public static final String ZK_SSL_KEY_STORE_LOCATION_PROP = "zookeeper.ssl.keystore.location";
-    public static final String ZK_SSL_KEY_STORE_PASSWORD_PROP = "zookeeper.ssl.keystore.password";
-    public static final String ZK_SSL_KEY_STORE_TYPE_PROP = "zookeeper.ssl.keystore.type";
-    public static final String ZK_SSL_TRUST_STORE_LOCATION_PROP = "zookeeper.ssl.truststore.location";
-    public static final String ZK_SSL_TRUST_STORE_PASSWORD_PROP = "zookeeper.ssl.truststore.password";
-    public static final String ZK_SSL_TRUST_STORE_TYPE_PROP = "zookeeper.ssl.truststore.type";
-    public static final String ZK_SSL_PROTOCOL_PROP = "zookeeper.ssl.protocol";
-    public static final String ZK_SSL_ENABLED_PROTOCOLS_PROP = "zookeeper.ssl.enabled.protocols";
-    public static final String ZK_SSL_CIPHER_SUITES_PROP = "zookeeper.ssl.cipher.suites";
-    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP = "zookeeper.ssl.endpoint.identification.algorithm";
-    public static final String ZK_SSL_CRL_ENABLE_PROP = "zookeeper.ssl.crl.enable";
-    public static final String ZK_SSL_OCSP_ENABLE_PROP = "zookeeper.ssl.ocsp.enable";
+    public static final String ZK_CONNECT_CONFIG = "zookeeper.connect";
+    public static final String ZK_SESSION_TIMEOUT_MS_CONFIG = "zookeeper.session.timeout.ms";
+    public static final String ZK_CONNECTION_TIMEOUT_MS_CONFIG = "zookeeper.connection.timeout.ms";
+    public static final String ZK_ENABLE_SECURE_ACLS_CONFIG = "zookeeper.set.acl";
+    public static final String ZK_MAX_IN_FLIGHT_REQUESTS_CONFIG = "zookeeper.max.in.flight.requests";
+    public static final String ZK_SSL_CLIENT_ENABLE_CONFIG = "zookeeper.ssl.client.enable";
+    public static final String ZK_CLIENT_CNXN_SOCKET_CONFIG = "zookeeper.clientCnxnSocket";
+    public static final String ZK_SSL_KEY_STORE_LOCATION_CONFIG = "zookeeper.ssl.keystore.location";
+    public static final String ZK_SSL_KEY_STORE_PASSWORD_CONFIG = "zookeeper.ssl.keystore.password";
+    public static final String ZK_SSL_KEY_STORE_TYPE_CONFIG = "zookeeper.ssl.keystore.type";
+    public static final String ZK_SSL_TRUST_STORE_LOCATION_CONFIG = "zookeeper.ssl.truststore.location";
+    public static final String ZK_SSL_TRUST_STORE_PASSWORD_CONFIG = "zookeeper.ssl.truststore.password";
+    public static final String ZK_SSL_TRUST_STORE_TYPE_CONFIG = "zookeeper.ssl.truststore.type";
+    public static final String ZK_SSL_PROTOCOL_CONFIG = "zookeeper.ssl.protocol";
+    public static final String ZK_SSL_ENABLED_PROTOCOLS_CONFIG = "zookeeper.ssl.enabled.protocols";
+    public static final String ZK_SSL_CIPHER_SUITES_CONFIG = "zookeeper.ssl.cipher.suites";
+    public static final String ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG = "zookeeper.ssl.endpoint.identification.algorithm";
+    public static final String ZK_SSL_CRL_ENABLE_CONFIG = "zookeeper.ssl.crl.enable";
+    public static final String ZK_SSL_OCSP_ENABLE_CONFIG = "zookeeper.ssl.ocsp.enable";
 
     public static final String ZK_CONNECT_DOC = "Specifies the ZooKeeper connection string in the form <code>hostname:port</code> where host and port are the " +
         "host and port of a ZooKeeper server. To allow connecting through other ZooKeeper nodes when that ZooKeeper machine is " +
@@ -50,7 +50,7 @@ public final class ZkConfigs {
         "The server can also have a ZooKeeper chroot path as part of its ZooKeeper connection string which puts its data under some path in the global ZooKeeper namespace. " +
         "For example to give a chroot path of <code>/chroot/path</code> you would give the connection string as <code>hostname1:port1,hostname2:port2,hostname3:port3/chroot/path</code>.";
     public static final String ZK_SESSION_TIMEOUT_MS_DOC = "Zookeeper session timeout";
-    public static final String ZK_CONNECTION_TIMEOUT_MS_DOC = "The max time that the client waits to establish a connection to ZooKeeper. If not set, the value in " + ZK_SESSION_TIMEOUT_MS_PROP + " is used";
+    public static final String ZK_CONNECTION_TIMEOUT_MS_DOC = "The max time that the client waits to establish a connection to ZooKeeper. If not set, the value in " + ZK_SESSION_TIMEOUT_MS_CONFIG + " is used";
     public static final String ZK_ENABLE_SECURE_ACLS_DOC = "Set client to use secure ACLs";
     public static final String ZK_MAX_IN_FLIGHT_REQUESTS_DOC = "The maximum number of unacknowledged requests the client will send to ZooKeeper before blocking.";
     public static final String ZK_SSL_CLIENT_ENABLE_DOC;
@@ -88,58 +88,58 @@ public final class ZkConfigs {
     static {
         Map<String, String> zkSslConfigToSystemPropertyMap = new HashMap<>();
 
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CLIENT_ENABLE_PROP, SECURE_CLIENT);
-        zkSslConfigToSystemPropertyMap.put(ZK_CLIENT_CNXN_SOCKET_PROP, ZOOKEEPER_CLIENT_CNXN_SOCKET);
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_LOCATION_PROP, "zookeeper.ssl.keyStore.location");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_PASSWORD_PROP, "zookeeper.ssl.keyStore.password");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_TYPE_PROP, "zookeeper.ssl.keyStore.type");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_LOCATION_PROP, "zookeeper.ssl.trustStore.location");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_PASSWORD_PROP, "zookeeper.ssl.trustStore.password");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_TYPE_PROP, "zookeeper.ssl.trustStore.type");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_PROTOCOL_PROP, "zookeeper.ssl.protocol");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENABLED_PROTOCOLS_PROP, "zookeeper.ssl.enabledProtocols");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CIPHER_SUITES_PROP, "zookeeper.ssl.ciphersuites");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP, "zookeeper.ssl.hostnameVerification");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CRL_ENABLE_PROP, "zookeeper.ssl.crl");
-        zkSslConfigToSystemPropertyMap.put(ZK_SSL_OCSP_ENABLE_PROP, "zookeeper.ssl.ocsp");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CLIENT_ENABLE_CONFIG, SECURE_CLIENT);
+        zkSslConfigToSystemPropertyMap.put(ZK_CLIENT_CNXN_SOCKET_CONFIG, ZOOKEEPER_CLIENT_CNXN_SOCKET);
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_LOCATION_CONFIG, "zookeeper.ssl.keyStore.location");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_PASSWORD_CONFIG, "zookeeper.ssl.keyStore.password");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_KEY_STORE_TYPE_CONFIG, "zookeeper.ssl.keyStore.type");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_LOCATION_CONFIG, "zookeeper.ssl.trustStore.location");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_PASSWORD_CONFIG, "zookeeper.ssl.trustStore.password");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_TRUST_STORE_TYPE_CONFIG, "zookeeper.ssl.trustStore.type");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_PROTOCOL_CONFIG, "zookeeper.ssl.protocol");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENABLED_PROTOCOLS_CONFIG, "zookeeper.ssl.enabledProtocols");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CIPHER_SUITES_CONFIG, "zookeeper.ssl.ciphersuites");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "zookeeper.ssl.hostnameVerification");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_CRL_ENABLE_CONFIG, "zookeeper.ssl.crl");
+        zkSslConfigToSystemPropertyMap.put(ZK_SSL_OCSP_ENABLE_CONFIG, "zookeeper.ssl.ocsp");
 
         ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP = Collections.unmodifiableMap(zkSslConfigToSystemPropertyMap);
 
         ZK_SSL_CLIENT_ENABLE_DOC = "Set client to use TLS when connecting to ZooKeeper." +
             " An explicit value overrides any value set via the <code>zookeeper.client.secure</code> system property (note the different name)." +
-            " Defaults to false if neither is set; when true, <code>" + ZK_CLIENT_CNXN_SOCKET_PROP + "</code> must be set (typically to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code>); other values to set may include " +
-            ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.keySet().stream().filter(x -> !x.equals(ZK_SSL_CLIENT_ENABLE_PROP) && !x.equals(ZK_CLIENT_CNXN_SOCKET_PROP)).sorted().collect(Collectors.joining("<code>", "</code>, <code>", "</code>"));
+            " Defaults to false if neither is set; when true, <code>" + ZK_CLIENT_CNXN_SOCKET_CONFIG + "</code> must be set (typically to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code>); other values to set may include " +
+            ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.keySet().stream().filter(x -> !x.equals(ZK_SSL_CLIENT_ENABLE_CONFIG) && !x.equals(ZK_CLIENT_CNXN_SOCKET_CONFIG)).sorted().collect(Collectors.joining("<code>", "</code>, <code>", "</code>"));
         ZK_CLIENT_CNXN_SOCKET_DOC = "Typically set to <code>org.apache.zookeeper.ClientCnxnSocketNetty</code> when using TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_CLIENT_CNXN_SOCKET_PROP) + "</code> system property.";
+            " Overrides any explicit value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_CLIENT_CNXN_SOCKET_CONFIG) + "</code> system property.";
         ZK_SSL_KEY_STORE_LOCATION_DOC = "Keystore location when using a client-side certificate with TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_LOCATION_PROP) + "</code> system property (note the camelCase).";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_LOCATION_CONFIG) + "</code> system property (note the camelCase).";
         ZK_SSL_KEY_STORE_PASSWORD_DOC = "Keystore password when using a client-side certificate with TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_PASSWORD_PROP) + "</code> system property (note the camelCase)." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_PASSWORD_CONFIG) + "</code> system property (note the camelCase)." +
             " Note that ZooKeeper does not support a key password different from the keystore password, so be sure to set the key password in the keystore to be identical to the keystore password; otherwise the connection attempt to Zookeeper will fail.";
         ZK_SSL_KEY_STORE_TYPE_DOC = "Keystore type when using a client-side certificate with TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_TYPE_PROP) + "</code> system property (note the camelCase)." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_KEY_STORE_TYPE_CONFIG) + "</code> system property (note the camelCase)." +
             " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the keystore.";
         ZK_SSL_TRUST_STORE_LOCATION_DOC = "Truststore location when using TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_LOCATION_PROP) + "</code> system property (note the camelCase).";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_LOCATION_CONFIG) + "</code> system property (note the camelCase).";
         ZK_SSL_TRUST_STORE_PASSWORD_DOC = "Truststore password when using TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_PASSWORD_PROP) + "</code> system property (note the camelCase).";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_PASSWORD_CONFIG) + "</code> system property (note the camelCase).";
         ZK_SSL_TRUST_STORE_TYPE_DOC = "Truststore type when using TLS connectivity to ZooKeeper." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_TYPE_PROP) + "</code> system property (note the camelCase)." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_TRUST_STORE_TYPE_CONFIG) + "</code> system property (note the camelCase)." +
             " The default value of <code>null</code> means the type will be auto-detected based on the filename extension of the truststore.";
         ZK_SSL_PROTOCOL_DOC = "Specifies the protocol to be used in ZooKeeper TLS negotiation." +
-            " An explicit value overrides any value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_PROTOCOL_PROP) + "</code> system property.";
+            " An explicit value overrides any value set via the same-named <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_PROTOCOL_CONFIG) + "</code> system property.";
         ZK_SSL_ENABLED_PROTOCOLS_DOC = "Specifies the enabled protocol(s) in ZooKeeper TLS negotiation (csv)." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENABLED_PROTOCOLS_PROP) + "</code> system property (note the camelCase)." +
-            " The default value of <code>null</code> means the enabled protocol will be the value of the <code>" + ZK_SSL_PROTOCOL_PROP + "</code> configuration property.";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENABLED_PROTOCOLS_CONFIG) + "</code> system property (note the camelCase)." +
+            " The default value of <code>null</code> means the enabled protocol will be the value of the <code>" + ZK_SSL_PROTOCOL_CONFIG + "</code> configuration property.";
         ZK_SSL_CIPHER_SUITES_DOC = "Specifies the enabled cipher suites to be used in ZooKeeper TLS negotiation (csv)." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CIPHER_SUITES_PROP) + "</code> system property (note the single word \"ciphersuites\")." +
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CIPHER_SUITES_CONFIG) + "</code> system property (note the single word \"ciphersuites\")." +
             " The default value of <code>null</code> means the list of enabled cipher suites is determined by the Java runtime being used.";
         ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC = "Specifies whether to enable hostname verification in the ZooKeeper TLS negotiation process, with (case-insensitively) \"https\" meaning ZooKeeper hostname verification is enabled and an explicit blank value meaning it is disabled (disabling it is only recommended for testing purposes)." +
-            " An explicit value overrides any \"true\" or \"false\" value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_PROP) + "</code> system property (note the different name and values; true implies https and false implies blank).";
+            " An explicit value overrides any \"true\" or \"false\" value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG) + "</code> system property (note the different name and values; true implies https and false implies blank).";
         ZK_SSL_CRL_ENABLE_DOC = "Specifies whether to enable Certificate Revocation List in the ZooKeeper TLS protocols." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CRL_ENABLE_PROP) + "</code> system property (note the shorter name).";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_CRL_ENABLE_CONFIG) + "</code> system property (note the shorter name).";
         ZK_SSL_OCSP_ENABLE_DOC = "Specifies whether to enable Online Certificate Status Protocol in the ZooKeeper TLS protocols." +
-            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_OCSP_ENABLE_PROP) + "</code> system property (note the shorter name).";
+            " Overrides any explicit value set via the <code>" + ZK_SSL_CONFIG_TO_SYSTEM_PROPERTY_MAP.get(ZK_SSL_OCSP_ENABLE_CONFIG) + "</code> system property (note the shorter name).";
     }
 
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
index 89af1d5729e84..161604d8c5384 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
@@ -110,7 +110,7 @@ public void start() throws IOException {
         zookeeper = new EmbeddedZookeeper();
         log.debug("ZooKeeper instance is running at {}", zKConnectString());
 
-        brokerConfig.put(ZkConfigs.ZK_CONNECT_PROP, zKConnectString());
+        brokerConfig.put(ZkConfigs.ZK_CONNECT_CONFIG, zKConnectString());
         putIfAbsent(brokerConfig, KafkaConfig.ListenersProp(), "PLAINTEXT://localhost:" + DEFAULT_BROKER_PORT);
         putIfAbsent(brokerConfig, KafkaConfig.DeleteTopicEnableProp(), true);
         putIfAbsent(brokerConfig, CleanerConfig.LOG_CLEANER_DEDUPE_BUFFER_SIZE_PROP, 2 * 1024 * 1024L);
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
index 4b57d576ae1c7..b0a80a8da1f8d 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/KafkaEmbedded.java
@@ -91,7 +91,7 @@ private Properties effectiveConfigFrom(final Properties initialConfig) {
         effectiveConfig.put(KafkaConfig.AutoCreateTopicsEnableProp(), true);
         effectiveConfig.put(KafkaConfig.MessageMaxBytesProp(), 1000000);
         effectiveConfig.put(KafkaConfig.ControlledShutdownEnableProp(), true);
-        effectiveConfig.put(ZkConfigs.ZK_SESSION_TIMEOUT_MS_PROP, 10000);
+        effectiveConfig.put(ZkConfigs.ZK_SESSION_TIMEOUT_MS_CONFIG, 10000);
 
         effectiveConfig.putAll(initialConfig);
         effectiveConfig.setProperty(KafkaConfig.LogDirProp(), logDir.getAbsolutePath());

From 8b274d8c1bfbfa6d4319ded884a11da790d7bf77 Mon Sep 17 00:00:00 2001
From: Walker Carlson <18128741+wcarlson5@users.noreply.github.com>
Date: Thu, 28 Mar 2024 10:30:18 -0500
Subject: [PATCH 225/258] KAFKA-7663: Reprocessing on user added global stores
 restore (#15414)

When custom processors are added via StreamBuilder#addGlobalStore they will now reprocess all records through the custom transformer instead of loading directly.

We do this so that users that transform the records will not get improperly formatted records down stream.

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../apache/kafka/streams/StreamsBuilder.java  |  13 +-
 .../org/apache/kafka/streams/Topology.java    |  12 +-
 .../internals/InternalStreamsBuilder.java     |   7 +-
 .../internals/graph/GlobalStoreNode.java      |   9 +-
 .../internals/graph/TableSourceNode.java      |   3 +-
 .../internals/GlobalStateManagerImpl.java     | 123 ++++++++++-
 .../internals/InternalTopologyBuilder.java    |  39 +++-
 .../internals/ProcessorTopology.java          |  10 +-
 .../internals/RecordDeserializer.java         |  71 ++++---
 .../internals/InMemoryKeyValueStore.java      |   2 +
 .../integration/GlobalStateReprocessTest.java | 201 ++++++++++++++++++
 .../GlobalThreadShutDownOrderTest.java        |  18 +-
 .../internals/GlobalStateManagerImplTest.java |  85 +++++++-
 .../internals/GlobalStreamThreadTest.java     |   4 +-
 .../InternalTopologyBuilderTest.java          |  31 ++-
 .../internals/ProcessorTopologyFactories.java |   7 +-
 .../internals/RecordCollectorTest.java        |  21 +-
 .../processor/internals/StreamTaskTest.java   |   6 +-
 18 files changed, 569 insertions(+), 93 deletions(-)
 create mode 100644 streams/src/test/java/org/apache/kafka/streams/integration/GlobalStateReprocessTest.java

diff --git a/streams/src/main/java/org/apache/kafka/streams/StreamsBuilder.java b/streams/src/main/java/org/apache/kafka/streams/StreamsBuilder.java
index 91e679d8e5170..c9f77d3c4cc12 100644
--- a/streams/src/main/java/org/apache/kafka/streams/StreamsBuilder.java
+++ b/streams/src/main/java/org/apache/kafka/streams/StreamsBuilder.java
@@ -534,9 +534,7 @@ public synchronized StreamsBuilder addStateStore(final StoreBuilder<?> builder)
      * of the input topic.
      * <p>
      * The provided {@link org.apache.kafka.streams.processor.ProcessorSupplier} will be used to create an {@link ProcessorNode} that will receive all
-     * records forwarded from the {@link SourceNode}. NOTE: you should not use the {@code Processor} to insert transformed records into
-     * the global state store. This store uses the source topic as changelog and during restore will insert records directly
-     * from the source.
+     * records forwarded from the {@link SourceNode}.
      * This {@link ProcessorNode} should be used to keep the {@link StateStore} up-to-date.
      * The default {@link TimestampExtractor} as specified in the {@link StreamsConfig config} is used.
      * <p>
@@ -567,7 +565,8 @@ public synchronized <K, V> StreamsBuilder addGlobalStore(final StoreBuilder<?> s
             new StoreBuilderWrapper(storeBuilder),
             topic,
             new ConsumedInternal<>(consumed),
-            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get())
+            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get()),
+            true
         );
         return this;
     }
@@ -585,9 +584,6 @@ public synchronized <K, V> StreamsBuilder addGlobalStore(final StoreBuilder<?> s
      * The supplier should always generate a new instance. Creating a single {@link Processor} object
      * and returning the same object reference in {@link ProcessorSupplier#get()} is a
      * violation of the supplier pattern and leads to runtime exceptions.
-     * NOTE: you should not use the {@link Processor} to insert transformed records into
-     * the global state store. This store uses the source topic as changelog and during restore will insert records directly
-     * from the source.
      * This {@link Processor} should be used to keep the {@link StateStore} up-to-date.
      * The default {@link TimestampExtractor} as specified in the {@link StreamsConfig config} is used.
      * <p>
@@ -611,7 +607,8 @@ public synchronized <KIn, VIn> StreamsBuilder addGlobalStore(final StoreBuilder<
             new StoreBuilderWrapper(storeBuilder),
             topic,
             new ConsumedInternal<>(consumed),
-            stateUpdateSupplier
+            stateUpdateSupplier,
+            true
         );
         return this;
     }
diff --git a/streams/src/main/java/org/apache/kafka/streams/Topology.java b/streams/src/main/java/org/apache/kafka/streams/Topology.java
index d1f0d1eb8bc8e..0a1ac5c54e875 100644
--- a/streams/src/main/java/org/apache/kafka/streams/Topology.java
+++ b/streams/src/main/java/org/apache/kafka/streams/Topology.java
@@ -865,7 +865,8 @@ public synchronized <K, V> Topology addGlobalStore(final StoreBuilder<?> storeBu
             valueDeserializer,
             topic,
             processorName,
-            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get())
+            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get()),
+            true
         );
         return this;
     }
@@ -917,7 +918,8 @@ public synchronized <K, V> Topology addGlobalStore(final StoreBuilder<?> storeBu
             valueDeserializer,
             topic,
             processorName,
-            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get())
+            () -> ProcessorAdapter.adapt(stateUpdateSupplier.get()),
+            true
         );
         return this;
     }
@@ -960,7 +962,8 @@ public synchronized <KIn, VIn> Topology addGlobalStore(final StoreBuilder<?> sto
             valueDeserializer,
             topic,
             processorName,
-            stateUpdateSupplier
+            stateUpdateSupplier,
+            true
         );
         return this;
     }
@@ -1005,7 +1008,8 @@ public synchronized <KIn, VIn> Topology addGlobalStore(final StoreBuilder<?> sto
             valueDeserializer,
             topic,
             processorName,
-            stateUpdateSupplier
+            stateUpdateSupplier,
+            true
         );
         return this;
     }
diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java
index 26b4f648fca66..ba072e535d39f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java
@@ -41,6 +41,7 @@
 import org.apache.kafka.streams.kstream.internals.graph.TableSourceNode;
 import org.apache.kafka.streams.kstream.internals.graph.VersionedSemanticsGraphNode;
 import org.apache.kafka.streams.kstream.internals.graph.WindowedStreamProcessorNode;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
 import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder;
 import org.apache.kafka.streams.processor.internals.StoreFactory;
 import org.apache.kafka.streams.state.KeyValueStore;
@@ -221,7 +222,8 @@ public synchronized void addStateStore(final StoreFactory builder) {
     public synchronized <KIn, VIn> void addGlobalStore(final StoreFactory storeFactory,
                                                        final String topic,
                                                        final ConsumedInternal<KIn, VIn> consumed,
-                                                       final org.apache.kafka.streams.processor.api.ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier) {
+                                                       final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier,
+                                                       final boolean reprocessOnRestore) {
         // explicitly disable logging for global stores
         storeFactory.withLoggingDisabled();
 
@@ -235,7 +237,8 @@ public synchronized <KIn, VIn> void addGlobalStore(final StoreFactory storeFacto
             topic,
             consumed,
             processorName,
-            stateUpdateSupplier
+            stateUpdateSupplier,
+            reprocessOnRestore
         );
 
         addGraphNode(root, globalStoreNode);
diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/GlobalStoreNode.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/GlobalStoreNode.java
index 7b7709cd6566c..9bc3aae11c316 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/GlobalStoreNode.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/GlobalStoreNode.java
@@ -29,6 +29,7 @@ public class GlobalStoreNode<KIn, VIn, S extends StateStore> extends StateStoreN
     private final ConsumedInternal<KIn, VIn> consumed;
     private final String processorName;
     private final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier;
+    private final boolean reprocessOnRestore;
 
 
     public GlobalStoreNode(final StoreFactory storeBuilder,
@@ -36,7 +37,8 @@ public GlobalStoreNode(final StoreFactory storeBuilder,
                            final String topic,
                            final ConsumedInternal<KIn, VIn> consumed,
                            final String processorName,
-                           final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier) {
+                           final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier,
+                           final boolean reprocessOnRestore) {
 
         super(storeBuilder);
         this.sourceName = sourceName;
@@ -44,6 +46,7 @@ public GlobalStoreNode(final StoreFactory storeBuilder,
         this.consumed = consumed;
         this.processorName = processorName;
         this.stateUpdateSupplier = stateUpdateSupplier;
+        this.reprocessOnRestore = reprocessOnRestore;
     }
 
     @Override
@@ -56,7 +59,8 @@ public void writeToTopology(final InternalTopologyBuilder topologyBuilder) {
                                        consumed.valueDeserializer(),
                                        topic,
                                        processorName,
-                                       stateUpdateSupplier);
+                                       stateUpdateSupplier,
+                                       reprocessOnRestore);
 
     }
 
@@ -66,6 +70,7 @@ public String toString() {
                "sourceName='" + sourceName + '\'' +
                ", topic='" + topic + '\'' +
                ", processorName='" + processorName + '\'' +
+               ", reprocessOnRestore='" + reprocessOnRestore + '\'' +
                "} ";
     }
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/TableSourceNode.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/TableSourceNode.java
index e5c5db58e3d4d..b75a466cb8aa8 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/TableSourceNode.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/graph/TableSourceNode.java
@@ -105,7 +105,8 @@ public void writeToTopology(final InternalTopologyBuilder topologyBuilder) {
                 consumedInternal().valueDeserializer(),
                 topicName,
                 processorParameters.processorName(),
-                (ProcessorSupplier<K, V, Void, Void>) processorParameters.processorSupplier()
+                (ProcessorSupplier<K, V, Void, Void>) processorParameters.processorSupplier(),
+                false
             );
         } else {
             topologyBuilder.addSource(consumedInternal().offsetResetPolicy(),
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImpl.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImpl.java
index f64599218188c..d42e90bf33f4c 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImpl.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImpl.java
@@ -28,6 +28,7 @@
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.errors.DeserializationExceptionHandler;
 import org.apache.kafka.streams.errors.ProcessorStateException;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.processor.CommitCallback;
@@ -35,6 +36,8 @@
 import org.apache.kafka.streams.processor.StateRestoreListener;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.StateStoreContext;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.processor.internals.Task.TaskType;
 import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
 import org.apache.kafka.streams.state.internals.RecordConverter;
@@ -53,8 +56,10 @@
 import java.util.Set;
 import java.util.function.Supplier;
 
+import static org.apache.kafka.streams.processor.internals.RecordDeserializer.handleDeserializationFailure;
 import static org.apache.kafka.streams.processor.internals.StateManagerUtil.CHECKPOINT_FILE_NAME;
 import static org.apache.kafka.streams.processor.internals.StateManagerUtil.converterForStore;
+import static org.apache.kafka.streams.processor.internals.metrics.TaskMetrics.droppedRecordsSensor;
 
 /**
  * This class is responsible for the initialization, restoration, closing, flushing etc
@@ -77,8 +82,8 @@ public class GlobalStateManagerImpl implements GlobalStateManager {
     private final Set<String> globalStoreNames = new HashSet<>();
     private final Set<String> globalNonPersistentStoresTopics = new HashSet<>();
     private final FixedOrderMap<String, Optional<StateStore>> globalStores = new FixedOrderMap<>();
-
     private InternalProcessorContext globalProcessorContext;
+    private DeserializationExceptionHandler deserializationExceptionHandler;
 
     public GlobalStateManagerImpl(final LogContext logContext,
                                   final Time time,
@@ -116,6 +121,7 @@ public GlobalStateManagerImpl(final LogContext logContext,
             config.getLong(StreamsConfig.POLL_MS_CONFIG) + requestTimeoutMs
         );
         taskTimeoutMs = config.getLong(StreamsConfig.TASK_TIMEOUT_MS_CONFIG);
+        deserializationExceptionHandler = config.defaultDeserializationExceptionHandler();
     }
 
     @Override
@@ -203,13 +209,23 @@ public void registerStore(final StateStore store,
         );
 
         try {
-            restoreState(
-                stateRestoreCallback,
-                topicPartitions,
-                highWatermarks,
-                store.name(),
-                converterForStore(store)
-            );
+            final Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>> reprocessFactory = topology
+                .storeNameToReprocessOnRestore().getOrDefault(store.name(), Optional.empty());
+            if (reprocessFactory.isPresent()) {
+                reprocessState(
+                    topicPartitions,
+                    highWatermarks,
+                    reprocessFactory.get(),
+                    store.name());
+            } else {
+                restoreState(
+                    stateRestoreCallback,
+                    topicPartitions,
+                    highWatermarks,
+                    store.name(),
+                    converterForStore(store)
+                );
+            }
         } finally {
             globalConsumer.unsubscribe();
         }
@@ -237,6 +253,97 @@ private List<TopicPartition> topicPartitionsForStore(final StateStore store) {
         return topicPartitions;
     }
 
+    //Visible for testing
+    public void setDeserializationExceptionHandler(final DeserializationExceptionHandler deserializationExceptionHandler) {
+        this.deserializationExceptionHandler = deserializationExceptionHandler;
+    }
+
+    @SuppressWarnings("unchecked")
+    private void reprocessState(final List<TopicPartition> topicPartitions,
+                                final Map<TopicPartition, Long> highWatermarks,
+                                final InternalTopologyBuilder.ReprocessFactory reprocessFactory,
+                                final String storeName) {
+        final Processor source = reprocessFactory.processorSupplier().get();
+        source.init(globalProcessorContext);
+
+        for (final TopicPartition topicPartition : topicPartitions) {
+            long currentDeadline = NO_DEADLINE;
+
+            globalConsumer.assign(Collections.singletonList(topicPartition));
+            long offset;
+            final Long checkpoint = checkpointFileCache.get(topicPartition);
+            if (checkpoint != null) {
+                globalConsumer.seek(topicPartition, checkpoint);
+                offset = checkpoint;
+            } else {
+                globalConsumer.seekToBeginning(Collections.singletonList(topicPartition));
+                offset = getGlobalConsumerOffset(topicPartition);
+            }
+            final Long highWatermark = highWatermarks.get(topicPartition);
+            stateRestoreListener.onRestoreStart(topicPartition, storeName, offset, highWatermark);
+
+            long restoreCount = 0L;
+
+            while (offset < highWatermark) {
+                // we add `request.timeout.ms` to `poll.ms` because `poll.ms` might be too short
+                // to give a fetch request a fair chance to actually complete and we don't want to
+                // start `task.timeout.ms` too early
+                //
+                // TODO with https://issues.apache.org/jira/browse/KAFKA-10315 we can just call
+                //      `poll(pollMS)` without adding the request timeout and do a more precise
+                //      timeout handling
+                final ConsumerRecords<byte[], byte[]> records = globalConsumer.poll(pollMsPlusRequestTimeout);
+                if (records.isEmpty()) {
+                    currentDeadline = maybeUpdateDeadlineOrThrow(currentDeadline);
+                } else {
+                    currentDeadline = NO_DEADLINE;
+                }
+
+                for (final ConsumerRecord<byte[], byte[]> record : records.records(topicPartition)) {
+                    final ProcessorRecordContext recordContext =
+                        new ProcessorRecordContext(
+                            record.timestamp(),
+                            record.offset(),
+                            record.partition(),
+                            record.topic(),
+                            record.headers());
+                    globalProcessorContext.setRecordContext(recordContext);
+
+                    try {
+                        if (record.key() != null) {
+                            source.process(new Record<>(
+                                reprocessFactory.keyDeserializer().deserialize(record.topic(), record.key()),
+                                reprocessFactory.valueDeserializer().deserialize(record.topic(), record.value()),
+                                record.timestamp(),
+                                record.headers()));
+                            restoreCount++;
+                        }
+                    } catch (final Exception deserializationException) {
+                        handleDeserializationFailure(
+                            deserializationExceptionHandler,
+                            globalProcessorContext,
+                            deserializationException,
+                            record,
+                            log,
+                            droppedRecordsSensor(
+                                Thread.currentThread().getName(),
+                                globalProcessorContext.taskId().toString(),
+                                globalProcessorContext.metrics()
+                            )
+                        );
+                    }
+                }
+
+                offset = getGlobalConsumerOffset(topicPartition);
+
+                stateRestoreListener.onBatchRestored(topicPartition, storeName, offset, restoreCount);
+            }
+            stateRestoreListener.onRestoreEnd(topicPartition, storeName, restoreCount);
+            checkpointFileCache.put(topicPartition, offset);
+
+        }
+    }
+
     private void restoreState(final StateRestoreCallback stateRestoreCallback,
                               final List<TopicPartition> topicPartitions,
                               final Map<TopicPartition, Long> highWatermarks,
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
index 87dc930154210..6870f09a827f2 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
@@ -123,6 +123,9 @@ public InternalTopologyBuilder(final TopologyConfig topologyConfigs) {
     // map from changelog topic name to its corresponding state store.
     private final Map<String, String> changelogTopicToStore = new HashMap<>();
 
+    // map of store name to restore behavior
+    private final Map<String, Optional<ReprocessFactory<?, ?, ?, ?>>> storeNameToReprocessOnRestore = new HashMap<>();
+
     // all global topics
     private final Set<String> globalTopics = new HashSet<>();
 
@@ -154,6 +157,32 @@ public InternalTopologyBuilder(final TopologyConfig topologyConfigs) {
 
     private boolean hasPersistentStores = false;
 
+    public static class ReprocessFactory<KIn, VIn, KOut, VOut> {
+
+        private final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier;
+        private final Deserializer<KIn> keyDeserializer;
+        private final Deserializer<VIn> valueDeserializer;
+
+        private ReprocessFactory(final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier,
+                                 final Deserializer<KIn> key,
+                                 final Deserializer<VIn> value) {
+            this.processorSupplier = processorSupplier;
+            this.keyDeserializer = key;
+            this.valueDeserializer = value;
+        }
+        public ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier() {
+            return processorSupplier;
+        }
+
+        public Deserializer<KIn> keyDeserializer() {
+            return keyDeserializer;
+        }
+
+        public Deserializer<VIn> valueDeserializer() {
+            return valueDeserializer;
+        }
+    }
+
     private static abstract class NodeFactory<KIn, VIn, KOut, VOut> {
         final String name;
         final String[] predecessors;
@@ -584,7 +613,8 @@ public final <KIn, VIn> void addGlobalStore(final StoreFactory storeFactory,
                                                 final Deserializer<VIn> valueDeserializer,
                                                 final String topic,
                                                 final String processorName,
-                                                final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier) {
+                                                final ProcessorSupplier<KIn, VIn, Void, Void> stateUpdateSupplier,
+                                                final boolean reprocessOnRestore) {
         Objects.requireNonNull(storeFactory, "store builder must not be null");
         ApiUtils.checkSupplier(stateUpdateSupplier);
         validateGlobalStoreArguments(sourceName,
@@ -613,6 +643,10 @@ public final <KIn, VIn> void addGlobalStore(final StoreFactory storeFactory,
             keyDeserializer,
             valueDeserializer)
         );
+        storeNameToReprocessOnRestore.put(storeFactory.name(),
+            reprocessOnRestore ?
+                Optional.of(new ReprocessFactory<>(stateUpdateSupplier, keyDeserializer, valueDeserializer))
+                : Optional.empty());
         nodeToSourceTopics.put(sourceName, Arrays.asList(topics));
         nodeGrouper.add(sourceName);
         nodeFactory.addStateStore(storeFactory.name());
@@ -996,7 +1030,8 @@ private ProcessorTopology build(final Set<String> nodeGroup) {
                                      new ArrayList<>(stateStoreMap.values()),
                                      new ArrayList<>(globalStateStores.values()),
                                      storeToChangelogTopic,
-                                     repartitionTopics);
+                                     repartitionTopics,
+                                     storeNameToReprocessOnRestore);
     }
 
     private void buildSinkNode(final Map<String, ProcessorNode<?, ?, ?, ?>> processorMap,
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorTopology.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorTopology.java
index d2383c7adbb1c..f8f6ec146de21 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorTopology.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorTopology.java
@@ -26,6 +26,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 public class ProcessorTopology {
@@ -42,6 +43,7 @@ public class ProcessorTopology {
     // the following contains entries for the entire topology, eg stores that do not belong to this ProcessorTopology
     private final List<StateStore> globalStateStores;
     private final Map<String, String> storeToChangelogTopic;
+    private final Map<String, Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>>> storeNameToReprocessOnRestore;
 
     public ProcessorTopology(final List<ProcessorNode<?, ?, ?, ?>> processorNodes,
                              final Map<String, SourceNode<?, ?>> sourceNodesByTopic,
@@ -49,7 +51,8 @@ public ProcessorTopology(final List<ProcessorNode<?, ?, ?, ?>> processorNodes,
                              final List<StateStore> stateStores,
                              final List<StateStore> globalStateStores,
                              final Map<String, String> storeToChangelogTopic,
-                             final Set<String> repartitionTopics) {
+                             final Set<String> repartitionTopics,
+                             final Map<String, Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>>> storeNameToReprocessOnRestore) {
         this.processorNodes = Collections.unmodifiableList(processorNodes);
         this.sourceNodesByTopic = new HashMap<>(sourceNodesByTopic);
         this.sinksByTopic = Collections.unmodifiableMap(sinksByTopic);
@@ -57,6 +60,7 @@ public ProcessorTopology(final List<ProcessorNode<?, ?, ?, ?>> processorNodes,
         this.globalStateStores = Collections.unmodifiableList(globalStateStores);
         this.storeToChangelogTopic = Collections.unmodifiableMap(storeToChangelogTopic);
         this.repartitionTopics = Collections.unmodifiableSet(repartitionTopics);
+        this.storeNameToReprocessOnRestore = storeNameToReprocessOnRestore;
 
         this.terminalNodes = new HashSet<>();
         for (final ProcessorNode<?, ?, ?, ?> node : processorNodes) {
@@ -103,6 +107,10 @@ public List<StateStore> stateStores() {
         return stateStores;
     }
 
+    public Map<String, Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>>> storeNameToReprocessOnRestore() {
+        return storeNameToReprocessOnRestore;
+    }
+
     public List<StateStore> globalStateStores() {
         return Collections.unmodifiableList(globalStateStores);
     }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordDeserializer.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordDeserializer.java
index 6ffe9551a3508..9c56a138acecd 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordDeserializer.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordDeserializer.java
@@ -29,7 +29,7 @@
 
 import static org.apache.kafka.streams.StreamsConfig.DEFAULT_DESERIALIZATION_EXCEPTION_HANDLER_CLASS_CONFIG;
 
-class RecordDeserializer {
+public class RecordDeserializer {
     private final Logger log;
     private final SourceNode<?, ?> sourceNode;
     private final Sensor droppedRecordsSensor;
@@ -68,40 +68,49 @@ ConsumerRecord<Object, Object> deserialize(final ProcessorContext<?, ?> processo
                 Optional.empty()
             );
         } catch (final Exception deserializationException) {
-            final DeserializationExceptionHandler.DeserializationHandlerResponse response;
-            try {
-                response = deserializationExceptionHandler.handle(
-                    (InternalProcessorContext<?, ?>) processorContext,
-                    rawRecord,
-                    deserializationException);
-            } catch (final Exception fatalUserException) {
-                log.error(
-                    "Deserialization error callback failed after deserialization error for record {}",
-                    rawRecord,
-                    deserializationException);
-                throw new StreamsException("Fatal user code error in deserialization error callback", fatalUserException);
-            }
+            handleDeserializationFailure(deserializationExceptionHandler, processorContext, deserializationException, rawRecord, log, droppedRecordsSensor);
+            return null; //  'handleDeserializationFailure' would either throw or swallow -- if we swallow we need to skip the record by returning 'null'
+        }
+    }
 
-            if (response == DeserializationExceptionHandler.DeserializationHandlerResponse.FAIL) {
-                throw new StreamsException("Deserialization exception handler is set to fail upon" +
-                    " a deserialization error. If you would rather have the streaming pipeline" +
-                    " continue after a deserialization error, please set the " +
-                    DEFAULT_DESERIALIZATION_EXCEPTION_HANDLER_CLASS_CONFIG + " appropriately.",
-                    deserializationException);
-            } else {
-                log.warn(
-                    "Skipping record due to deserialization error. topic=[{}] partition=[{}] offset=[{}]",
-                    rawRecord.topic(),
-                    rawRecord.partition(),
-                    rawRecord.offset(),
-                    deserializationException
-                );
-                droppedRecordsSensor.record();
-                return null;
-            }
+    public static void handleDeserializationFailure(final DeserializationExceptionHandler deserializationExceptionHandler,
+                                                    final ProcessorContext<?, ?> processorContext,
+                                                    final Exception deserializationException,
+                                                    final ConsumerRecord<byte[], byte[]> rawRecord,
+                                                    final Logger log,
+                                                    final Sensor droppedRecordsSensor) {
+        final DeserializationExceptionHandler.DeserializationHandlerResponse response;
+        try {
+            response = deserializationExceptionHandler.handle(
+                (InternalProcessorContext<?, ?>) processorContext,
+                rawRecord,
+                deserializationException);
+        } catch (final Exception fatalUserException) {
+            log.error(
+                "Deserialization error callback failed after deserialization error for record {}",
+                rawRecord,
+                deserializationException);
+            throw new StreamsException("Fatal user code error in deserialization error callback", fatalUserException);
+        }
+        if (response == DeserializationExceptionHandler.DeserializationHandlerResponse.FAIL) {
+            throw new StreamsException("Deserialization exception handler is set to fail upon" +
+                " a deserialization error. If you would rather have the streaming pipeline" +
+                " continue after a deserialization error, please set the " +
+                DEFAULT_DESERIALIZATION_EXCEPTION_HANDLER_CLASS_CONFIG + " appropriately.",
+                deserializationException);
+        } else {
+            log.warn(
+                "Skipping record due to deserialization error. topic=[{}] partition=[{}] offset=[{}]",
+                rawRecord.topic(),
+                rawRecord.partition(),
+                rawRecord.offset(),
+                deserializationException
+            );
+            droppedRecordsSensor.record();
         }
     }
 
+
     SourceNode<?, ?> sourceNode() {
         return sourceNode;
     }
diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
index 652db52f0c73b..b33e307259b76 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
@@ -76,6 +76,8 @@ public void init(final ProcessorContext context,
                 false
             );
             // register the store
+            open = true;
+
             context.register(
                 root,
                 (RecordBatchingStateRestoreCallback) records -> {
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/GlobalStateReprocessTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalStateReprocessTest.java
new file mode 100644
index 0000000000000..20e7631857ec0
--- /dev/null
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalStateReprocessTest.java
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.streams.integration;
+
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.common.serialization.LongSerializer;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.StoreQueryParameters;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
+import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.processor.api.ContextualProcessor;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.streams.state.KeyValueStore;
+import org.apache.kafka.streams.state.QueryableStoreTypes;
+import org.apache.kafka.streams.state.ReadOnlyKeyValueStore;
+import org.apache.kafka.streams.state.Stores;
+import org.apache.kafka.streams.state.internals.KeyValueStoreBuilder;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+
+import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.safeUniqueTestName;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
+
+@Timeout(600)
+@Tag("integration")
+public class GlobalStateReprocessTest {
+    private static final int NUM_BROKERS = 1;
+    private static final Properties BROKER_CONFIG;
+
+    static {
+        BROKER_CONFIG = new Properties();
+        BROKER_CONFIG.put("transaction.state.log.replication.factor", (short) 1);
+        BROKER_CONFIG.put("transaction.state.log.min.isr", 1);
+    }
+
+    public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(NUM_BROKERS, BROKER_CONFIG);
+
+    @BeforeAll
+    public static void startCluster() throws IOException {
+        CLUSTER.start();
+    }
+
+    @AfterAll
+    public static void closeCluster() {
+        CLUSTER.stop();
+    }
+
+
+    private final MockTime mockTime = CLUSTER.time;
+    private final String globalStore = "globalStore";
+    private StreamsBuilder builder;
+    private Properties streamsConfiguration;
+    private KafkaStreams kafkaStreams;
+    private String globalStoreTopic;
+
+
+    @BeforeEach
+    public void before(final TestInfo testInfo) throws Exception {
+        builder = new StreamsBuilder();
+
+        createTopics();
+        streamsConfiguration = new Properties();
+        final String safeTestName = safeUniqueTestName(testInfo);
+        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, "app-" + safeTestName);
+        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
+        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
+        streamsConfiguration.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
+        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
+
+        final KeyValueStoreBuilder<String, Long> storeBuilder = new KeyValueStoreBuilder<>(
+            Stores.persistentKeyValueStore(globalStore),
+            Serdes.String(),
+            Serdes.Long(),
+            mockTime);
+
+        final ProcessorSupplier<String, Long, Void, Void> processorSupplier;
+        processorSupplier = () -> new ContextualProcessor<String, Long, Void, Void>() {
+            @Override
+            public void process(final Record<String, Long> record) {
+                final KeyValueStore<String, Long> stateStore =
+                    context().getStateStore(storeBuilder.name());
+                stateStore.put(
+                    record.key() + "- this is the right value.",
+                    record.value()
+                );
+            }
+        };
+
+        builder.addGlobalStore(
+            storeBuilder,
+            globalStoreTopic,
+            Consumed.with(Serdes.String(), Serdes.Long()),
+            processorSupplier
+        );
+    }
+
+    @AfterEach
+    public void after() throws Exception {
+        if (kafkaStreams != null) {
+            kafkaStreams.close();
+        }
+        IntegrationTestUtils.purgeLocalStreamsState(streamsConfiguration);
+    }
+
+    @Test
+    public void shouldReprocessWithUserProvidedStore() throws Exception {
+        kafkaStreams = new KafkaStreams(builder.build(), streamsConfiguration);
+        populateTopics(globalStoreTopic);
+
+        kafkaStreams.start();
+
+        TestUtils.waitForCondition(
+            () -> !storeContents(kafkaStreams).isEmpty(),
+            30000,
+            "Has not processed record within 30 seconds");
+
+        assertThat(storeContents(kafkaStreams).get(0), containsString("- this is the right value."));
+
+
+        kafkaStreams.close();
+        kafkaStreams.cleanUp();
+
+        kafkaStreams = new KafkaStreams(builder.build(), streamsConfiguration);
+        kafkaStreams.start();
+
+        TestUtils.waitForCondition(
+            () -> !storeContents(kafkaStreams).isEmpty(),
+            30000,
+            "Has not processed record within 30 seconds");
+
+        assertThat(storeContents(kafkaStreams).get(0), containsString("- this is the right value."));
+    }
+
+    private void createTopics() throws Exception {
+        globalStoreTopic = "global-store-topic";
+        CLUSTER.createTopic(globalStoreTopic);
+    }
+
+    private void populateTopics(final String topicName) throws Exception {
+        IntegrationTestUtils.produceKeyValuesSynchronously(
+            topicName,
+            Collections.singletonList(new KeyValue<>("A", 1L)),
+            TestUtils.producerConfig(
+                CLUSTER.bootstrapServers(),
+                StringSerializer.class,
+                LongSerializer.class,
+                new Properties()),
+            mockTime);
+    }
+
+    private List<String> storeContents(final KafkaStreams streams) {
+        final ArrayList<String> keySet = new ArrayList<>();
+        final ReadOnlyKeyValueStore<String, Long> keyValueStore =
+            streams.store(StoreQueryParameters.fromNameAndType(globalStore, QueryableStoreTypes.keyValueStore()));
+        final KeyValueIterator<String, Long> range = keyValueStore.reverseAll();
+        while (range.hasNext()) {
+            keySet.add(range.next().key);
+        }
+        range.close();
+        return keySet;
+    }
+}
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
index 18db670d90d9c..e486565ddf102 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
@@ -30,13 +30,14 @@
 import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
 import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
 import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.processor.api.ContextualProcessor;
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
 import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.state.KeyValueStore;
 import org.apache.kafka.streams.state.Stores;
 import org.apache.kafka.streams.state.internals.KeyValueStoreBuilder;
-import org.apache.kafka.test.MockApiProcessorSupplier;
 import org.apache.kafka.test.TestUtils;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -126,11 +127,24 @@ public void before(final TestInfo testInfo) throws Exception {
             Serdes.Long(),
             mockTime);
 
+        final ProcessorSupplier<String, Long, Void, Void> processorSupplier;
+        processorSupplier = () -> new ContextualProcessor<String, Long, Void, Void>() {
+            @Override
+            public void process(final Record<String, Long> record) {
+                final KeyValueStore<String, Long> stateStore =
+                    context().getStateStore(storeBuilder.name());
+                stateStore.put(
+                    record.key(),
+                    record.value()
+                );
+            }
+        };
+
         builder.addGlobalStore(
             storeBuilder,
             globalStoreTopic,
             Consumed.with(Serdes.String(), Serdes.Long()),
-            new MockApiProcessorSupplier<>()
+            processorSupplier
         );
 
         builder
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImplTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImplTest.java
index 0644f1b705a5b..25e54de816566 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImplTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImplTest.java
@@ -23,15 +23,20 @@
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.errors.DeserializationExceptionHandler;
+import org.apache.kafka.streams.errors.LogAndContinueExceptionHandler;
 import org.apache.kafka.streams.errors.ProcessorStateException;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.processor.StateRestoreCallback;
 import org.apache.kafka.streams.processor.StateStore;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
 import org.apache.kafka.streams.state.TimestampedBytesStore;
 import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
 import org.apache.kafka.streams.state.internals.WrappedStateStore;
@@ -55,6 +60,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -75,6 +81,9 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class GlobalStateManagerImplTest {
 
@@ -86,30 +95,37 @@ public class GlobalStateManagerImplTest {
     private final String storeName2 = "t2-store";
     private final String storeName3 = "t3-store";
     private final String storeName4 = "t4-store";
+    private final String storeName5 = "t5-store";
     private final TopicPartition t1 = new TopicPartition("t1", 1);
     private final TopicPartition t2 = new TopicPartition("t2", 1);
     private final TopicPartition t3 = new TopicPartition("t3", 1);
     private final TopicPartition t4 = new TopicPartition("t4", 1);
+    private final TopicPartition t5 = new TopicPartition("t5", 1);
     private GlobalStateManagerImpl stateManager;
     private StateDirectory stateDirectory;
     private StreamsConfig streamsConfig;
-    private NoOpReadOnlyStore<Object, Object> store1, store2, store3, store4;
+    private NoOpReadOnlyStore<Object, Object> store1, store2, store3, store4, store5;
     private MockConsumer<byte[], byte[]> consumer;
     private File checkpointFile;
     private ProcessorTopology topology;
     private InternalMockProcessorContext processorContext;
+    private Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>> optionalMockReprocessFactory;
+    private DeserializationExceptionHandler deserializationExceptionHandler;
 
     static ProcessorTopology withGlobalStores(final List<StateStore> stateStores,
-                                              final Map<String, String> storeToChangelogTopic) {
+                                              final Map<String, String> storeToChangelogTopic,
+                                              final Map<String, Optional<InternalTopologyBuilder.ReprocessFactory<?, ?, ?, ?>>> reprocessFactoryMap) {
         return new ProcessorTopology(Collections.emptyList(),
                                      Collections.emptyMap(),
                                      Collections.emptyMap(),
                                      Collections.emptyList(),
                                      stateStores,
                                      storeToChangelogTopic,
-                                     Collections.emptySet());
+                                     Collections.emptySet(),
+                                     reprocessFactoryMap);
     }
 
+    @SuppressWarnings("unchecked")
     @Before
     public void before() {
         final Map<String, String> storeToTopic = new HashMap<>();
@@ -118,14 +134,25 @@ public void before() {
         storeToTopic.put(storeName2, t2.topic());
         storeToTopic.put(storeName3, t3.topic());
         storeToTopic.put(storeName4, t4.topic());
+        storeToTopic.put(storeName5, t5.topic());
 
         store1 = new NoOpReadOnlyStore<>(storeName1, true);
         store2 = new ConverterStore<>(storeName2, true);
         store3 = new NoOpReadOnlyStore<>(storeName3);
         store4 = new NoOpReadOnlyStore<>(storeName4);
-
-        topology = withGlobalStores(asList(store1, store2, store3, store4), storeToTopic);
-
+        store5 = new NoOpReadOnlyStore<>(storeName5);
+
+        optionalMockReprocessFactory = mock(Optional.class);
+        when(optionalMockReprocessFactory.isPresent()).thenReturn(false);
+        topology = withGlobalStores(asList(store1, store2, store3, store4, store5), storeToTopic,
+            mkMap(
+                mkEntry(storeName1, Optional.empty()),
+                mkEntry(storeName2, Optional.empty()),
+                mkEntry(storeName3, Optional.empty()),
+                mkEntry(storeName4, Optional.empty()),
+                mkEntry(storeName5, optionalMockReprocessFactory)
+            )
+        );
         streamsConfig = new StreamsConfig(new Properties() {
             {
                 put(StreamsConfig.APPLICATION_ID_CONFIG, "appId");
@@ -223,7 +250,7 @@ public void shouldInitializeStateStores() {
     @Test
     public void shouldReturnInitializedStoreNames() {
         final Set<String> storeNames = stateManager.initialize();
-        assertEquals(Utils.mkSet(storeName1, storeName2, storeName3, storeName4), storeNames);
+        assertEquals(Utils.mkSet(storeName1, storeName2, storeName3, storeName4, storeName5), storeNames);
     }
 
     @Test
@@ -691,7 +718,7 @@ public synchronized long position(final TopicPartition partition) {
                 return numberOfCalls.incrementAndGet();
             }
         };
-        initializeConsumer(0, 0, t1, t2, t3, t4);
+        initializeConsumer(0, 0, t1, t2, t3, t4, t5);
 
         streamsConfig = new StreamsConfig(mkMap(
             mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
@@ -857,7 +884,7 @@ public synchronized long position(final TopicPartition partition) {
                 return numberOfCalls.incrementAndGet();
             }
         };
-        initializeConsumer(0, 0, t1, t2, t3, t4);
+        initializeConsumer(0, 0, t1, t2, t3, t4, t5);
 
         streamsConfig = new StreamsConfig(mkMap(
             mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
@@ -1018,7 +1045,7 @@ public synchronized long position(final TopicPartition partition) {
                 throw new TimeoutException("KABOOM!");
             }
         };
-        initializeConsumer(0, 0, t1, t2, t3, t4);
+        initializeConsumer(0, 0, t1, t2, t3, t4, t5);
 
         streamsConfig = new StreamsConfig(mkMap(
             mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
@@ -1092,6 +1119,44 @@ public synchronized ConsumerRecords<byte[], byte[]> poll(final Duration timeout)
         assertThat(time.milliseconds() - startTime, equalTo(331_100L));
     }
 
+    @SuppressWarnings("unchecked")
+    private void setUpReprocessing() {
+        final InternalTopologyBuilder.ReprocessFactory reprocessFactory = mock(InternalTopologyBuilder.ReprocessFactory.class);
+        final ProcessorSupplier processorSupplier = mock(ProcessorSupplier.class);
+        final Processor processor = mock(Processor.class);
+        final Deserializer deserializer = mock(Deserializer.class);
+
+        when(optionalMockReprocessFactory.isPresent()).thenReturn(true);
+        when(optionalMockReprocessFactory.get()).thenReturn(reprocessFactory);
+        when(reprocessFactory.processorSupplier()).thenReturn(processorSupplier);
+        when(processorSupplier.get()).thenReturn(processor);
+        when(reprocessFactory.keyDeserializer()).thenReturn(deserializer);
+        when(reprocessFactory.valueDeserializer()).thenReturn(deserializer);
+        when(deserializer.deserialize(any(), any())).thenThrow(new StreamsException("fail"));
+    }
+
+    @Test
+    public void shouldFailOnDeserializationErrorsWhenReprocessing() {
+        setUpReprocessing();
+        initializeConsumer(2, 0, t5);
+
+        stateManager.initialize();
+
+        assertThrows(StreamsException.class, () -> stateManager.registerStore(store5, stateRestoreCallback, null));
+    }
+
+    @Test
+    public void shouldSkipOnDeserializationErrorsWhenReprocessing() {
+        stateManager.setDeserializationExceptionHandler(new LogAndContinueExceptionHandler());
+        setUpReprocessing();
+        initializeConsumer(2, 0, t5);
+
+        stateManager.initialize();
+
+        stateManager.registerStore(store5, stateRestoreCallback, null);
+        assertEquals(0, stateRestoreCallback.restored.size());
+    }
+
     private void writeCorruptCheckpoint() throws IOException {
         final File checkpointFile = new File(stateManager.baseDir(), StateManagerUtil.CHECKPOINT_FILE_NAME);
         try (final OutputStream stream = Files.newOutputStream(checkpointFile.toPath())) {
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStreamThreadTest.java
index 5214a0989f1d5..c1b83aaef85ca 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStreamThreadTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/GlobalStreamThreadTest.java
@@ -113,7 +113,9 @@ public void process(final Record<Object, Object> record) {
             null,
             GLOBAL_STORE_TOPIC_NAME,
             "processorName",
-            processorSupplier);
+            processorSupplier,
+            false
+        );
 
         baseDirectoryName = TestUtils.tempDirectory().getAbsolutePath();
         final HashMap<String, Object> properties = new HashMap<>();
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
index 63f00b8cb7f5f..bbb4625d56440 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
@@ -221,7 +221,9 @@ public void testAddGlobalStoreWithBadSupplier() {
                         null,
                         "globalTopic",
                         "global-processor",
-                () -> processor)
+                () -> processor,
+                false
+            )
         );
         assertThat(exception.getMessage(), containsString("#get() must return a new object each time it is called."));
     }
@@ -329,7 +331,8 @@ public void testPatternSourceTopicsWithGlobalTopics() {
             null,
             "globalTopic",
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
         builder.initializeSubscription();
 
@@ -353,7 +356,8 @@ public void testNameSourceTopicsWithGlobalTopics() {
             null,
             "globalTopic",
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
         builder.initializeSubscription();
 
@@ -465,7 +469,8 @@ public void shouldNotAllowToAddStoresWithSameNameWhenFirstStoreIsGlobal() {
             null,
             "global-topic",
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
 
         final TopologyException exception = assertThrows(
@@ -496,7 +501,8 @@ public void shouldNotAllowToAddStoresWithSameNameWhenSecondStoreIsGlobal() {
                 null,
                 "global-topic",
                 "global-processor",
-                new MockApiProcessorSupplier<>()
+                new MockApiProcessorSupplier<>(),
+                false
             )
         );
 
@@ -521,7 +527,8 @@ public void shouldNotAllowToAddGlobalStoresWithSameName() {
             null,
             "global-topic",
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
 
         final TopologyException exception = assertThrows(
@@ -534,7 +541,8 @@ public void shouldNotAllowToAddGlobalStoresWithSameName() {
                 null,
                 "global-topic",
                 "global-processor-2",
-                new MockApiProcessorSupplier<>()
+                new MockApiProcessorSupplier<>(),
+                false
             )
         );
 
@@ -729,7 +737,8 @@ public void shouldAllowIncrementalBuilds() {
             null,
             "globalTopic",
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
         newNodeGroups = builder.nodeGroups();
         assertNotEquals(oldNodeGroups, newNodeGroups);
@@ -1153,7 +1162,8 @@ public void shouldNotAllowToAddGlobalStoreWithSourceNameEqualsProcessorName() {
             null,
             "anyTopicName",
             sameNameForSourceAndProcessor,
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         ));
     }
 
@@ -1297,7 +1307,8 @@ public void shouldConnectGlobalStateStoreToInputTopic() {
             null,
             globalTopic,
             "global-processor",
-            new MockApiProcessorSupplier<>()
+            new MockApiProcessorSupplier<>(),
+            false
         );
         builder.initializeSubscription();
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyFactories.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyFactories.java
index ccc78a48731c5..68b16e261aebe 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyFactories.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyFactories.java
@@ -35,7 +35,8 @@ public static ProcessorTopology with(final List<ProcessorNode<?, ?, ?, ?>> proce
                                      stateStoresByName,
                                      Collections.emptyList(),
                                      storeToChangelogTopic,
-                                     Collections.emptySet());
+                                     Collections.emptySet(),
+                                     Collections.emptyMap());
     }
 
     static ProcessorTopology withLocalStores(final List<StateStore> stateStores,
@@ -46,7 +47,9 @@ static ProcessorTopology withLocalStores(final List<StateStore> stateStores,
                                      stateStores,
                                      Collections.emptyList(),
                                      storeToChangelogTopic,
-                                     Collections.emptySet());
+                                     Collections.emptySet(),
+                                     Collections.emptyMap());
+
     }
 
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
index bf943d5b5e6d1..60503c46895a8 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
@@ -168,7 +168,8 @@ public void setup() {
             emptyList(),
             emptyList(),
             emptyMap(),
-            emptySet()
+            emptySet(),
+            emptyMap()
         );
         collector = new RecordCollectorImpl(
             logContext,
@@ -333,7 +334,8 @@ public Optional<Set<Integer>> partitions(final String topic, final String key, f
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
@@ -400,7 +402,8 @@ public Optional<Set<Integer>> partitions(final String topic, final String key, f
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
@@ -467,7 +470,8 @@ public Optional<Set<Integer>> partitions(final String topic, final String key, f
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
@@ -546,7 +550,8 @@ public Optional<Set<Integer>> partitions(final String topic, final String key, f
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
@@ -599,7 +604,8 @@ public void shouldUseDefaultPartitionerAsPartitionReturnsNull() {
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
@@ -649,7 +655,8 @@ public void shouldUseDefaultPartitionerAsStreamPartitionerIsNull() {
                 emptyList(),
                 emptyList(),
                 emptyMap(),
-                emptySet()
+                emptySet(),
+                emptyMap()
         );
         collector = new RecordCollectorImpl(
                 logContext,
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
index eaafef53f538d..dea8e0efcbc69 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
@@ -212,7 +212,8 @@ private static ProcessorTopology withRepartitionTopics(final List<ProcessorNode<
                                      emptyList(),
                                      emptyList(),
                                      emptyMap(),
-                                     repartitionTopics);
+                                     repartitionTopics,
+                                     emptyMap());
     }
 
     private static ProcessorTopology withSources(final List<ProcessorNode<?, ?, ?, ?>> processorNodes,
@@ -223,7 +224,8 @@ private static ProcessorTopology withSources(final List<ProcessorNode<?, ?, ?, ?
                                      emptyList(),
                                      emptyList(),
                                      emptyMap(),
-                                     Collections.emptySet());
+                                     Collections.emptySet(),
+                                     emptyMap());
     }
 
     private static StreamsConfig createConfig() {

From 71bcac3b6ada9872dcdec48f72372d6c5b041c0a Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Thu, 28 Mar 2024 13:27:41 -0700
Subject: [PATCH 226/258] MINOR: AbstractConfig cleanup (#15597)

Signed-off-by: Greg Harris <greg.harris@aiven.io>

Reviewers: Chris Egerton <chrise@aiven.io>, Mickael Maison <mickael.maison@gmail.com>, Omnia G H Ibrahim <o.g.h.ibrahim@gmail.com>, Matthias J. Sax <matthias@confluent.io>
---
 .../kafka/common/config/AbstractConfig.java   |  51 ++++++--
 .../org/apache/kafka/common/utils/Utils.java  |  18 ++-
 .../common/config/AbstractConfigTest.java     | 122 ++++++++++++++----
 .../connect/mirror/MirrorClientConfig.java    |   2 +-
 .../kafka/connect/cli/AbstractConnectCli.java |   2 +-
 .../kafka/connect/runtime/WorkerConfig.java   |   8 +-
 .../runtime/rest/RestServerConfig.java        |   2 +-
 .../kafka/connect/runtime/WorkerTest.java     |   1 +
 .../admin/BrokerApiVersionsCommand.scala      |   2 +-
 .../controller/PartitionStateMachine.scala    |   2 +-
 .../kafka/server/DynamicBrokerConfig.scala    |   6 +-
 .../main/scala/kafka/server/KafkaConfig.scala |   2 +-
 .../PartitionStateMachineTest.scala           |   4 +-
 13 files changed, 166 insertions(+), 56 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/config/AbstractConfig.java b/clients/src/main/java/org/apache/kafka/common/config/AbstractConfig.java
index 84bae97a03a8c..aeb7f07a29c7b 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/AbstractConfig.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/AbstractConfig.java
@@ -25,6 +25,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -33,6 +34,8 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 /**
  * A convenient base class for configurations to extend.
@@ -58,6 +61,8 @@ public class AbstractConfig {
 
     private final ConfigDef definition;
 
+    public static final String AUTOMATIC_CONFIG_PROVIDERS_PROPERTY = "org.apache.kafka.automatic.config.providers";
+
     public static final String CONFIG_PROVIDERS_CONFIG = "config.providers";
 
     private static final String CONFIG_PROVIDERS_PARAM = ".param.";
@@ -101,14 +106,11 @@ public class AbstractConfig {
      *                            the constructor to resolve any variables in {@code originals}; may be null or empty
      * @param doLog               whether the configurations should be logged
      */
-    @SuppressWarnings({"unchecked", "this-escape"})
+    @SuppressWarnings({"this-escape"})
     public AbstractConfig(ConfigDef definition, Map<?, ?> originals, Map<String, ?> configProviderProps, boolean doLog) {
-        /* check that all the keys are really strings */
-        for (Map.Entry<?, ?> entry : originals.entrySet())
-            if (!(entry.getKey() instanceof String))
-                throw new ConfigException(entry.getKey().toString(), entry.getValue(), "Key must be a string.");
+        Map<String, Object> originalMap = Utils.castToStringObjectMap(originals);
 
-        this.originals = resolveConfigVariables(configProviderProps, (Map<String, Object>) originals);
+        this.originals = resolveConfigVariables(configProviderProps, originalMap);
         this.values = definition.parse(this.originals);
         Map<String, Object> configUpdates = postProcessParsedConfig(Collections.unmodifiableMap(this.values));
         for (Map.Entry<String, Object> update : configUpdates.entrySet()) {
@@ -521,6 +523,7 @@ private Map<String, String> extractPotentialVariables(Map<?, ?> configMap) {
     private Map<String, ?> resolveConfigVariables(Map<String, ?> configProviderProps, Map<String, Object> originals) {
         Map<String, String> providerConfigString;
         Map<String, ?> configProperties;
+        Predicate<String> classNameFilter;
         Map<String, Object> resolvedOriginals = new HashMap<>();
         // As variable configs are strings, parse the originals and obtain the potential variable configs.
         Map<String, String> indirectVariables = extractPotentialVariables(originals);
@@ -529,11 +532,13 @@ private Map<String, String> extractPotentialVariables(Map<?, ?> configMap) {
         if (configProviderProps == null || configProviderProps.isEmpty()) {
             providerConfigString = indirectVariables;
             configProperties = originals;
+            classNameFilter = automaticConfigProvidersFilter();
         } else {
             providerConfigString = extractPotentialVariables(configProviderProps);
             configProperties = configProviderProps;
+            classNameFilter = ignored -> true;
         }
-        Map<String, ConfigProvider> providers = instantiateConfigProviders(providerConfigString, configProperties);
+        Map<String, ConfigProvider> providers = instantiateConfigProviders(providerConfigString, configProperties, classNameFilter);
 
         if (!providers.isEmpty()) {
             ConfigTransformer configTransformer = new ConfigTransformer(providers);
@@ -547,6 +552,17 @@ private Map<String, String> extractPotentialVariables(Map<?, ?> configMap) {
         return new ResolvingMap<>(resolvedOriginals, originals);
     }
 
+    private Predicate<String> automaticConfigProvidersFilter() {
+        String systemProperty = System.getProperty(AUTOMATIC_CONFIG_PROVIDERS_PROPERTY);
+        if (systemProperty == null) {
+            return ignored -> true;
+        } else {
+            return Arrays.stream(systemProperty.split(","))
+                    .map(String::trim)
+                    .collect(Collectors.toSet())::contains;
+        }
+    }
+
     private Map<String, Object> configProviderProperties(String configProviderPrefix, Map<String, ?> providerConfigProperties) {
         Map<String, Object> result = new HashMap<>();
         for (Map.Entry<String, ?> entry : providerConfigProperties.entrySet()) {
@@ -567,9 +583,14 @@ private Map<String, Object> configProviderProperties(String configProviderPrefix
      *
      * @param indirectConfigs          The map of potential variable configs
      * @param providerConfigProperties The map of config provider configs
-     * @return map map of config provider name and its instance.
+     * @param classNameFilter          Filter for config provider class names
+     * @return map of config provider name and its instance.
      */
-    private Map<String, ConfigProvider> instantiateConfigProviders(Map<String, String> indirectConfigs, Map<String, ?> providerConfigProperties) {
+    private Map<String, ConfigProvider> instantiateConfigProviders(
+            Map<String, String> indirectConfigs,
+            Map<String, ?> providerConfigProperties,
+            Predicate<String> classNameFilter
+    ) {
         final String configProviders = indirectConfigs.get(CONFIG_PROVIDERS_CONFIG);
 
         if (configProviders == null || configProviders.isEmpty()) {
@@ -580,9 +601,15 @@ private Map<String, ConfigProvider> instantiateConfigProviders(Map<String, Strin
 
         for (String provider : configProviders.split(",")) {
             String providerClass = providerClassProperty(provider);
-            if (indirectConfigs.containsKey(providerClass))
-                providerMap.put(provider, indirectConfigs.get(providerClass));
-
+            if (indirectConfigs.containsKey(providerClass)) {
+                String providerClassName = indirectConfigs.get(providerClass);
+                if (classNameFilter.test(providerClassName)) {
+                    providerMap.put(provider, providerClassName);
+                } else {
+                    throw new ConfigException(providerClassName + " is not allowed. Update System property '"
+                            + AUTOMATIC_CONFIG_PROVIDERS_PROPERTY + "' to allow " + providerClassName);
+                }
+            }
         }
         // Instantiate Config Providers
         Map<String, ConfigProvider> configProviderInstances = new HashMap<>();
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
index c316b7a181662..3fd3833de1e44 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
@@ -1502,13 +1502,23 @@ public static <K, V> Map<K, V> filterMap(final Map<K, V> map, final Predicate<En
      * @return a map including all elements in properties
      */
     public static Map<String, Object> propsToMap(Properties properties) {
-        Map<String, Object> map = new HashMap<>(properties.size());
-        for (Map.Entry<Object, Object> entry : properties.entrySet()) {
+        return castToStringObjectMap(properties);
+    }
+
+    /**
+     * Cast a map with arbitrary type keys to be keyed on String.
+     * @param inputMap A map with unknown type keys
+     * @return A map with the same contents as the input map, but with String keys
+     * @throws ConfigException if any key is not a String
+     */
+    public static Map<String, Object> castToStringObjectMap(Map<?, ?> inputMap) {
+        Map<String, Object> map = new HashMap<>(inputMap.size());
+        for (Map.Entry<?, ?> entry : inputMap.entrySet()) {
             if (entry.getKey() instanceof String) {
                 String k = (String) entry.getKey();
-                map.put(k, properties.get(k));
+                map.put(k, entry.getValue());
             } else {
-                throw new ConfigException(entry.getKey().toString(), entry.getValue(), "Key must be a string.");
+                throw new ConfigException(String.valueOf(entry.getKey()), entry.getValue(), "Key must be a string.");
             }
         }
         return map;
diff --git a/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java b/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java
index 5859dc1dc1278..bf018aebbfcb9 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.provider.FileConfigProvider;
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.metrics.FakeMetricsReporter;
 import org.apache.kafka.common.metrics.JmxReporter;
@@ -26,7 +27,10 @@
 import org.apache.kafka.common.security.TestSecurityConfig;
 import org.apache.kafka.common.config.provider.MockVaultConfigProvider;
 import org.apache.kafka.common.config.provider.MockFileConfigProvider;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.MockConsumerInterceptor;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.util.Arrays;
@@ -46,6 +50,23 @@
 
 public class AbstractConfigTest {
 
+    private String propertyValue;
+
+    @BeforeEach
+    public void setup() {
+        propertyValue = System.getProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY);
+        System.clearProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY);
+    }
+
+    @AfterEach
+    public void teardown() {
+        if (propertyValue != null) {
+            System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, propertyValue);
+        } else {
+            System.clearProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY);
+        }
+    }
+
     @Test
     public void testConfiguredInstances() {
         testValidInputs("    ");
@@ -254,12 +275,7 @@ private void testInvalidInputs(String configValue) {
         Properties props = new Properties();
         props.put(TestConfig.METRIC_REPORTER_CLASSES_CONFIG, configValue);
         TestConfig config = new TestConfig(props);
-        try {
-            config.getConfiguredInstances(TestConfig.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
-            fail("Expected a config exception due to invalid props :" + props);
-        } catch (KafkaException e) {
-            // this is good
-        }
+        assertThrows(KafkaException.class, () -> config.getConfiguredInstances(TestConfig.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class));
     }
 
     @Test
@@ -349,16 +365,6 @@ protected Class<?> findClass(String name) throws ClassNotFoundException {
         }
     }
 
-    @SuppressWarnings("unchecked")
-    public Map<String, ?> convertPropertiesToMap(Map<?, ?> props) {
-        for (Map.Entry<?, ?> entry : props.entrySet()) {
-            if (!(entry.getKey() instanceof String))
-                throw new ConfigException(entry.getKey().toString(), entry.getValue(),
-                    "Key must be a string.");
-        }
-        return (Map<String, ?>) props;
-    }
-
     @Test
     public void testOriginalWithOverrides() {
         Properties props = new Properties();
@@ -389,6 +395,43 @@ public void testOriginalsWithConfigProvidersProps() {
         MockFileConfigProvider.assertClosed(id);
     }
 
+    @Test
+    public void testOriginalsWithConfigProvidersPropsExcluded() {
+        System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, MockVaultConfigProvider.class.getName() + " , " + FileConfigProvider.class.getName());
+        Properties props = new Properties();
+
+        // Test Case: Config provider that is not an allowed class
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        props.put("config.providers.file.param.testId", id);
+        props.put("prefix.ssl.truststore.location.number", 5);
+        props.put("sasl.kerberos.service.name", "service name");
+        props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
+        props.put("sasl.kerberos.password", "${file:/usr/kerberos:password}");
+        assertThrows(ConfigException.class, () -> new TestIndirectConfigResolution(props, Collections.emptyMap()));
+    }
+
+    @Test
+    public void testOriginalsWithConfigProvidersPropsIncluded() {
+        System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, MockFileConfigProvider.class.getName() + " , " + FileConfigProvider.class.getName());
+        Properties props = new Properties();
+
+        // Test Case: Config provider that is an allowed class
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class", MockFileConfigProvider.class.getName());
+        String id = UUID.randomUUID().toString();
+        props.put("config.providers.file.param.testId", id);
+        props.put("prefix.ssl.truststore.location.number", 5);
+        props.put("sasl.kerberos.service.name", "service name");
+        props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
+        props.put("sasl.kerberos.password", "${file:/usr/kerberos:password}");
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, Collections.emptyMap());
+        assertEquals("testKey", config.originals().get("sasl.kerberos.key"));
+        assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
+        MockFileConfigProvider.assertClosed(id);
+    }
+
     @Test
     public void testConfigProvidersPropsAsParam() {
         // Test Case: Valid Test Case for ConfigProviders as a separate variable
@@ -400,7 +443,7 @@ public void testConfigProvidersPropsAsParam() {
         Properties props = new Properties();
         props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
         props.put("sasl.kerberos.password", "${file:/usr/kerberos:password}");
-        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, Utils.castToStringObjectMap(providers));
         assertEquals("testKey", config.originals().get("sasl.kerberos.key"));
         assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
         MockFileConfigProvider.assertClosed(id);
@@ -417,7 +460,7 @@ public void testImmutableOriginalsWithConfigProvidersProps() {
         Properties props = new Properties();
         props.put("sasl.kerberos.key", "${file:/usr/kerberos:key}");
         Map<?, ?> immutableMap = Collections.unmodifiableMap(props);
-        Map<String, ?> provMap = convertPropertiesToMap(providers);
+        Map<String, ?> provMap = Utils.castToStringObjectMap(providers);
         TestIndirectConfigResolution config = new TestIndirectConfigResolution(immutableMap, provMap);
         assertEquals("testKey", config.originals().get("sasl.kerberos.key"));
         MockFileConfigProvider.assertClosed(id);
@@ -437,7 +480,7 @@ public void testAutoConfigResolutionWithMultipleConfigProviders() {
         props.put("sasl.kerberos.password", "${file:/usr/kerberos:password}");
         props.put("sasl.truststore.key", "${vault:/usr/truststore:truststoreKey}");
         props.put("sasl.truststore.password", "${vault:/usr/truststore:truststorePassword}");
-        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, Utils.castToStringObjectMap(providers));
         assertEquals("testKey", config.originals().get("sasl.kerberos.key"));
         assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
         assertEquals("testTruststoreKey", config.originals().get("sasl.truststore.key"));
@@ -453,12 +496,33 @@ public void testAutoConfigResolutionWithInvalidConfigProviderClass() {
         props.put("config.providers.file.class",
             "org.apache.kafka.common.config.provider.InvalidConfigProvider");
         props.put("testKey", "${test:/foo/bar/testpath:testKey}");
-        try {
-            new TestIndirectConfigResolution(props);
-            fail("Expected a config exception due to invalid props :" + props);
-        } catch (KafkaException e) {
-            // this is good
-        }
+        assertThrows(KafkaException.class, () -> new TestIndirectConfigResolution(props));
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithInvalidConfigProviderClassExcluded() {
+        String invalidConfigProvider = "org.apache.kafka.common.config.provider.InvalidConfigProvider";
+        System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, "");
+        // Test Case: Any config provider specified while the system property is empty
+        Properties props = new Properties();
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class", invalidConfigProvider);
+        props.put("testKey", "${test:/foo/bar/testpath:testKey}");
+        KafkaException e = assertThrows(KafkaException.class, () -> new TestIndirectConfigResolution(props, Collections.emptyMap()));
+        assertTrue(e.getMessage().contains(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY));
+    }
+
+    @Test
+    public void testAutoConfigResolutionWithInvalidConfigProviderClassIncluded() {
+        String invalidConfigProvider = "org.apache.kafka.common.config.provider.InvalidConfigProvider";
+        System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, invalidConfigProvider);
+        // Test Case: Invalid config provider specified, but is also included in the system property
+        Properties props = new Properties();
+        props.put("config.providers", "file");
+        props.put("config.providers.file.class", invalidConfigProvider);
+        props.put("testKey", "${test:/foo/bar/testpath:testKey}");
+        KafkaException e = assertThrows(KafkaException.class, () -> new TestIndirectConfigResolution(props, Collections.emptyMap()));
+        assertFalse(e.getMessage().contains(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY));
     }
 
     @Test
@@ -496,13 +560,15 @@ public void testAutoConfigResolutionWithDuplicateConfigProvider() {
         props.put("config.providers", "file");
         props.put("config.providers.file.class", MockVaultConfigProvider.class.getName());
 
-        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, Utils.castToStringObjectMap(providers));
         assertEquals("${file:/usr/kerberos:key}", config.originals().get("sasl.kerberos.key"));
     }
 
     @Test
     public void testConfigProviderConfigurationWithConfigParams() {
-        // Test Case: Valid Test Case With Multiple ConfigProviders as a separate variable
+        // should have no effect
+        System.setProperty(AbstractConfig.AUTOMATIC_CONFIG_PROVIDERS_PROPERTY, MockFileConfigProvider.class.getName());
+        // Test Case: Specify a config provider not allowed, but passed via the trusted providers argument
         Properties providers = new Properties();
         providers.put("config.providers", "vault");
         providers.put("config.providers.vault.class", MockVaultConfigProvider.class.getName());
@@ -512,7 +578,7 @@ public void testConfigProviderConfigurationWithConfigParams() {
         props.put("sasl.truststore.key", "${vault:/usr/truststore:truststoreKey}");
         props.put("sasl.truststore.password", "${vault:/usr/truststore:truststorePassword}");
         props.put("sasl.truststore.location", "${vault:/usr/truststore:truststoreLocation}");
-        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, Utils.castToStringObjectMap(providers));
         assertEquals("/usr/vault", config.originals().get("sasl.truststore.location"));
     }
 
diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClientConfig.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClientConfig.java
index 477459895c56c..053e594fbeb1d 100644
--- a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClientConfig.java
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/MirrorClientConfig.java
@@ -76,7 +76,7 @@ public class MirrorClientConfig extends AbstractConfig {
     public static final String PRODUCER_CLIENT_PREFIX = "producer.";
 
     MirrorClientConfig(Map<?, ?> props) {
-        super(CONFIG_DEF, props, true);
+        super(CONFIG_DEF, props, Utils.castToStringObjectMap(props), true);
     }
 
     public ReplicationPolicy replicationPolicy() {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java
index de666c7bd60a2..c770a19624a1d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java
@@ -125,7 +125,7 @@ public Connect startConnect(Map<String, String> workerProps, String... extraArgs
 
         RestClient restClient = new RestClient(config);
 
-        ConnectRestServer restServer = new ConnectRestServer(config.rebalanceTimeout(), restClient, workerProps);
+        ConnectRestServer restServer = new ConnectRestServer(config.rebalanceTimeout(), restClient, config.originals());
         restServer.initializeServer();
 
         URI advertisedUrl = restServer.advertisedUrl();
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfig.java
index fe28918a29a2b..6de49ebd55084 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfig.java
@@ -438,9 +438,15 @@ public static PluginDiscoveryMode pluginDiscovery(Map<String, String> props) {
 
     @SuppressWarnings("this-escape")
     public WorkerConfig(ConfigDef definition, Map<String, String> props) {
-        super(definition, props);
+        super(definition, props, Utils.castToStringObjectMap(props), true);
         logInternalConverterRemovalWarnings(props);
         logPluginPathConfigProviderWarning(props);
     }
 
+    @Override
+    public Map<String, Object> originals() {
+        Map<String, Object> map = super.originals();
+        map.remove(AbstractConfig.CONFIG_PROVIDERS_CONFIG);
+        return map;
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServerConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServerConfig.java
index 0d6d06a4a596c..4b8b5acf93519 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServerConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestServerConfig.java
@@ -258,7 +258,7 @@ public Integer rebalanceTimeoutMs() {
     }
 
     protected RestServerConfig(ConfigDef configDef, Map<?, ?> props) {
-        super(configDef, props);
+        super(configDef, props, Utils.castToStringObjectMap(props), true);
     }
 
     // Visible for testing
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
index c09fa093fb32e..4579794a2c4ff 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTest.java
@@ -1167,6 +1167,7 @@ public void testAdminConfigsClientOverridesWithAllPolicy() {
 
         Map<String, Object> connConfig = Collections.singletonMap("metadata.max.age.ms", "10000");
         Map<String, String> expectedConfigs = new HashMap<>(workerProps);
+        expectedConfigs.remove(AbstractConfig.CONFIG_PROVIDERS_CONFIG);
         expectedConfigs.put("bootstrap.servers", "localhost:9092");
         expectedConfigs.put("client.id", "testid");
         expectedConfigs.put("metadata.max.age.ms", "10000");
diff --git a/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala b/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
index 45df36f7a1f87..6cb273f066f21 100644
--- a/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
+++ b/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
@@ -261,7 +261,7 @@ object BrokerApiVersionsCommand {
       config
     }
 
-    class AdminConfig(originals: Map[_,_]) extends AbstractConfig(AdminConfigDef, originals.asJava, false)
+    class AdminConfig(originals: Map[_,_]) extends AbstractConfig(AdminConfigDef, originals.asJava, Utils.castToStringObjectMap(originals.asJava), false)
 
     def create(props: Properties): AdminClient = create(props.asScala.toMap)
 
diff --git a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
index 5dedad426406b..51634291b6d56 100755
--- a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
@@ -486,7 +486,7 @@ class ZkPartitionStateMachine(config: KafkaConfig,
     } else {
       val (logConfigs, failed) = zkClient.getLogConfigs(
         partitionsWithNoLiveInSyncReplicas.iterator.map { case (partition, _) => partition.topic }.toSet,
-        config.originals()
+        config.extractLogConfigMap
       )
 
       partitionsWithNoLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
index fefa385558206..cae70f795321b 100755
--- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
@@ -197,7 +197,7 @@ object DynamicBrokerConfig {
 
   private[server] def resolveVariableConfigs(propsOriginal: Properties): Properties = {
     val props = new Properties
-    val config = new AbstractConfig(new ConfigDef(), propsOriginal, false)
+    val config = new AbstractConfig(new ConfigDef(), propsOriginal, Utils.castToStringObjectMap(propsOriginal), false)
     config.originals.forEach { (key, value) =>
       if (!key.startsWith(AbstractConfig.CONFIG_PROVIDERS_CONFIG)) {
         props.put(key, value)
@@ -740,13 +740,13 @@ class DynamicLogConfig(logManager: LogManager, server: KafkaBroker) extends Brok
     val originalLogConfig = logManager.currentDefaultConfig
     val originalUncleanLeaderElectionEnable = originalLogConfig.uncleanLeaderElectionEnable
     val newBrokerDefaults = new util.HashMap[String, Object](originalLogConfig.originals)
-    newConfig.valuesFromThisConfig.forEach { (k, v) =>
+    newConfig.extractLogConfigMap.forEach { (k, v) =>
       if (DynamicLogConfig.ReconfigurableConfigs.contains(k)) {
         DynamicLogConfig.KafkaConfigToLogConfigName.get(k).foreach { configName =>
           if (v == null)
              newBrokerDefaults.remove(configName)
           else
-            newBrokerDefaults.put(configName, v.asInstanceOf[AnyRef])
+            newBrokerDefaults.put(configName, v)
         }
       }
     }
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index c6d6ad83379c9..af0827d0e42f8 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -1291,7 +1291,7 @@ object KafkaConfig {
 }
 
 class KafkaConfig private(doLog: Boolean, val props: java.util.Map[_, _], dynamicConfigOverride: Option[DynamicBrokerConfig])
-  extends AbstractConfig(KafkaConfig.configDef, props, doLog) with Logging {
+  extends AbstractConfig(KafkaConfig.configDef, props, Utils.castToStringObjectMap(props), doLog) with Logging {
 
   def this(props: java.util.Map[_, _]) = this(true, KafkaConfig.populateSynonyms(props), None)
   def this(props: java.util.Map[_, _], doLog: Boolean) = this(doLog, KafkaConfig.populateSynonyms(props), None)
diff --git a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
index 10cbe58904564..183e8657e0d44 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
@@ -258,7 +258,7 @@ class PartitionStateMachineTest {
       .thenReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
         TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
-    when(mockZkClient.getLogConfigs(Set.empty, config.originals()))
+    when(mockZkClient.getLogConfigs(Set.empty, config.extractLogConfigMap))
       .thenReturn((Map(partition.topic -> new LogConfig(new Properties)), Map.empty[String, Exception]))
     val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withPartitionEpoch(2)
@@ -434,7 +434,7 @@ class PartitionStateMachineTest {
     }
     prepareMockToGetTopicPartitionsStatesRaw()
     def prepareMockToGetLogConfigs(): Unit = {
-      when(mockZkClient.getLogConfigs(Set.empty, config.originals())).thenReturn((Map.empty[String, LogConfig], Map.empty[String, Exception]))
+      when(mockZkClient.getLogConfigs(Set.empty, config.extractLogConfigMap)).thenReturn((Map.empty[String, LogConfig], Map.empty[String, Exception]))
     }
     prepareMockToGetLogConfigs()
 

From bf5e04e41626322eb60402266e7f28a262db4b4c Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Thu, 28 Mar 2024 20:07:42 -0700
Subject: [PATCH 227/258] KAFKA-16349: Prevent race conditions in Exit class
 from stopping test JVM (#15484)

Signed-off-by: Greg Harris <greg.harris@aiven.io>
Reviewers: Chris Egerton <chrise@aiven.io>
---
 checkstyle/import-control-server-common.xml   |   2 +-
 .../org/apache/kafka/common/utils/Exit.java   |  33 ++++-
 core/src/main/scala/kafka/utils/Exit.scala    |  20 +++
 .../kafka/utils/ShutdownableThreadTest.scala  |  69 ---------
 .../server/util/ShutdownableThreadTest.java   | 133 ++++++++++++++++++
 5 files changed, 185 insertions(+), 72 deletions(-)
 delete mode 100644 core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
 create mode 100644 server-common/src/test/java/org/apache/kafka/server/util/ShutdownableThreadTest.java

diff --git a/checkstyle/import-control-server-common.xml b/checkstyle/import-control-server-common.xml
index 2c5c652e979e8..5af34ca4a1e62 100644
--- a/checkstyle/import-control-server-common.xml
+++ b/checkstyle/import-control-server-common.xml
@@ -105,10 +105,10 @@
             <allow pkg="org.apache.kafka.server.util.json" />
 
             <allow class="org.apache.kafka.server.util.TopicFilter.IncludeList" />
+            <allow class="org.apache.kafka.test.TestUtils" />
             <subpackage name="timer">
                 <allow class="org.apache.kafka.server.util.MockTime" />
                 <allow class="org.apache.kafka.server.util.ShutdownableThread" />
-                <allow class="org.apache.kafka.test.TestUtils" />
             </subpackage>
         </subpackage>
     </subpackage>
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Exit.java b/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
index 6d503ca2ad5c3..1ef1ee15b5bf9 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
@@ -43,6 +43,14 @@ public interface ShutdownHookAdder {
             Runtime.getRuntime().addShutdownHook(new Thread(runnable));
     };
 
+    private static final Procedure NOOP_HALT_PROCEDURE = (statusCode, message) -> {
+        throw new IllegalStateException("Halt called after resetting procedures; possible race condition present in test");
+    };
+
+    private static final Procedure NOOP_EXIT_PROCEDURE = (statusCode, message) -> {
+        throw new IllegalStateException("Exit called after resetting procedures; possible race condition present in test");
+    };
+
     private volatile static Procedure exitProcedure = DEFAULT_EXIT_PROCEDURE;
     private volatile static Procedure haltProcedure = DEFAULT_HALT_PROCEDURE;
     private volatile static ShutdownHookAdder shutdownHookAdder = DEFAULT_SHUTDOWN_HOOK_ADDER;
@@ -67,26 +75,47 @@ public static void addShutdownHook(String name, Runnable runnable) {
         shutdownHookAdder.addShutdownHook(name, runnable);
     }
 
+    /**
+     * For testing only, do not call in main code.
+     */
     public static void setExitProcedure(Procedure procedure) {
         exitProcedure = procedure;
     }
 
+    /**
+     * For testing only, do not call in main code.
+     */
     public static void setHaltProcedure(Procedure procedure) {
         haltProcedure = procedure;
     }
 
+    /**
+     * For testing only, do not call in main code.
+     */
     public static void setShutdownHookAdder(ShutdownHookAdder shutdownHookAdder) {
         Exit.shutdownHookAdder = shutdownHookAdder;
     }
 
+    /**
+     * For testing only, do not call in main code.
+     * <p>Clears the procedure set in {@link #setExitProcedure(Procedure)}, but does not restore system default behavior of exiting the JVM.
+     */
     public static void resetExitProcedure() {
-        exitProcedure = DEFAULT_EXIT_PROCEDURE;
+        exitProcedure = NOOP_EXIT_PROCEDURE;
     }
 
+    /**
+     * For testing only, do not call in main code.
+     * <p>Clears the procedure set in {@link #setHaltProcedure(Procedure)}, but does not restore system default behavior of exiting the JVM.
+     */
     public static void resetHaltProcedure() {
-        haltProcedure = DEFAULT_HALT_PROCEDURE;
+        haltProcedure = NOOP_HALT_PROCEDURE;
     }
 
+    /**
+     * For testing only, do not call in main code.
+     * <p>Restores the system default shutdown hook behavior.
+     */
     public static void resetShutdownHookAdder() {
         shutdownHookAdder = DEFAULT_SHUTDOWN_HOOK_ADDER;
     }
diff --git a/core/src/main/scala/kafka/utils/Exit.scala b/core/src/main/scala/kafka/utils/Exit.scala
index eddd929af5547..84027100ec3a7 100644
--- a/core/src/main/scala/kafka/utils/Exit.scala
+++ b/core/src/main/scala/kafka/utils/Exit.scala
@@ -38,22 +38,42 @@ object Exit {
     JExit.addShutdownHook(name, () => shutdownHook)
   }
 
+  /**
+   * For testing only, do not call in main code.
+   */
   def setExitProcedure(exitProcedure: (Int, Option[String]) => Nothing): Unit =
     JExit.setExitProcedure(functionToProcedure(exitProcedure))
 
+  /**
+   * For testing only, do not call in main code.
+   */
   def setHaltProcedure(haltProcedure: (Int, Option[String]) => Nothing): Unit =
     JExit.setHaltProcedure(functionToProcedure(haltProcedure))
 
+  /**
+   * For testing only, do not call in main code.
+   */
   def setShutdownHookAdder(shutdownHookAdder: (String, => Unit) => Unit): Unit = {
     JExit.setShutdownHookAdder((name, runnable) => shutdownHookAdder(name, runnable.run()))
   }
 
+  /**
+   * For testing only, do not call in main code.
+   * <p>Clears the procedure set in [[setExitProcedure]], but does not restore system default behavior of exiting the JVM.
+   */
   def resetExitProcedure(): Unit =
     JExit.resetExitProcedure()
 
+  /**
+   * For testing only, do not call in main code.
+   * <p>Clears the procedure set in [[setHaltProcedure]], but does not restore system default behavior of exiting the JVM.
+   */
   def resetHaltProcedure(): Unit =
     JExit.resetHaltProcedure()
 
+  /**
+   * For testing only, do not call in main code.
+   */
   def resetShutdownHookAdder(): Unit =
     JExit.resetShutdownHookAdder()
 
diff --git a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala b/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
deleted file mode 100644
index da998903d2746..0000000000000
--- a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
-  * Licensed to the Apache Software Foundation (ASF) under one or more
-  * contributor license agreements.  See the NOTICE file distributed with
-  * this work for additional information regarding copyright ownership.
-  * The ASF licenses this file to You under the Apache License, Version 2.0
-  * (the "License"); you may not use this file except in compliance with
-  * the License.  You may obtain a copy of the License at
-  *
-  * http://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software
-  * distributed under the License is distributed on an "AS IS" BASIS,
-  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  * See the License for the specific language governing permissions and
-  * limitations under the License.
-  */
-package kafka.utils
-
-import java.util.concurrent.{CountDownLatch, TimeUnit}
-import org.apache.kafka.common.internals.FatalExitError
-import org.apache.kafka.server.util.ShutdownableThread
-import org.junit.jupiter.api.{AfterEach, Test}
-import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertTrue}
-
-class ShutdownableThreadTest {
-
-  @AfterEach
-  def tearDown(): Unit = Exit.resetExitProcedure()
-
-  @Test
-  def testShutdownWhenCalledAfterThreadStart(): Unit = {
-    @volatile var statusCodeOption: Option[Int] = None
-    Exit.setExitProcedure { (statusCode, _) =>
-      statusCodeOption = Some(statusCode)
-      // Sleep until interrupted to emulate the fact that `System.exit()` never returns
-      Thread.sleep(Long.MaxValue)
-      throw new AssertionError
-    }
-    val latch = new CountDownLatch(1)
-    val thread = new ShutdownableThread("shutdownable-thread-test") {
-      override def doWork(): Unit = {
-        latch.countDown()
-        throw new FatalExitError
-      }
-    }
-    thread.start()
-    assertTrue(latch.await(10, TimeUnit.SECONDS), "doWork was not invoked")
-
-    thread.shutdown()
-    TestUtils.waitUntilTrue(() => statusCodeOption.isDefined, "Status code was not set by exit procedure")
-    assertEquals(1, statusCodeOption.get)
-  }
-
-  @Test
-  def testIsThreadStarted(): Unit = {
-    val latch = new CountDownLatch(1)
-    val thread = new ShutdownableThread("shutdownable-thread-test") {
-      override def doWork(): Unit = {
-        latch.countDown()
-      }
-    }
-    assertFalse(thread.isStarted)
-    thread.start()
-    latch.await()
-    assertTrue(thread.isStarted)
-
-    thread.shutdown()
-  }
-}
diff --git a/server-common/src/test/java/org/apache/kafka/server/util/ShutdownableThreadTest.java b/server-common/src/test/java/org/apache/kafka/server/util/ShutdownableThreadTest.java
new file mode 100644
index 0000000000000..e59acc9767ec2
--- /dev/null
+++ b/server-common/src/test/java/org/apache/kafka/server/util/ShutdownableThreadTest.java
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.server.util;
+
+import org.apache.kafka.common.internals.FatalExitError;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class ShutdownableThreadTest {
+
+    @AfterEach
+    public void tearDown() {
+        Exit.resetExitProcedure();
+    }
+
+    @Test
+    public void testShutdownWhenCalledAfterThreadStart() throws InterruptedException {
+        AtomicReference<Optional<Integer>> statusCodeOption = new AtomicReference<>(Optional.empty());
+        Exit.setExitProcedure((statusCode, ignored) -> {
+            statusCodeOption.set(Optional.of(statusCode));
+            // Sleep until interrupted to emulate the fact that `System.exit()` never returns
+            Utils.sleep(Long.MAX_VALUE);
+            throw new AssertionError();
+        });
+        CountDownLatch latch = new CountDownLatch(1);
+        ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-test") {
+            @Override
+            public void doWork() {
+                latch.countDown();
+                throw new FatalExitError();
+            }
+        };
+        thread.start();
+        assertTrue(latch.await(10, TimeUnit.SECONDS), "doWork was not invoked");
+
+        thread.shutdown();
+        TestUtils.waitForCondition(() -> statusCodeOption.get().isPresent(), "Status code was not set by exit procedure");
+        assertEquals(1, statusCodeOption.get().get());
+    }
+
+    @Test
+    public void testIsThreadStarted() throws InterruptedException {
+        CountDownLatch latch = new CountDownLatch(1);
+        ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-test") {
+            @Override
+            public void doWork() {
+                latch.countDown();
+            }
+        };
+        assertFalse(thread.isStarted());
+        thread.start();
+        assertTrue(latch.await(10, TimeUnit.SECONDS), "doWork was not invoked");
+        assertTrue(thread.isStarted());
+
+        thread.shutdown();
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    public void testShutdownWhenTestTimesOut(boolean isInterruptible) {
+        // Mask the exit procedure so the contents of the "test" can't shut down the JVM
+        Exit.setExitProcedure((statusCode, ignored) -> {
+            throw new FatalExitError();
+        });
+        // This latch will be triggered only after the "test" finishes
+        CountDownLatch afterTest = new CountDownLatch(1);
+        try {
+            // This is the "test", which uses a ShutdownableThread with exit procedures masked.
+            CountDownLatch startupLatch = new CountDownLatch(1);
+            ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-timeout", isInterruptible) {
+                @Override
+                public void doWork() {
+                    // Tell the test that we finished starting, and are ready to shut down.
+                    startupLatch.countDown();
+                    if (isInterruptible) {
+                        // Swallow the interruption sent by the thread
+                        try {
+                            afterTest.await();
+                        } catch (InterruptedException ignored) {
+                        }
+                    }
+                    // Trigger a fatal exit, after the test has completed and the exit procedure masking has been cleaned up
+                    try {
+                        afterTest.await();
+                    } catch (InterruptedException ignored) {
+                    }
+                    throw new FatalExitError();
+                }
+            };
+            thread.start();
+            startupLatch.await();
+            // Interrupt ourselves, as if the test was interrupted for timing out
+            Thread.currentThread().interrupt();
+            thread.shutdown();
+            fail("Shutdown should have been interrupted");
+        } catch (InterruptedException ignored) {
+            // Swallow the interruption, so that the surrounding test passes
+        } finally {
+            Exit.resetExitProcedure();
+        }
+        // After the test has stopped waiting for the thread and cleaned up the Exit procedure, let the thread continue
+        afterTest.countDown();
+    }
+}

From 7c3a596688501223f2b53ed05fca7c8b9d39c3ca Mon Sep 17 00:00:00 2001
From: John Yu <yujuan476@gmail.com>
Date: Fri, 29 Mar 2024 20:28:19 +0800
Subject: [PATCH 228/258] KAFKA-16397 Use ByteBufferOutputStream to avoid array
 copy (#15589)

Reviewers: Apoorv Mittal <amittal@confluent.io>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../telemetry/internals/ClientTelemetryUtils.java      | 10 +++-------
 .../telemetry/internals/ClientTelemetryUtilsTest.java  |  5 +++--
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtils.java b/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtils.java
index 03ec43813043c..6f3cbd18d941b 100644
--- a/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtils.java
@@ -29,8 +29,6 @@
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -204,16 +202,14 @@ public static byte[] compress(byte[] raw, CompressionType compressionType) throw
     public static ByteBuffer decompress(byte[] metrics, CompressionType compressionType) {
         ByteBuffer data = ByteBuffer.wrap(metrics);
         try (InputStream in = compressionType.wrapForInput(data, RecordBatch.CURRENT_MAGIC_VALUE, BufferSupplier.create());
-            ByteArrayOutputStream out = new ByteArrayOutputStream()) {
-
+            ByteBufferOutputStream out = new ByteBufferOutputStream(512)) {
             byte[] bytes = new byte[data.capacity() * 2];
             int nRead;
             while ((nRead = in.read(bytes, 0, bytes.length)) != -1) {
                 out.write(bytes, 0, nRead);
             }
-
-            out.flush();
-            return ByteBuffer.wrap(out.toByteArray());
+            out.buffer().flip();
+            return out.buffer();
         } catch (IOException e) {
             throw new KafkaException("Failed to decompress metrics data", e);
         }
diff --git a/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtilsTest.java
index 27c49d412a9ce..7f9f18be6beb2 100644
--- a/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/telemetry/internals/ClientTelemetryUtilsTest.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.CompressionType;
+import org.apache.kafka.common.utils.Utils;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
@@ -132,9 +133,9 @@ public void testCompressDecompress(CompressionType compressionType) throws IOExc
         } else {
             assertArrayEquals(testString, compressed);
         }
-
         ByteBuffer decompressed = ClientTelemetryUtils.decompress(compressed, compressionType);
         assertNotNull(decompressed);
-        assertArrayEquals(testString, decompressed.array());
+        byte[] actualResult = Utils.toArray(decompressed);
+        assertArrayEquals(testString, actualResult);
     }
 }
\ No newline at end of file

From 7317159f0cf1bbf4e3e83d1729bb0e682a9ad05c Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Fri, 29 Mar 2024 08:51:42 -0700
Subject: [PATCH 229/258] KAFKA-16303: Add upgrade notes to 3.5.0, 3.5.2, and
 3.7.0 about MM2 offset translation (#15423)

Signed-off-by: Greg Harris <greg.harris@aiven.io>
Reviewers: Mickael Maison <mickael.maison@gmail.com>
---
 docs/upgrade.html | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/docs/upgrade.html b/docs/upgrade.html
index 6488ead802334..6e17d6fc4b279 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -102,6 +102,11 @@ <h5><a id="upgrade_370_notable" href="#upgrade_370_notable">Notable changes in 3
         <li>Kafka Streams ships multiple KIPs for IQv2 support.
             See the <a href="/{{version}}/documentation/streams/upgrade-guide#streams_api_changes_370">Kafka Streams upgrade section</a> for more details.
         </li>
+        <li>
+            In versions 3.5.0, 3.5.1, 3.5.2, 3.6.0, and 3.6.1, MirrorMaker 2 offset translation may not reach the end
+            of a replicated topic after the upstream consumers commit at the end of the source topic.
+            This was addressed in <a href="https://issues.apache.org/jira/browse/KAFKA-15906">KAFKA-15906</a>.
+        </li>
         <li>All the notable changes are present in the <a href="https://kafka.apache.org/blog#apache_kafka_370_release_announcement">blog post announcing the 3.7.0 release.</a>
         </li>
     </ul>
@@ -225,6 +230,11 @@ <h5><a id="upgrade_352_notable" href="#upgrade_352_notable">Notable changes in 3
         during rolling upgrade. This issue will impact the availability of the topic partition.
         See <a href="https://issues.apache.org/jira/browse/KAFKA-15353">KAFKA-15353</a> for more details.
     </li>
+    <li>
+        In 3.5.0 and 3.5.1, there was an issue where MirrorMaker 2 offset translation produced an earlier offset than needed,
+        substantially increasing the re-delivery of data when starting a consumer from the downstream consumer offsets.
+        See <a href="https://issues.apache.org/jira/browse/KAFKA-15202">KAFKA-15202</a> for more details.
+    </li>
 </ul>
 
 <h4><a id="upgrade_3_5_1" href="#upgrade_3_5_1">Upgrading to 3.5.1 from any version 0.8.x through 3.4.x</a></h4>
@@ -334,6 +344,12 @@ <h5><a id="upgrade_350_notable" href="#upgrade_350_notable">Notable changes in 3
             The 'kafka.tools' package is deprecated and will change to 'org.apache.kafka.tools' in the next major release.
             See <a href="https://issues.apache.org/jira/browse/KAFKA-14525">KAFKA-14525</a> for more details.
         </li>
+        <li>In versions earlier than 3.5.0 and 3.4.1, MirrorMaker 2 offset translation could incorrectly translate offsets for topics using compaction, transactional producers, and filter SMTs.
+            In 3.5.0 and 3.4.1, offset translation has changed for all topics in order to ensure at-least-once delivery when a consumer is failed-over to the translated offsets.
+            Translated offsets will be earlier than in previous versions, so consumers using downstream offsets may initially have more lag, and re-deliver more data after failing-over.
+            See <a href="https://issues.apache.org/jira/browse/KAFKA-12468">KAFKA-12468</a> and linked tickets, and <a href="https://github.com/apache/kafka/pull/13178">PR #13178</a> for more details.
+            Further improvements to the offset translation  are included in later releases to reduce the lag introduced by this change, so consider upgrading MM2 to the latest version available.
+        </li>
     </ul>
 
 <h4><a id="upgrade_3_4_0" href="#upgrade_3_4_0">Upgrading to 3.4.0 from any version 0.8.x through 3.3.x</a></h4>

From 2118d858053df9494ae0c314c3d32a9f7756e9ac Mon Sep 17 00:00:00 2001
From: Hector Geraldino <hgeraldino@gmail.com>
Date: Fri, 29 Mar 2024 18:39:36 -0400
Subject: [PATCH 230/258] KAFKA-16223 Replace EasyMock/PowerMock with Mockito
 for KafkaConfigBackingStoreTest (1/3) (#15520)

Reviewers: Greg Harris <greg.harris@aiven.io>
---
 .../storage/KafkaConfigBackingStore.java      |   7 +-
 .../KafkaConfigBackingStoreMockitoTest.java   | 503 ++++++++++++++++++
 .../storage/KafkaConfigBackingStoreTest.java  | 265 ---------
 3 files changed, 509 insertions(+), 266 deletions(-)
 create mode 100644 connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreMockitoTest.java

diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
index 735020173996c..48ae74b0813f4 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java
@@ -293,7 +293,7 @@ public static String LOGGER_CLUSTER_KEY(String namespace) {
     private final String topic;
     // Data is passed to the log already serialized. We use a converter to handle translating to/from generic Connect
     // format to serialized form
-    private final KafkaBasedLog<String, byte[]> configLog;
+    private KafkaBasedLog<String, byte[]> configLog;
     // Connector -> # of tasks
     final Map<String, Integer> connectorTaskCounts = new HashMap<>();
     // Connector and task configs: name or id -> config map
@@ -329,6 +329,11 @@ public static String LOGGER_CLUSTER_KEY(String namespace) {
     private final Map<String, Object> fencableProducerProps;
     private final Time time;
 
+    //VisibleForTesting
+    void setConfigLog(KafkaBasedLog<String, byte[]> configLog) {
+        this.configLog = configLog;
+    }
+
     @Deprecated
     public KafkaConfigBackingStore(Converter converter, DistributedConfig config, WorkerConfigTransformer configTransformer) {
         this(converter, config, configTransformer, null, "connect-distributed-");
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreMockitoTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreMockitoTest.java
new file mode 100644
index 0000000000000..37d1c55066f0b
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreMockitoTest.java
@@ -0,0 +1,503 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.storage;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.IsolationLevel;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.data.Field;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.runtime.RestartRequest;
+import org.apache.kafka.connect.runtime.TargetState;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.KafkaBasedLog;
+import org.apache.kafka.connect.util.TestFuture;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.mockito.stubbing.Answer;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.ISOLATION_LEVEL_CONFIG;
+import static org.apache.kafka.clients.producer.ProducerConfig.CLIENT_ID_CONFIG;
+import static org.apache.kafka.clients.producer.ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG;
+import static org.apache.kafka.clients.producer.ProducerConfig.TRANSACTIONAL_ID_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG;
+import static org.apache.kafka.connect.storage.KafkaConfigBackingStore.INCLUDE_TASKS_FIELD_NAME;
+import static org.apache.kafka.connect.storage.KafkaConfigBackingStore.ONLY_FAILED_FIELD_NAME;
+import static org.apache.kafka.connect.storage.KafkaConfigBackingStore.RESTART_KEY;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.StrictStubs.class)
+public class KafkaConfigBackingStoreMockitoTest {
+    private static final String CLIENT_ID_BASE = "test-client-id-";
+    private static final String TOPIC = "connect-configs";
+    private static final short TOPIC_REPLICATION_FACTOR = 5;
+    private static final Map<String, String> DEFAULT_CONFIG_STORAGE_PROPS = new HashMap<>();
+
+    static {
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.CONFIG_TOPIC_CONFIG, TOPIC);
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, Short.toString(TOPIC_REPLICATION_FACTOR));
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.GROUP_ID_CONFIG, "connect");
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "status-topic");
+        DEFAULT_CONFIG_STORAGE_PROPS.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "broker1:9092,broker2:9093");
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DEFAULT_CONFIG_STORAGE_PROPS.put(DistributedConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+    }
+
+    private static final List<String> CONNECTOR_IDS = Arrays.asList("connector1", "connector2");
+    private static final List<String> CONNECTOR_CONFIG_KEYS = Arrays.asList("connector-connector1", "connector-connector2");
+    private static final List<String> TARGET_STATE_KEYS = Arrays.asList("target-state-connector1", "target-state-connector2");
+
+
+    private static final String CONNECTOR_1_NAME = "connector1";
+    private static final String CONNECTOR_2_NAME = "connector2";
+    private static final List<String> RESTART_CONNECTOR_KEYS = Arrays.asList(RESTART_KEY(CONNECTOR_1_NAME), RESTART_KEY(CONNECTOR_2_NAME));
+
+    private static final Struct ONLY_FAILED_MISSING_STRUCT = new Struct(KafkaConfigBackingStore.RESTART_REQUEST_V0).put(INCLUDE_TASKS_FIELD_NAME, false);
+    private static final Struct INCLUDE_TASKS_MISSING_STRUCT = new Struct(KafkaConfigBackingStore.RESTART_REQUEST_V0).put(ONLY_FAILED_FIELD_NAME, true);
+    private static final List<Struct> RESTART_REQUEST_STRUCTS = Arrays.asList(
+            new Struct(KafkaConfigBackingStore.RESTART_REQUEST_V0).put(ONLY_FAILED_FIELD_NAME, true).put(INCLUDE_TASKS_FIELD_NAME, false),
+            ONLY_FAILED_MISSING_STRUCT,
+            INCLUDE_TASKS_MISSING_STRUCT);
+
+
+    // Need some placeholders -- the contents don't matter here, just that they are restored properly
+    private static final List<Map<String, String>> SAMPLE_CONFIGS = Arrays.asList(
+            Collections.singletonMap("config-key-one", "config-value-one"),
+            Collections.singletonMap("config-key-two", "config-value-two"),
+            Collections.singletonMap("config-key-three", "config-value-three")
+    );
+
+    // The exact format doesn't matter here since both conversions are mocked
+    private static final List<byte[]> CONFIGS_SERIALIZED = Arrays.asList(
+            "config-bytes-1".getBytes(), "config-bytes-2".getBytes(), "config-bytes-3".getBytes(),
+            "config-bytes-4".getBytes(), "config-bytes-5".getBytes(), "config-bytes-6".getBytes(),
+            "config-bytes-7".getBytes(), "config-bytes-8".getBytes(), "config-bytes-9".getBytes()
+    );
+
+    private static final List<byte[]> TARGET_STATES_SERIALIZED = Arrays.asList(
+            "started".getBytes(), "paused".getBytes(), "stopped".getBytes()
+    );
+    @Mock
+    private Converter converter;
+    @Mock
+    private ConfigBackingStore.UpdateListener configUpdateListener;
+    private Map<String, String> props = new HashMap<>(DEFAULT_CONFIG_STORAGE_PROPS);
+    private DistributedConfig config;
+    @Mock
+    KafkaBasedLog<String, byte[]> configLog;
+    @Mock
+    Future<RecordMetadata> producerFuture;
+    private KafkaConfigBackingStore configStorage;
+
+    private final ArgumentCaptor<String> capturedTopic = ArgumentCaptor.forClass(String.class);
+    @SuppressWarnings("unchecked")
+    private final ArgumentCaptor<Map<String, Object>> capturedConsumerProps = ArgumentCaptor.forClass(Map.class);
+    @SuppressWarnings("unchecked")
+    private final ArgumentCaptor<Map<String, Object>> capturedProducerProps = ArgumentCaptor.forClass(Map.class);
+    @SuppressWarnings("unchecked")
+    private final ArgumentCaptor<Supplier<TopicAdmin>> capturedAdminSupplier = ArgumentCaptor.forClass(Supplier.class);
+    private final ArgumentCaptor<NewTopic> capturedNewTopic = ArgumentCaptor.forClass(NewTopic.class);
+    @SuppressWarnings("unchecked")
+    private final ArgumentCaptor<Callback<ConsumerRecord<String, byte[]>>> capturedConsumedCallback = ArgumentCaptor.forClass(Callback.class);
+
+    private final MockTime time = new MockTime();
+    private long logOffset = 0;
+
+    private void createStore() {
+        config = Mockito.spy(new DistributedConfig(props));
+        doReturn("test-cluster").when(config).kafkaClusterId();
+        configStorage = Mockito.spy(
+                new KafkaConfigBackingStore(
+                        converter, config, null, null, CLIENT_ID_BASE, time)
+        );
+        configStorage.setConfigLog(configLog);
+        configStorage.setUpdateListener(configUpdateListener);
+    }
+
+    @Before
+    public void setUp() {
+        createStore();
+    }
+
+    @Test
+    public void testStartStop() {
+        props.put("config.storage.min.insync.replicas", "3");
+        props.put("config.storage.max.message.bytes", "1001");
+        createStore();
+
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+
+        verifyConfigure();
+        assertEquals(TOPIC, capturedTopic.getValue());
+        assertEquals("org.apache.kafka.common.serialization.StringSerializer", capturedProducerProps.getValue().get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG));
+        assertEquals("org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG));
+        assertEquals("org.apache.kafka.common.serialization.StringDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG));
+        assertEquals("org.apache.kafka.common.serialization.ByteArrayDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG));
+
+        assertEquals(TOPIC, capturedNewTopic.getValue().name());
+        assertEquals(1, capturedNewTopic.getValue().numPartitions());
+        assertEquals(TOPIC_REPLICATION_FACTOR, capturedNewTopic.getValue().replicationFactor());
+        assertEquals("3", capturedNewTopic.getValue().configs().get("min.insync.replicas"));
+        assertEquals("1001", capturedNewTopic.getValue().configs().get("max.message.bytes"));
+
+        configStorage.start();
+        configStorage.stop();
+
+        verify(configLog).start();
+        verify(configLog).stop();
+    }
+
+    @Test
+    public void testSnapshotCannotMutateInternalState() {
+        props.put("config.storage.min.insync.replicas", "3");
+        props.put("config.storage.max.message.bytes", "1001");
+        createStore();
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+        verifyConfigure();
+
+        configStorage.start();
+        ClusterConfigState snapshot = configStorage.snapshot();
+        assertNotSame(snapshot.connectorTaskCounts, configStorage.connectorTaskCounts);
+        assertNotSame(snapshot.connectorConfigs, configStorage.connectorConfigs);
+        assertNotSame(snapshot.connectorTargetStates, configStorage.connectorTargetStates);
+        assertNotSame(snapshot.taskConfigs, configStorage.taskConfigs);
+        assertNotSame(snapshot.connectorTaskCountRecords, configStorage.connectorTaskCountRecords);
+        assertNotSame(snapshot.connectorTaskConfigGenerations, configStorage.connectorTaskConfigGenerations);
+        assertNotSame(snapshot.connectorsPendingFencing, configStorage.connectorsPendingFencing);
+        assertNotSame(snapshot.inconsistentConnectors, configStorage.inconsistent);
+    }
+
+    @Test
+    public void testPutConnectorConfig() throws Exception {
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+        verifyConfigure();
+        configStorage.start();
+
+        // Null before writing
+        ClusterConfigState configState = configStorage.snapshot();
+        assertEquals(-1, configState.offset());
+        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
+
+        String configKey = CONNECTOR_CONFIG_KEYS.get(1);
+        String targetStateKey = TARGET_STATE_KEYS.get(1);
+
+        doAnswer(expectReadToEnd(Collections.singletonMap(CONNECTOR_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(0))))
+                .doAnswer(expectReadToEnd(Collections.singletonMap(CONNECTOR_CONFIG_KEYS.get(1), CONFIGS_SERIALIZED.get(1))))
+                // Config deletion
+                .doAnswer(expectReadToEnd(new LinkedHashMap<String, byte[]>() {{
+                            put(configKey, null);
+                            put(targetStateKey, null);
+                        }})
+                ).when(configLog).readToEnd();
+
+        // Writing should block until it is written and read back from Kafka
+        expectConvertWriteRead(
+                CONNECTOR_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
+                "properties", SAMPLE_CONFIGS.get(0));
+
+        configStorage.putConnectorConfig(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), null);
+        configState = configStorage.snapshot();
+
+        assertEquals(1, configState.offset());
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
+        verify(configUpdateListener).onConnectorConfigUpdate(CONNECTOR_IDS.get(0));
+
+        // Second should also block and all configs should still be available
+        expectConvertWriteRead(
+                CONNECTOR_CONFIG_KEYS.get(1), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(1),
+                "properties", SAMPLE_CONFIGS.get(1));
+
+        configStorage.putConnectorConfig(CONNECTOR_IDS.get(1), SAMPLE_CONFIGS.get(1), null);
+        configState = configStorage.snapshot();
+
+        assertEquals(2, configState.offset());
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertEquals(SAMPLE_CONFIGS.get(1), configState.connectorConfig(CONNECTOR_IDS.get(1)));
+        verify(configUpdateListener).onConnectorConfigUpdate(CONNECTOR_IDS.get(1));
+
+        // Config deletion
+        expectConvertWriteRead(configKey, KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, null, null, null);
+        expectConvertWriteRead(targetStateKey, KafkaConfigBackingStore.TARGET_STATE_V0, null, null, null);
+
+        // Deletion should remove the second one we added
+        configStorage.removeConnectorConfig(CONNECTOR_IDS.get(1));
+        configState = configStorage.snapshot();
+
+        assertEquals(4, configState.offset());
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
+        assertNull(configState.targetState(CONNECTOR_IDS.get(1)));
+        verify(configUpdateListener).onConnectorConfigRemove(CONNECTOR_IDS.get(1));
+
+        configStorage.stop();
+        verify(configLog).stop();
+    }
+
+    @Test
+    public void testPutConnectorConfigWithTargetState() throws Exception {
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+        verifyConfigure();
+        configStorage.start();
+
+        // Null before writing
+        ClusterConfigState configState = configStorage.snapshot();
+        assertEquals(-1, configState.offset());
+        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(0)));
+        assertNull(configState.targetState(CONNECTOR_IDS.get(0)));
+
+        doAnswer(expectReadToEnd(new LinkedHashMap<String, byte[]>() {{
+                    put(TARGET_STATE_KEYS.get(0), TARGET_STATES_SERIALIZED.get(2));
+                    put(CONNECTOR_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(0));
+                }})
+        ).when(configLog).readToEnd();
+
+        // We expect to write the target state first, followed by the config write and then a read to end
+        expectConvertWriteRead(
+                TARGET_STATE_KEYS.get(0), KafkaConfigBackingStore.TARGET_STATE_V1, TARGET_STATES_SERIALIZED.get(2),
+                "state.v2", TargetState.STOPPED.name());
+
+        expectConvertWriteRead(
+                CONNECTOR_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
+                "properties", SAMPLE_CONFIGS.get(0));
+
+        // Writing should block until it is written and read back from Kafka
+        configStorage.putConnectorConfig(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), TargetState.STOPPED);
+        configState = configStorage.snapshot();
+        assertEquals(2, configState.offset());
+        assertEquals(TargetState.STOPPED, configState.targetState(CONNECTOR_IDS.get(0)));
+        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
+
+        // We don't expect the config update listener's onConnectorTargetStateChange hook to be invoked
+        verify(configUpdateListener, never()).onConnectorTargetStateChange(anyString());
+
+        verify(configUpdateListener).onConnectorConfigUpdate(CONNECTOR_IDS.get(0));
+
+        configStorage.stop();
+        verify(configLog).stop();
+    }
+
+    @Test
+    public void testRecordToRestartRequest() {
+        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
+                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
+        Struct struct = RESTART_REQUEST_STRUCTS.get(0);
+        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
+        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
+        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
+        assertEquals(struct.getBoolean(INCLUDE_TASKS_FIELD_NAME), restartRequest.includeTasks());
+        assertEquals(struct.getBoolean(ONLY_FAILED_FIELD_NAME), restartRequest.onlyFailed());
+    }
+
+    @Test
+    public void testRecordToRestartRequestOnlyFailedInconsistent() {
+        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
+                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
+        Struct struct = ONLY_FAILED_MISSING_STRUCT;
+        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
+        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
+        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
+        assertEquals(struct.getBoolean(INCLUDE_TASKS_FIELD_NAME), restartRequest.includeTasks());
+        assertFalse(restartRequest.onlyFailed());
+    }
+
+    @Test
+    public void testRecordToRestartRequestIncludeTasksInconsistent() {
+        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
+                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
+        Struct struct = INCLUDE_TASKS_MISSING_STRUCT;
+        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
+        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
+        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
+        assertFalse(restartRequest.includeTasks());
+        assertEquals(struct.getBoolean(ONLY_FAILED_FIELD_NAME), restartRequest.onlyFailed());
+    }
+
+    @Test
+    public void testFencableProducerPropertiesOverrideUserSuppliedValues() {
+        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "preparing");
+        String groupId = "my-other-connect-cluster";
+        props.put(GROUP_ID_CONFIG, groupId);
+        props.put(TRANSACTIONAL_ID_CONFIG, "my-custom-transactional-id");
+        props.put(ENABLE_IDEMPOTENCE_CONFIG, "false");
+        createStore();
+
+        Map<String, Object> fencableProducerProperties = configStorage.fencableProducerProps(config);
+        assertEquals("connect-cluster-" + groupId, fencableProducerProperties.get(TRANSACTIONAL_ID_CONFIG));
+        assertEquals("true", fencableProducerProperties.get(ENABLE_IDEMPOTENCE_CONFIG));
+    }
+
+    @Test
+    public void testConsumerPropertiesDoNotOverrideUserSuppliedValuesWithoutExactlyOnceSourceEnabled() {
+        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "preparing");
+        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
+        createStore();
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+        verifyConfigure();
+
+        assertEquals(
+                IsolationLevel.READ_UNCOMMITTED.toString(),
+                capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
+        );
+    }
+
+    @Test
+    public void testClientIds() {
+        props = new HashMap<>(DEFAULT_CONFIG_STORAGE_PROPS);
+        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "enabled");
+        createStore();
+
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
+        verifyConfigure();
+
+        Map<String, Object> fencableProducerProps = configStorage.fencableProducerProps(config);
+
+        final String expectedClientId = CLIENT_ID_BASE + "configs";
+        assertEquals(expectedClientId, capturedProducerProps.getValue().get(CLIENT_ID_CONFIG));
+        assertEquals(expectedClientId, capturedConsumerProps.getValue().get(CLIENT_ID_CONFIG));
+        assertEquals(expectedClientId + "-leader", fencableProducerProps.get(CLIENT_ID_CONFIG));
+    }
+
+    private void verifyConfigure() {
+        verify(configStorage).createKafkaBasedLog(capturedTopic.capture(), capturedProducerProps.capture(),
+                capturedConsumerProps.capture(), capturedConsumedCallback.capture(),
+                capturedNewTopic.capture(), capturedAdminSupplier.capture(),
+                any(WorkerConfig.class), any(Time.class));
+    }
+
+    // If non-empty, deserializations should be a LinkedHashMap
+    private void expectStart(final List<ConsumerRecord<String, byte[]>> preexistingRecords,
+                             final Map<byte[], Struct> deserializations) {
+        doAnswer(invocation -> {
+            for (ConsumerRecord<String, byte[]> rec : preexistingRecords)
+                capturedConsumedCallback.getValue().onCompletion(null, rec);
+            return null;
+        }).when(configLog).start();
+
+        for (Map.Entry<byte[], Struct> deserializationEntry : deserializations.entrySet()) {
+            // Note null schema because default settings for internal serialization are schema-less
+            when(converter.toConnectData(TOPIC, deserializationEntry.getKey()))
+                    .thenReturn(new SchemaAndValue(null, structToMap(deserializationEntry.getValue())));
+        }
+    }
+
+    private void expectPartitionCount(int partitionCount) {
+        when(configLog.partitionCount()).thenReturn(partitionCount);
+    }
+
+    // Expect a conversion & write to the underlying log, followed by a subsequent read when the data is consumed back
+    // from the log. Validate the data that is captured when the conversion is performed matches the specified data
+    // (by checking a single field's value)
+    private void expectConvertWriteRead(final String configKey, final Schema valueSchema, final byte[] serialized,
+                                        final String dataFieldName, final Object dataFieldValue) throws Exception {
+        final ArgumentCaptor<Struct> capturedRecord = ArgumentCaptor.forClass(Struct.class);
+        if (serialized != null)
+            when(converter.fromConnectData(eq(TOPIC), eq(valueSchema), capturedRecord.capture()))
+                    .thenReturn(serialized);
+
+        when(configLog.sendWithReceipt(configKey, serialized)).thenReturn(producerFuture);
+        when(producerFuture.get(anyLong(), any(TimeUnit.class))).thenReturn(null);
+        when(converter.toConnectData(TOPIC, serialized)).thenAnswer((Answer<SchemaAndValue>) invocation -> {
+            if (dataFieldName != null)
+                assertEquals(dataFieldValue, capturedRecord.getValue().get(dataFieldName));
+            // Note null schema because default settings for internal serialization are schema-less
+            return new SchemaAndValue(null, serialized == null ? null : structToMap(capturedRecord.getValue()));
+        });
+    }
+
+    // This map needs to maintain ordering
+    private Answer<Future<Void>> expectReadToEnd(final Map<String, byte[]> serializedConfigs) {
+        return invocation -> {
+            TestFuture<Void> future = new TestFuture<>();
+            for (Map.Entry<String, byte[]> entry : serializedConfigs.entrySet()) {
+                capturedConsumedCallback.getValue().onCompletion(null,
+                        new ConsumerRecord<>(TOPIC, 0, logOffset++, 0L, TimestampType.CREATE_TIME, 0, 0,
+                                entry.getKey(), entry.getValue(), new RecordHeaders(), Optional.empty()));
+            }
+            future.resolveOnGet((Void) null);
+            return future;
+        };
+    }
+
+    // Generates a Map representation of Struct. Only does shallow traversal, so nested structs are not converted
+    private Map<String, Object> structToMap(Struct struct) {
+        if (struct == null)
+            return null;
+
+        HashMap<String, Object> result = new HashMap<>();
+        for (Field field : struct.schema().fields())
+            result.put(field.name(), struct.get(field));
+        return result;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
index 39068a990459f..bca2a73f1c7bc 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaConfigBackingStoreTest.java
@@ -18,10 +18,8 @@
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.admin.NewTopic;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.Producer;
-import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.IsolationLevel;
 import org.apache.kafka.common.config.ConfigException;
@@ -72,7 +70,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
-import static org.apache.kafka.clients.CommonClientConfigs.CLIENT_ID_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.ISOLATION_LEVEL_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.TRANSACTIONAL_ID_CONFIG;
@@ -84,7 +81,6 @@
 import static org.apache.kafka.connect.storage.KafkaConfigBackingStore.RESTART_KEY;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -230,174 +226,6 @@ public void setUp() {
         createStore();
     }
 
-    @Test
-    public void testStartStop() throws Exception {
-        props.put("config.storage.min.insync.replicas", "3");
-        props.put("config.storage.max.message.bytes", "1001");
-        createStore();
-        expectConfigure();
-        expectStart(Collections.emptyList(), Collections.emptyMap());
-        expectPartitionCount(1);
-        expectStop();
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-
-        assertEquals(TOPIC, capturedTopic.getValue());
-        assertEquals("org.apache.kafka.common.serialization.StringSerializer", capturedProducerProps.getValue().get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG));
-        assertEquals("org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG));
-        assertEquals("org.apache.kafka.common.serialization.StringDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG));
-        assertEquals("org.apache.kafka.common.serialization.ByteArrayDeserializer", capturedConsumerProps.getValue().get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG));
-
-        assertEquals(TOPIC, capturedNewTopic.getValue().name());
-        assertEquals(1, capturedNewTopic.getValue().numPartitions());
-        assertEquals(TOPIC_REPLICATION_FACTOR, capturedNewTopic.getValue().replicationFactor());
-        assertEquals("3", capturedNewTopic.getValue().configs().get("min.insync.replicas"));
-        assertEquals("1001", capturedNewTopic.getValue().configs().get("max.message.bytes"));
-        configStorage.start();
-        configStorage.stop();
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testSnapshotCannotMutateInternalState() throws Exception {
-        props.put("config.storage.min.insync.replicas", "3");
-        props.put("config.storage.max.message.bytes", "1001");
-        createStore();
-        expectConfigure();
-        expectStart(Collections.emptyList(), Collections.emptyMap());
-        expectPartitionCount(1);
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-
-        configStorage.start();
-        ClusterConfigState snapshot = configStorage.snapshot();
-        assertNotSame(snapshot.connectorTaskCounts, configStorage.connectorTaskCounts);
-        assertNotSame(snapshot.connectorConfigs, configStorage.connectorConfigs);
-        assertNotSame(snapshot.connectorTargetStates, configStorage.connectorTargetStates);
-        assertNotSame(snapshot.taskConfigs, configStorage.taskConfigs);
-        assertNotSame(snapshot.connectorTaskCountRecords, configStorage.connectorTaskCountRecords);
-        assertNotSame(snapshot.connectorTaskConfigGenerations, configStorage.connectorTaskConfigGenerations);
-        assertNotSame(snapshot.connectorsPendingFencing, configStorage.connectorsPendingFencing);
-        assertNotSame(snapshot.inconsistentConnectors, configStorage.inconsistent);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPutConnectorConfig() throws Exception {
-        expectConfigure();
-        expectStart(Collections.emptyList(), Collections.emptyMap());
-
-        expectConvertWriteAndRead(
-                CONNECTOR_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
-                "properties", SAMPLE_CONFIGS.get(0));
-        configUpdateListener.onConnectorConfigUpdate(CONNECTOR_IDS.get(0));
-        EasyMock.expectLastCall();
-
-        expectConvertWriteAndRead(
-                CONNECTOR_CONFIG_KEYS.get(1), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(1),
-                "properties", SAMPLE_CONFIGS.get(1));
-        configUpdateListener.onConnectorConfigUpdate(CONNECTOR_IDS.get(1));
-        EasyMock.expectLastCall();
-
-        // Config deletion
-        expectConnectorRemoval(CONNECTOR_CONFIG_KEYS.get(1), TARGET_STATE_KEYS.get(1));
-        configUpdateListener.onConnectorConfigRemove(CONNECTOR_IDS.get(1));
-        EasyMock.expectLastCall();
-
-        expectPartitionCount(1);
-        expectStop();
-
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-        configStorage.start();
-
-        // Null before writing
-        ClusterConfigState configState = configStorage.snapshot();
-        assertEquals(-1, configState.offset());
-        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(0)));
-        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
-
-        // Writing should block until it is written and read back from Kafka
-        configStorage.putConnectorConfig(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), null);
-        configState = configStorage.snapshot();
-        assertEquals(1, configState.offset());
-        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
-        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
-
-        // Second should also block and all configs should still be available
-        configStorage.putConnectorConfig(CONNECTOR_IDS.get(1), SAMPLE_CONFIGS.get(1), null);
-        configState = configStorage.snapshot();
-        assertEquals(2, configState.offset());
-        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
-        assertEquals(SAMPLE_CONFIGS.get(1), configState.connectorConfig(CONNECTOR_IDS.get(1)));
-
-        // Deletion should remove the second one we added
-        configStorage.removeConnectorConfig(CONNECTOR_IDS.get(1));
-        configState = configStorage.snapshot();
-        assertEquals(4, configState.offset());
-        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
-        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(1)));
-        assertNull(configState.targetState(CONNECTOR_IDS.get(1)));
-
-        configStorage.stop();
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testPutConnectorConfigWithTargetState() throws Exception {
-        expectConfigure();
-        expectStart(Collections.emptyList(), Collections.emptyMap());
-
-        // We expect to write the target state first, followed by the config write and then a read to end
-
-        expectConvertWriteRead(
-            TARGET_STATE_KEYS.get(0), KafkaConfigBackingStore.TARGET_STATE_V1, TARGET_STATES_SERIALIZED.get(2),
-            "state.v2", TargetState.STOPPED.name());
-        // We don't expect the config update listener's onConnectorTargetStateChange hook to be invoked
-
-        expectConvertWriteRead(
-            CONNECTOR_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
-            "properties", SAMPLE_CONFIGS.get(0));
-        configUpdateListener.onConnectorConfigUpdate(CONNECTOR_IDS.get(0));
-        EasyMock.expectLastCall();
-
-        LinkedHashMap<String, byte[]> recordsToRead = new LinkedHashMap<>();
-        recordsToRead.put(TARGET_STATE_KEYS.get(0), TARGET_STATES_SERIALIZED.get(2));
-        recordsToRead.put(CONNECTOR_CONFIG_KEYS.get(0), CONFIGS_SERIALIZED.get(0));
-        expectReadToEnd(recordsToRead);
-
-        expectPartitionCount(1);
-        expectStop();
-
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-        configStorage.start();
-
-        // Null before writing
-        ClusterConfigState configState = configStorage.snapshot();
-        assertEquals(-1, configState.offset());
-        assertNull(configState.connectorConfig(CONNECTOR_IDS.get(0)));
-        assertNull(configState.targetState(CONNECTOR_IDS.get(0)));
-
-        // Writing should block until it is written and read back from Kafka
-        configStorage.putConnectorConfig(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), TargetState.STOPPED);
-        configState = configStorage.snapshot();
-        assertEquals(2, configState.offset());
-        assertEquals(TargetState.STOPPED, configState.targetState(CONNECTOR_IDS.get(0)));
-        assertEquals(SAMPLE_CONFIGS.get(0), configState.connectorConfig(CONNECTOR_IDS.get(0)));
-
-        configStorage.stop();
-
-        PowerMock.verifyAll();
-    }
-
     @Test
     public void testPutConnectorConfigProducerError() throws Exception {
         expectConfigure();
@@ -1438,42 +1266,6 @@ private void testPutRestartRequest(RestartRequest restartRequest) throws Excepti
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testRecordToRestartRequest() {
-        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
-                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
-        Struct struct = RESTART_REQUEST_STRUCTS.get(0);
-        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
-        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
-        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
-        assertEquals(struct.getBoolean(INCLUDE_TASKS_FIELD_NAME), restartRequest.includeTasks());
-        assertEquals(struct.getBoolean(ONLY_FAILED_FIELD_NAME), restartRequest.onlyFailed());
-    }
-
-    @Test
-    public void testRecordToRestartRequestOnlyFailedInconsistent() {
-        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
-                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
-        Struct struct = ONLY_FAILED_MISSING_STRUCT;
-        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
-        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
-        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
-        assertEquals(struct.getBoolean(INCLUDE_TASKS_FIELD_NAME), restartRequest.includeTasks());
-        assertFalse(restartRequest.onlyFailed());
-    }
-
-    @Test
-    public void testRecordToRestartRequestIncludeTasksInconsistent() {
-        ConsumerRecord<String, byte[]> record = new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0, 0, RESTART_CONNECTOR_KEYS.get(0),
-                CONFIGS_SERIALIZED.get(0), new RecordHeaders(), Optional.empty());
-        Struct struct = INCLUDE_TASKS_MISSING_STRUCT;
-        SchemaAndValue schemaAndValue = new SchemaAndValue(struct.schema(), structToMap(struct));
-        RestartRequest restartRequest = configStorage.recordToRestartRequest(record, schemaAndValue);
-        assertEquals(CONNECTOR_1_NAME, restartRequest.connectorName());
-        assertFalse(restartRequest.includeTasks());
-        assertEquals(struct.getBoolean(ONLY_FAILED_FIELD_NAME), restartRequest.onlyFailed());
-    }
-
     @Test
     public void testRestoreRestartRequestInconsistentState() throws Exception {
         // Restoring data should notify only of the latest values after loading is complete. This also validates
@@ -1612,24 +1404,6 @@ public void testFencableProducerPropertiesInsertedByDefault() throws Exception {
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testFencableProducerPropertiesOverrideUserSuppliedValues() throws Exception {
-        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "preparing");
-        String groupId = "my-other-connect-cluster";
-        props.put(GROUP_ID_CONFIG, groupId);
-        props.put(TRANSACTIONAL_ID_CONFIG, "my-custom-transactional-id");
-        props.put(ENABLE_IDEMPOTENCE_CONFIG, "false");
-        createStore();
-
-        PowerMock.replayAll();
-
-        Map<String, Object> fencableProducerProperties = configStorage.fencableProducerProps(config);
-        assertEquals("connect-cluster-" + groupId, fencableProducerProperties.get(TRANSACTIONAL_ID_CONFIG));
-        assertEquals("true", fencableProducerProperties.get(ENABLE_IDEMPOTENCE_CONFIG));
-
-        PowerMock.verifyAll();
-    }
-
     @Test
     public void testConsumerPropertiesInsertedByDefaultWithExactlyOnceSourceEnabled() throws Exception {
         props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "enabled");
@@ -1684,45 +1458,6 @@ public void testConsumerPropertiesNotInsertedByDefaultWithoutExactlyOnceSourceEn
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testConsumerPropertiesDoNotOverrideUserSuppliedValuesWithoutExactlyOnceSourceEnabled() throws Exception {
-        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "preparing");
-        props.put(ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
-        createStore();
-
-        expectConfigure();
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-
-        assertEquals(
-                IsolationLevel.READ_UNCOMMITTED.toString(),
-                capturedConsumerProps.getValue().get(ISOLATION_LEVEL_CONFIG)
-        );
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testClientIds() throws Exception {
-        props = new HashMap<>(DEFAULT_CONFIG_STORAGE_PROPS);
-        props.put(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG, "enabled");
-        createStore();
-
-        expectConfigure();
-        PowerMock.replayAll();
-
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, config);
-        Map<String, Object> fencableProducerProps = configStorage.fencableProducerProps(config);
-
-        final String expectedClientId = CLIENT_ID_BASE + "configs";
-        assertEquals(expectedClientId, capturedProducerProps.getValue().get(CLIENT_ID_CONFIG));
-        assertEquals(expectedClientId, capturedConsumerProps.getValue().get(CLIENT_ID_CONFIG));
-        assertEquals(expectedClientId + "-leader", fencableProducerProps.get(CLIENT_ID_CONFIG));
-
-        PowerMock.verifyAll();
-    }
-
     private void expectConfigure() throws Exception {
         PowerMock.expectPrivate(configStorage, "createKafkaBasedLog",
                 EasyMock.capture(capturedTopic), EasyMock.capture(capturedProducerProps),

From 6e4a098055d5ef852022a76a6ea48ccb880f6342 Mon Sep 17 00:00:00 2001
From: Calvin Liu <83986057+CalvinConfluent@users.noreply.github.com>
Date: Fri, 29 Mar 2024 16:54:55 -0700
Subject: [PATCH 231/258] KAFKA-16217: Stop the abort transaction try loop when
 closing producers (#15541)

This is a mitigation fix for the https://issues.apache.org/jira/browse/KAFKA-16217. Exceptions should not block closing the producers.
This PR reverts a part of the change #13591

Reviewers: Kirk True <ktrue@confluent.io>, Justine Olshan <jolshan@confluent.io>
---
 .../clients/producer/internals/Sender.java    |  5 +++--
 .../producer/internals/SenderTest.java        | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
index 99bc1d68b0b22..c4e2b73e8b91b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/Sender.java
@@ -270,13 +270,14 @@ public void run() {
         while (!forceClose && transactionManager != null && transactionManager.hasOngoingTransaction()) {
             if (!transactionManager.isCompleting()) {
                 log.info("Aborting incomplete transaction due to shutdown");
-
                 try {
                     // It is possible for the transaction manager to throw errors when aborting. Catch these
                     // so as not to interfere with the rest of the shutdown logic.
                     transactionManager.beginAbort();
                 } catch (Exception e) {
-                    log.error("Error in kafka producer I/O thread while aborting transaction: ", e);
+                    log.error("Error in kafka producer I/O thread while aborting transaction when during closing: ", e);
+                    // Force close in case the transactionManager is in error states.
+                    forceClose = true;
                 }
             }
             try {
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
index 4ef9dab4d096c..eb01d1d5841d7 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/SenderTest.java
@@ -132,9 +132,11 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class SenderTest {
     private static final int MAX_REQUEST_SIZE = 1024 * 1024;
@@ -3274,6 +3276,26 @@ public void testProducerBatchRetriesWhenPartitionLeaderChanges() throws Exceptio
         }
     }
 
+    // This test is expected to run fast. If timeout, the sender is not able to close properly.
+    @Timeout(5)
+    @Test
+    public void testSenderShouldCloseWhenTransactionManagerInErrorState() throws Exception {
+        metrics.close();
+        Map<String, String> clientTags = Collections.singletonMap("client-id", "clientA");
+        metrics = new Metrics(new MetricConfig().tags(clientTags));
+        TransactionManager transactionManager = mock(TransactionManager.class);
+        SenderMetricsRegistry metricsRegistry = new SenderMetricsRegistry(metrics);
+        Sender sender = new Sender(logContext, client, metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
+                1, metricsRegistry, time, REQUEST_TIMEOUT, RETRY_BACKOFF_MS, transactionManager, apiVersions);
+        when(transactionManager.hasOngoingTransaction()).thenReturn(true);
+        when(transactionManager.beginAbort()).thenThrow(new IllegalStateException());
+        sender.initiateClose();
+
+        // The sender should directly get closed.
+        sender.run();
+        verify(transactionManager, times(1)).close();
+    }
+
     /**
      * Test the scenario that FetchResponse returns NOT_LEADER_OR_FOLLOWER, indicating change in leadership, but it
      * does not contain new leader info(defined in KIP-951).

From d8673b26bf20634a2caf784159701062212cb103 Mon Sep 17 00:00:00 2001
From: Nikolay <nizhikov@apache.org>
Date: Sat, 30 Mar 2024 06:54:22 +0300
Subject: [PATCH 232/258] KAFKA-15899 [1/2] Move kafka.security package from
 core to server module (#15572)

1) This PR moves kafka.security classes from core to server module.
2) AclAuthorizer not moved, because it has heavy dependencies on core classes that not rewrited from scala at the moment.
3) AclAuthorizer will be deleted as part of ZK removal

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 build.gradle                                  |   1 +
 checkstyle/import-control-jmh-benchmarks.xml  |   1 +
 .../org/apache/kafka/common/utils/Utils.java  |   9 +
 .../main/scala/kafka/admin/AclCommand.scala   |  10 +-
 .../scala/kafka/network/SocketServer.scala    |   2 +-
 .../kafka/security/CredentialProvider.scala   |  62 ---
 .../security/authorizer/AclAuthorizer.scala   |  18 +-
 .../kafka/security/authorizer/AclEntry.scala  | 147 -------
 .../main/scala/kafka/server/AuthHelper.scala  |   4 +-
 .../scala/kafka/server/BrokerServer.scala     |   2 +-
 .../scala/kafka/server/ConfigHandler.scala    |   2 +-
 .../scala/kafka/server/ControllerServer.scala |   3 +-
 .../main/scala/kafka/server/KafkaBroker.scala |   2 +-
 .../main/scala/kafka/server/KafkaServer.scala |   2 +-
 .../server/metadata/ScramPublisher.scala      |   2 +-
 .../scala/kafka/tools/TestRaftServer.scala    |   2 +-
 .../main/scala/kafka/zk/KafkaZkClient.scala   |   4 +-
 core/src/main/scala/kafka/zk/ZkData.scala     |  10 +-
 .../zk/migration/ZkAclMigrationClient.scala   |   3 +-
 .../AbstractAuthorizerIntegrationTest.scala   |   4 +-
 .../kafka/api/AuthorizerIntegrationTest.scala | 364 +++++++++---------
 .../kafka/api/BaseAdminIntegrationTest.scala  |   8 +-
 .../DescribeAuthorizedOperationsTest.scala    |  13 +-
 .../kafka/api/EndToEndAuthorizationTest.scala |  12 +-
 .../api/GroupAuthorizerIntegrationTest.scala  |   4 +-
 .../api/PlaintextAdminIntegrationTest.scala   |   4 +-
 .../api/SaslSslAdminIntegrationTest.scala     |   6 +-
 .../kafka/zk/ZkMigrationIntegrationTest.scala |  10 +-
 .../scala/kafka/zk/LiteralAclStoreTest.scala  |   4 +-
 .../unit/kafka/admin/AclCommandTest.scala     |   9 +-
 .../unit/kafka/network/SocketServerTest.scala |   2 +-
 .../AclAuthorizerWithZkSaslTest.scala         |   6 +-
 .../security/authorizer/AclEntryTest.scala    |  15 +-
 .../security/authorizer/AuthorizerTest.scala  |  60 +--
 .../authorizer/BaseAuthorizerTest.scala       |  10 +-
 .../DelegationTokenManagerTest.scala          |   6 +-
 .../server/DescribeClusterRequestTest.scala   |   8 +-
 .../scala/unit/kafka/utils/TestUtils.scala    |   2 +-
 .../unit/kafka/zk/KafkaZkClientTest.scala     |  15 +-
 .../migration/ZkAclMigrationClientTest.scala  |  19 +-
 .../kafka/jmh/acl/AuthorizerBenchmark.java    |   2 +-
 .../kafka/security/CredentialProvider.java    |  69 ++++
 .../kafka/security/authorizer/AclEntry.java   | 209 ++++++++++
 .../group/AuthorizerIntegrationTest.java      |   3 +-
 44 files changed, 617 insertions(+), 533 deletions(-)
 delete mode 100644 core/src/main/scala/kafka/security/CredentialProvider.scala
 delete mode 100644 core/src/main/scala/kafka/security/authorizer/AclEntry.scala
 create mode 100644 server/src/main/java/org/apache/kafka/security/CredentialProvider.java
 create mode 100644 server/src/main/java/org/apache/kafka/security/authorizer/AclEntry.java

diff --git a/build.gradle b/build.gradle
index 0564d27a01380..14ec93512b692 100644
--- a/build.gradle
+++ b/build.gradle
@@ -829,6 +829,7 @@ project(':server') {
     implementation project(':transaction-coordinator')
     implementation project(':raft')
     implementation libs.metrics
+    implementation libs.jacksonDatabind
 
     implementation libs.slf4jApi
 
diff --git a/checkstyle/import-control-jmh-benchmarks.xml b/checkstyle/import-control-jmh-benchmarks.xml
index 1160e3f67dfde..d3484802a78ef 100644
--- a/checkstyle/import-control-jmh-benchmarks.xml
+++ b/checkstyle/import-control-jmh-benchmarks.xml
@@ -47,6 +47,7 @@
     <allow pkg="kafka.message"/>
     <allow pkg="org.mockito"/>
     <allow pkg="kafka.security.authorizer"/>
+    <allow pkg="org.apache.kafka.security.authorizer"/>
     <allow pkg="org.apache.kafka.server"/>
     <allow pkg="org.apache.kafka.storage"/>
     <allow pkg="org.apache.kafka.clients"/>
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
index 3fd3833de1e44..0eac34cd1cd16 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Utils.java
@@ -1685,6 +1685,15 @@ public static <V> Map<String, V> entriesWithPrefix(Map<String, V> map, String pr
         return result;
     }
 
+    /**
+     * Checks requirement. Throw {@link IllegalArgumentException} if {@code requirement} failed.
+     * @param requirement Requirement to check.
+     */
+    public static void require(boolean requirement) {
+        if (!requirement)
+            throw new IllegalArgumentException("requirement failed");
+    }
+
     /**
      * A runnable that can throw checked exception.
      */
diff --git a/core/src/main/scala/kafka/admin/AclCommand.scala b/core/src/main/scala/kafka/admin/AclCommand.scala
index 0415598b4ed5b..f1fc27cb4bf3e 100644
--- a/core/src/main/scala/kafka/admin/AclCommand.scala
+++ b/core/src/main/scala/kafka/admin/AclCommand.scala
@@ -20,7 +20,7 @@ package kafka.admin
 import java.util.Properties
 import joptsimple._
 import joptsimple.util.EnumConverter
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.utils._
 import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
 import org.apache.kafka.common.acl._
@@ -30,7 +30,7 @@ import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceP
 import org.apache.kafka.common.security.JaasUtils
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{Utils, SecurityUtils => JSecurityUtils}
-import org.apache.kafka.security.authorizer.AuthorizerUtils
+import org.apache.kafka.security.authorizer.{AclEntry, AuthorizerUtils}
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
@@ -436,7 +436,7 @@ object AclCommand extends Logging {
     if (opts.options.has(hostOptionSpec))
       opts.options.valuesOf(hostOptionSpec).asScala.map(_.trim).toSet
     else if (opts.options.has(principalOptionSpec))
-      Set[String](AclEntry.WildcardHost)
+      Set[String](AclEntry.WILDCARD_HOST)
     else
       Set.empty[String]
   }
@@ -486,7 +486,7 @@ object AclCommand extends Logging {
 
   private def validateOperation(opts: AclCommandOptions, resourceToAcls: Map[ResourcePatternFilter, Set[AccessControlEntry]]): Unit = {
     for ((resource, acls) <- resourceToAcls) {
-      val validOps = AclEntry.supportedOperations(resource.resourceType) + AclOperation.ALL
+      val validOps = AclEntry.supportedOperations(resource.resourceType).asScala.toSet + AclOperation.ALL
       if ((acls.map(_.operation) -- validOps).nonEmpty)
         CommandLineUtils.printUsageAndExit(opts.parser, s"ResourceType ${resource.resourceType} only supports operations ${validOps.map(JSecurityUtils.operationName).mkString(", ")}")
     }
@@ -566,7 +566,7 @@ object AclCommand extends Logging {
     val listOpt: OptionSpecBuilder = parser.accepts("list", "List ACLs for the specified resource, use --topic <topic> or --group <group> or --cluster to specify a resource.")
 
     val operationsOpt: OptionSpec[String] = parser.accepts("operation", "Operation that is being allowed or denied. Valid operation names are: " + Newline +
-      AclEntry.AclOperations.map("\t" + JSecurityUtils.operationName(_)).mkString(Newline) + Newline)
+      AclEntry.ACL_OPERATIONS.asScala.map("\t" + JSecurityUtils.operationName(_)).mkString(Newline) + Newline)
       .withRequiredArg
       .ofType(classOf[String])
       .defaultsTo(JSecurityUtils.operationName(AclOperation.ALL))
diff --git a/core/src/main/scala/kafka/network/SocketServer.scala b/core/src/main/scala/kafka/network/SocketServer.scala
index e90af2e18ca02..ea316983d0c9b 100644
--- a/core/src/main/scala/kafka/network/SocketServer.scala
+++ b/core/src/main/scala/kafka/network/SocketServer.scala
@@ -30,7 +30,6 @@ import kafka.network.ConnectionQuotas._
 import kafka.network.Processor._
 import kafka.network.RequestChannel.{CloseConnectionResponse, EndThrottlingResponse, NoOpResponse, SendResponse, StartThrottlingResponse}
 import kafka.network.SocketServer._
-import kafka.security.CredentialProvider
 import kafka.server.{ApiVersionManager, BrokerReconfigurable, KafkaConfig}
 import org.apache.kafka.common.message.ApiMessageType.ListenerType
 import kafka.utils._
@@ -47,6 +46,7 @@ import org.apache.kafka.common.requests.{ApiVersionsRequest, RequestContext, Req
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.utils.{KafkaThread, LogContext, Time, Utils}
 import org.apache.kafka.common.{Endpoint, KafkaException, MetricName, Reconfigurable}
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.metrics.KafkaMetricsGroup
 import org.apache.kafka.server.util.FutureUtils
 import org.slf4j.event.Level
diff --git a/core/src/main/scala/kafka/security/CredentialProvider.scala b/core/src/main/scala/kafka/security/CredentialProvider.scala
deleted file mode 100644
index 6c58ff0a5dd3f..0000000000000
--- a/core/src/main/scala/kafka/security/CredentialProvider.scala
+++ /dev/null
@@ -1,62 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.security
-
-import java.util.{Collection, Properties}
-import org.apache.kafka.clients.admin.{ScramMechanism => AdminScramMechanism}
-import org.apache.kafka.common.security.authenticator.CredentialCache
-import org.apache.kafka.common.security.scram.ScramCredential
-import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramMechanism}
-import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
-
-class CredentialProvider(scramMechanisms: Collection[String], val tokenCache: DelegationTokenCache) {
-
-  val credentialCache = new CredentialCache
-  ScramCredentialUtils.createCache(credentialCache, scramMechanisms)
-
-  def updateCredentials(username: String, config: Properties): Unit = {
-    for (mechanism <- ScramMechanism.values()) {
-      val cache = credentialCache.cache(mechanism.mechanismName, classOf[ScramCredential])
-      if (cache != null) {
-        config.getProperty(mechanism.mechanismName) match {
-          case null => cache.remove(username)
-          case c => cache.put(username, ScramCredentialUtils.credentialFromString(c))
-        }
-      }
-    }
-  }
-
-  def updateCredential(
-    mechanism: AdminScramMechanism,
-    name: String,
-    credential: ScramCredential
-  ): Unit = {
-    val cache = credentialCache.cache(mechanism.mechanismName(), classOf[ScramCredential])
-    cache.put(name, credential)
-  }
-
-  def removeCredentials(
-    mechanism: AdminScramMechanism,
-    name: String
-  ): Unit = {
-    val cache = credentialCache.cache(mechanism.mechanismName(), classOf[ScramCredential])
-    if (cache != null) {
-      cache.remove(name)
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
index fc2d30f8e8cb1..f091a43802b02 100644
--- a/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -18,9 +18,7 @@ package kafka.security.authorizer
 
 import java.{lang, util}
 import java.util.concurrent.{CompletableFuture, CompletionStage}
-
 import com.typesafe.scalalogging.Logger
-import kafka.security.authorizer.AclEntry.ResourceSeparator
 import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.utils._
 import kafka.utils.Implicits._
@@ -29,11 +27,13 @@ import org.apache.kafka.common.Endpoint
 import org.apache.kafka.common.acl._
 import org.apache.kafka.common.acl.AclOperation._
 import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+import org.apache.kafka.security.authorizer.AclEntry.RESOURCE_SEPARATOR
 import org.apache.kafka.common.errors.{ApiException, InvalidRequestException, UnsupportedVersionException}
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.resource._
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{SecurityUtils, Time}
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
 import org.apache.kafka.server.authorizer._
 import org.apache.kafka.server.common.MetadataVersion.IBP_2_0_IV1
@@ -410,8 +410,8 @@ class AclAuthorizer extends Authorizer with Logging {
                                 principal: String, host: String, op: AclOperation, permission: AclPermissionType,
                                 resourceType: ResourceType, patternType: PatternType): ArrayBuffer[Set[String]] = {
     val matched = ArrayBuffer[immutable.Set[String]]()
-    for (p <- Set(principal, AclEntry.WildcardPrincipalString);
-         h <- Set(host, AclEntry.WildcardHost);
+    for (p <- Set(principal, AclEntry.WILDCARD_PRINCIPAL_STRING);
+         h <- Set(host, AclEntry.WILDCARD_HOST);
          o <- Set(op, AclOperation.ALL)) {
       val resourceTypeKey = ResourceTypeKey(
         new AccessControlEntry(p, h, o, permission), resourceType, patternType)
@@ -426,8 +426,8 @@ class AclAuthorizer extends Authorizer with Logging {
   private def hasMatchingResources(resourceSnapshot: immutable.Map[ResourceTypeKey, immutable.Set[String]],
                                    principal: String, host: String, op: AclOperation, permission: AclPermissionType,
                                    resourceType: ResourceType, patternType: PatternType): Boolean = {
-    for (p <- Set(principal, AclEntry.WildcardPrincipalString);
-         h <- Set(host, AclEntry.WildcardHost);
+    for (p <- Set(principal, AclEntry.WILDCARD_PRINCIPAL_STRING);
+         h <- Set(host, AclEntry.WILDCARD_HOST);
          o <- Set(op, AclOperation.ALL)) {
           val resourceTypeKey = ResourceTypeKey(
             new AccessControlEntry(p, h, o, permission), resourceType, patternType)
@@ -566,9 +566,9 @@ class AclAuthorizer extends Authorizer with Logging {
                                 acls: AclSeqs): Boolean = {
     acls.find { acl =>
       acl.permissionType == permissionType &&
-        (acl.kafkaPrincipal == principal || acl.kafkaPrincipal == AclEntry.WildcardPrincipal) &&
+        (acl.kafkaPrincipal == principal || acl.kafkaPrincipal == AclEntry.WILDCARD_PRINCIPAL) &&
         (operation == acl.operation || acl.operation == AclOperation.ALL) &&
-        (acl.host == host || acl.host == AclEntry.WildcardHost)
+        (acl.host == host || acl.host == AclEntry.WILDCARD_HOST)
     }.exists { acl =>
       authorizerLogger.debug(s"operation = $operation on resource = $resource from host = $host is $permissionType based on acl = $acl")
       true
@@ -603,7 +603,7 @@ class AclAuthorizer extends Authorizer with Logging {
       val operation = SecurityUtils.operationName(action.operation)
       val host = requestContext.clientAddress.getHostAddress
       val resourceType = SecurityUtils.resourceTypeName(action.resourcePattern.resourceType)
-      val resource = s"$resourceType$ResourceSeparator${action.resourcePattern.patternType}$ResourceSeparator${action.resourcePattern.name}"
+      val resource = s"$resourceType$RESOURCE_SEPARATOR${action.resourcePattern.patternType}$RESOURCE_SEPARATOR${action.resourcePattern.name}"
       val authResult = if (authorized) "Allowed" else "Denied"
       val apiKey = if (ApiKeys.hasId(requestContext.requestType)) ApiKeys.forId(requestContext.requestType).name else requestContext.requestType
       val refCount = action.resourceReferenceCount
diff --git a/core/src/main/scala/kafka/security/authorizer/AclEntry.scala b/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
deleted file mode 100644
index be54746c99117..0000000000000
--- a/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
+++ /dev/null
@@ -1,147 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.security.authorizer
-
-import kafka.utils.Json
-import org.apache.kafka.common.acl.{AccessControlEntry, AclOperation, AclPermissionType}
-import org.apache.kafka.common.acl.AclOperation.{READ, WRITE, CREATE, DESCRIBE, DELETE, ALTER, DESCRIBE_CONFIGS, ALTER_CONFIGS, CLUSTER_ACTION, IDEMPOTENT_WRITE, CREATE_TOKENS, DESCRIBE_TOKENS}
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.resource.{ResourcePattern, ResourceType}
-import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.common.utils.SecurityUtils
-
-import scala.jdk.CollectionConverters._
-
-object AclEntry {
-  val WildcardPrincipal: KafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*")
-  val WildcardPrincipalString: String = WildcardPrincipal.toString
-  val WildcardHost: String = "*"
-  val WildcardResource: String = ResourcePattern.WILDCARD_RESOURCE
-
-  val ResourceSeparator: String = ":"
-  val ResourceTypes: Set[ResourceType] = ResourceType.values.toSet
-    .filterNot(t => t == ResourceType.UNKNOWN || t == ResourceType.ANY)
-  val AclOperations: Set[AclOperation] = AclOperation.values.toSet
-    .filterNot(t => t == AclOperation.UNKNOWN || t == AclOperation.ANY)
-
-  private val PrincipalKey = "principal"
-  private val PermissionTypeKey = "permissionType"
-  private val OperationKey = "operation"
-  private val HostsKey = "host"
-  val VersionKey: String = "version"
-  val CurrentVersion: Int = 1
-  private val AclsKey = "acls"
-
-  def apply(principal: KafkaPrincipal,
-            permissionType: AclPermissionType,
-            host: String,
-            operation: AclOperation): AclEntry = {
-    new AclEntry(new AccessControlEntry(if (principal == null) null else principal.toString,
-      host, operation, permissionType))
-  }
-
-  /**
-   * Parse JSON representation of ACLs
-   * @param bytes of acls json string
-   *
-   * <p>
-      {
-        "version": 1,
-        "acls": [
-          {
-            "host":"host1",
-            "permissionType": "Deny",
-            "operation": "Read",
-            "principal": "User:alice"
-          }
-        ]
-      }
-   * </p>
-   *
-   * @return set of AclEntry objects from the JSON string
-   */
-  def fromBytes(bytes: Array[Byte]): Set[AclEntry] = {
-    if (bytes == null || bytes.isEmpty)
-      return collection.immutable.Set.empty[AclEntry]
-
-    Json.parseBytes(bytes).map(_.asJsonObject).map { js =>
-      //the acl json version.
-      require(js(VersionKey).to[Int] == CurrentVersion)
-      js(AclsKey).asJsonArray.iterator.map(_.asJsonObject).map { itemJs =>
-        val principal = SecurityUtils.parseKafkaPrincipal(itemJs(PrincipalKey).to[String])
-        val permissionType = SecurityUtils.permissionType(itemJs(PermissionTypeKey).to[String])
-        val host = itemJs(HostsKey).to[String]
-        val operation = SecurityUtils.operation(itemJs(OperationKey).to[String])
-        AclEntry(principal, permissionType, host, operation)
-      }.toSet
-    }.getOrElse(Set.empty)
-  }
-
-  def toJsonCompatibleMap(acls: Set[AclEntry]): Map[String, Any] = {
-    Map(AclEntry.VersionKey -> AclEntry.CurrentVersion, AclEntry.AclsKey -> acls.map(acl => acl.toMap.asJava).toList.asJava)
-  }
-
-  def supportedOperations(resourceType: ResourceType): Set[AclOperation] = {
-    resourceType match {
-      case ResourceType.TOPIC => Set(READ, WRITE, CREATE, DESCRIBE, DELETE, ALTER, DESCRIBE_CONFIGS, ALTER_CONFIGS)
-      case ResourceType.GROUP => Set(READ, DESCRIBE, DELETE)
-      case ResourceType.CLUSTER => Set(CREATE, CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE, ALTER, DESCRIBE)
-      case ResourceType.TRANSACTIONAL_ID => Set(DESCRIBE, WRITE)
-      case ResourceType.DELEGATION_TOKEN => Set(DESCRIBE)
-      case ResourceType.USER => Set(CREATE_TOKENS, DESCRIBE_TOKENS)
-      case _ => throw new IllegalArgumentException("Not a concrete resource type")
-    }
-  }
-
-  def authorizationError(resourceType: ResourceType): Errors = {
-    resourceType match {
-      case ResourceType.TOPIC => Errors.TOPIC_AUTHORIZATION_FAILED
-      case ResourceType.GROUP => Errors.GROUP_AUTHORIZATION_FAILED
-      case ResourceType.CLUSTER => Errors.CLUSTER_AUTHORIZATION_FAILED
-      case ResourceType.TRANSACTIONAL_ID => Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
-      case ResourceType.DELEGATION_TOKEN => Errors.DELEGATION_TOKEN_AUTHORIZATION_FAILED
-      case _ => throw new IllegalArgumentException("Authorization error type not known")
-    }
-  }
-}
-
-class AclEntry(val ace: AccessControlEntry)
-  extends AccessControlEntry(ace.principal, ace.host, ace.operation, ace.permissionType) {
-
-  val kafkaPrincipal: KafkaPrincipal = if (principal == null)
-    null
-  else
-    SecurityUtils.parseKafkaPrincipal(principal)
-
-  def toMap: Map[String, Any] = {
-    Map(AclEntry.PrincipalKey -> principal,
-      AclEntry.PermissionTypeKey -> SecurityUtils.permissionTypeName(permissionType),
-      AclEntry.OperationKey -> SecurityUtils.operationName(operation),
-      AclEntry.HostsKey -> host)
-  }
-
-  override def hashCode(): Int = ace.hashCode()
-
-  override def equals(o: scala.Any): Boolean = super.equals(o) // to keep spotbugs happy
-
-  override def toString: String = {
-    "%s has %s permission for operations: %s from hosts: %s".format(principal, permissionType.name, operation, host)
-  }
-
-}
-
diff --git a/core/src/main/scala/kafka/server/AuthHelper.scala b/core/src/main/scala/kafka/server/AuthHelper.scala
index 1d988bdea60a9..6e779ce9007a9 100644
--- a/core/src/main/scala/kafka/server/AuthHelper.scala
+++ b/core/src/main/scala/kafka/server/AuthHelper.scala
@@ -20,7 +20,6 @@ package kafka.server
 import java.lang.{Byte => JByte}
 import java.util.Collections
 import kafka.network.RequestChannel
-import kafka.security.authorizer.AclEntry
 import kafka.utils.CoreUtils
 import org.apache.kafka.clients.admin.EndpointType
 import org.apache.kafka.common.acl.AclOperation
@@ -34,6 +33,7 @@ import org.apache.kafka.common.resource.Resource.CLUSTER_NAME
 import org.apache.kafka.common.resource.ResourceType.CLUSTER
 import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.authorizer.{Action, AuthorizationResult, Authorizer}
 
 import scala.collection.Seq
@@ -60,7 +60,7 @@ class AuthHelper(authorizer: Option[Authorizer]) {
   }
 
   def authorizedOperations(request: RequestChannel.Request, resource: Resource): Int = {
-    val supportedOps = AclEntry.supportedOperations(resource.resourceType).toList
+    val supportedOps = AclEntry.supportedOperations(resource.resourceType).asScala.toList
     val authorizedOps = authorizer match {
       case Some(authZ) =>
         val resourcePattern = new ResourcePattern(resource.resourceType, resource.name, PatternType.LITERAL)
diff --git a/core/src/main/scala/kafka/server/BrokerServer.scala b/core/src/main/scala/kafka/server/BrokerServer.scala
index 8ac568a5e3a7e..be9b53c044704 100644
--- a/core/src/main/scala/kafka/server/BrokerServer.scala
+++ b/core/src/main/scala/kafka/server/BrokerServer.scala
@@ -24,7 +24,6 @@ import kafka.log.LogManager
 import kafka.log.remote.RemoteLogManager
 import kafka.network.{DataPlaneAcceptor, SocketServer}
 import kafka.raft.KafkaRaftManager
-import kafka.security.CredentialProvider
 import kafka.server.metadata.{AclPublisher, BrokerMetadataPublisher, ClientQuotaMetadataManager, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicConfigPublisher, KRaftMetadataCache, ScramPublisher}
 import kafka.utils.CoreUtils
 import org.apache.kafka.common.config.ConfigException
@@ -42,6 +41,7 @@ import org.apache.kafka.coordinator.group.{GroupCoordinator, GroupCoordinatorCon
 import org.apache.kafka.image.publisher.MetadataPublisher
 import org.apache.kafka.metadata.{BrokerState, ListenerInfo, VersionRange}
 import org.apache.kafka.raft.RaftConfig
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.{AssignmentsManager, ClientMetricsManager, NodeToControllerChannelManager}
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.common.{ApiMessageAndVersion, DirectoryEventHandler, TopicIdPartition}
diff --git a/core/src/main/scala/kafka/server/ConfigHandler.scala b/core/src/main/scala/kafka/server/ConfigHandler.scala
index 05db9b3bc15a9..e9f6df007b2c9 100644
--- a/core/src/main/scala/kafka/server/ConfigHandler.scala
+++ b/core/src/main/scala/kafka/server/ConfigHandler.scala
@@ -23,7 +23,6 @@ import DynamicConfig.Broker._
 import kafka.controller.KafkaController
 import kafka.log.UnifiedLog
 import kafka.network.ConnectionQuotas
-import kafka.security.CredentialProvider
 import kafka.server.Constants._
 import kafka.server.QuotaFactory.QuotaManagers
 import kafka.utils.Implicits._
@@ -33,6 +32,7 @@ import org.apache.kafka.common.config.internals.QuotaConfigs
 import org.apache.kafka.common.metrics.Quota
 import org.apache.kafka.common.metrics.Quota._
 import org.apache.kafka.common.utils.Sanitizer
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.ClientMetricsManager
 import org.apache.kafka.server.config.ZooKeeperInternals
 import org.apache.kafka.storage.internals.log.{LogConfig, ThrottledReplicaListValidator}
diff --git a/core/src/main/scala/kafka/server/ControllerServer.scala b/core/src/main/scala/kafka/server/ControllerServer.scala
index a882e54e52da8..ac74eb858c38d 100644
--- a/core/src/main/scala/kafka/server/ControllerServer.scala
+++ b/core/src/main/scala/kafka/server/ControllerServer.scala
@@ -21,7 +21,6 @@ import kafka.metrics.LinuxIoMetricsCollector
 import kafka.migration.MigrationPropagator
 import kafka.network.{DataPlaneAcceptor, SocketServer}
 import kafka.raft.KafkaRaftManager
-import kafka.security.CredentialProvider
 import kafka.server.KafkaConfig.{AlterConfigPolicyClassNameProp, CreateTopicPolicyClassNameProp}
 import kafka.server.QuotaFactory.QuotaManagers
 
@@ -44,7 +43,7 @@ import org.apache.kafka.metadata.bootstrap.BootstrapMetadata
 import org.apache.kafka.metadata.migration.{KRaftMigrationDriver, LegacyPropagator}
 import org.apache.kafka.metadata.publisher.FeaturesPublisher
 import org.apache.kafka.raft.RaftConfig
-import org.apache.kafka.security.PasswordEncoder
+import org.apache.kafka.security.{CredentialProvider, PasswordEncoder}
 import org.apache.kafka.server.NodeToControllerChannelManager
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.common.ApiMessageAndVersion
diff --git a/core/src/main/scala/kafka/server/KafkaBroker.scala b/core/src/main/scala/kafka/server/KafkaBroker.scala
index e281087f12f51..ff86729ab3416 100644
--- a/core/src/main/scala/kafka/server/KafkaBroker.scala
+++ b/core/src/main/scala/kafka/server/KafkaBroker.scala
@@ -22,7 +22,6 @@ import kafka.log.LogManager
 import kafka.log.remote.RemoteLogManager
 import kafka.metrics.LinuxIoMetricsCollector
 import kafka.network.SocketServer
-import kafka.security.CredentialProvider
 import kafka.utils.Logging
 import org.apache.kafka.common.ClusterResource
 import org.apache.kafka.common.internals.ClusterResourceListeners
@@ -32,6 +31,7 @@ import org.apache.kafka.common.security.token.delegation.internals.DelegationTok
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.coordinator.group.GroupCoordinator
 import org.apache.kafka.metadata.BrokerState
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.NodeToControllerChannelManager
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.metrics.{KafkaMetricsGroup, KafkaYammerMetrics}
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index 03f4fe9363267..233a6b8eb989b 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -27,7 +27,6 @@ import kafka.log.remote.RemoteLogManager
 import kafka.metrics.KafkaMetricsReporter
 import kafka.network.{ControlPlaneAcceptor, DataPlaneAcceptor, RequestChannel, SocketServer}
 import kafka.raft.KafkaRaftManager
-import kafka.security.CredentialProvider
 import kafka.server.metadata.{OffsetTrackingListener, ZkConfigRepository, ZkMetadataCache}
 import kafka.utils._
 import kafka.zk.{AdminZkClient, BrokerInfo, KafkaZkClient}
@@ -53,6 +52,7 @@ import org.apache.kafka.metadata.properties.MetaPropertiesEnsemble.VerificationF
 import org.apache.kafka.metadata.properties.{MetaProperties, MetaPropertiesEnsemble}
 import org.apache.kafka.metadata.{BrokerState, MetadataRecordSerde, VersionRange}
 import org.apache.kafka.raft.RaftConfig
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.NodeToControllerChannelManager
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.common.MetadataVersion._
diff --git a/core/src/main/scala/kafka/server/metadata/ScramPublisher.scala b/core/src/main/scala/kafka/server/metadata/ScramPublisher.scala
index bacac660e7b2f..09789249571f0 100644
--- a/core/src/main/scala/kafka/server/metadata/ScramPublisher.scala
+++ b/core/src/main/scala/kafka/server/metadata/ScramPublisher.scala
@@ -17,11 +17,11 @@
 
 package kafka.server.metadata
 
-import kafka.security.CredentialProvider
 import kafka.server.KafkaConfig
 import kafka.utils.Logging
 import org.apache.kafka.image.loader.LoaderManifest
 import org.apache.kafka.image.{MetadataDelta, MetadataImage}
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.fault.FaultHandler
 
 
diff --git a/core/src/main/scala/kafka/tools/TestRaftServer.scala b/core/src/main/scala/kafka/tools/TestRaftServer.scala
index 420b935dc299e..092ec67b21995 100644
--- a/core/src/main/scala/kafka/tools/TestRaftServer.scala
+++ b/core/src/main/scala/kafka/tools/TestRaftServer.scala
@@ -22,7 +22,6 @@ import java.util.concurrent.{CompletableFuture, CountDownLatch, LinkedBlockingDe
 import joptsimple.{OptionException, OptionSpec}
 import kafka.network.{DataPlaneAcceptor, SocketServer}
 import kafka.raft.{KafkaRaftManager, RaftManager}
-import kafka.security.CredentialProvider
 import kafka.server.{KafkaConfig, KafkaRequestHandlerPool, SimpleApiVersionManager}
 import kafka.utils.{CoreUtils, Exit, Logging}
 import org.apache.kafka.common.errors.InvalidConfigurationException
@@ -37,6 +36,7 @@ import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{TopicPartition, Uuid, protocol}
 import org.apache.kafka.raft.errors.NotLeaderException
 import org.apache.kafka.raft.{Batch, BatchReader, LeaderAndEpoch, RaftClient, RaftConfig}
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.common.{Features, MetadataVersion}
 import org.apache.kafka.server.common.serialization.RecordSerde
 import org.apache.kafka.server.fault.ProcessTerminatingFaultHandler
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index e2303efa90824..e23450b954d76 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -22,7 +22,6 @@ import kafka.api.LeaderAndIsr
 import kafka.cluster.Broker
 import kafka.controller.{KafkaController, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.security.authorizer.AclAuthorizer.{NoAcls, VersionedAcls}
-import kafka.security.authorizer.AclEntry
 import kafka.server.KafkaConfig
 import kafka.utils.Logging
 import kafka.zk.TopicZNode.TopicIdReplicaAssignment
@@ -34,6 +33,7 @@ import org.apache.kafka.common.security.token.delegation.{DelegationToken, Token
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{KafkaException, TopicPartition, Uuid}
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.server.metrics.KafkaMetricsGroup
 import org.apache.kafka.storage.internals.log.LogConfig
@@ -1292,7 +1292,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
   def createAclPaths(): Unit = {
     ZkAclStore.stores.foreach(store => {
       createRecursive(store.aclPath, throwIfPathExists = false)
-      AclEntry.ResourceTypes.foreach(resourceType => createRecursive(store.path(resourceType), throwIfPathExists = false))
+      AclEntry.RESOURCE_TYPES.forEach(resourceType => createRecursive(store.path(resourceType), throwIfPathExists = false))
     })
 
     ZkAclChangeStore.stores.foreach(store => createRecursive(store.aclChangePath, throwIfPathExists = false))
diff --git a/core/src/main/scala/kafka/zk/ZkData.scala b/core/src/main/scala/kafka/zk/ZkData.scala
index 3e830f97f39bd..d4bbc28ea2c7d 100644
--- a/core/src/main/scala/kafka/zk/ZkData.scala
+++ b/core/src/main/scala/kafka/zk/ZkData.scala
@@ -26,7 +26,6 @@ import kafka.cluster.{Broker, EndPoint}
 import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
 import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.security.authorizer.AclAuthorizer.VersionedAcls
-import kafka.security.authorizer.AclEntry
 import kafka.server.DelegationTokenManagerZk
 import kafka.utils.Json
 import kafka.utils.json.JsonObject
@@ -41,6 +40,7 @@ import org.apache.kafka.common.utils.{SecurityUtils, Time}
 import org.apache.kafka.common.{KafkaException, TopicPartition, Uuid}
 import org.apache.kafka.metadata.LeaderRecoveryState
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.common.{MetadataVersion, ProducerIdsBlock}
 import org.apache.kafka.server.common.MetadataVersion.{IBP_0_10_0_IV1, IBP_2_7_IV0}
 import org.apache.kafka.server.config.ConfigType
@@ -713,13 +713,13 @@ case object LiteralAclChangeStore extends ZkAclChangeStore {
     if (resource.patternType != PatternType.LITERAL)
       throw new IllegalArgumentException("Only literal resource patterns can be encoded")
 
-    val legacyName = resource.resourceType.toString + AclEntry.ResourceSeparator + resource.name
+    val legacyName = resource.resourceType.toString + AclEntry.RESOURCE_SEPARATOR + resource.name
     legacyName.getBytes(UTF_8)
   }
 
   def decode(bytes: Array[Byte]): ResourcePattern = {
     val string = new String(bytes, UTF_8)
-    string.split(AclEntry.ResourceSeparator, 2) match {
+    string.split(AclEntry.RESOURCE_SEPARATOR, 2) match {
         case Array(resourceType, resourceName, _*) => new ResourcePattern(ResourceType.fromString(resourceType), resourceName, PatternType.LITERAL)
         case _ => throw new IllegalArgumentException("expected a string in format ResourceType:ResourceName but got " + string)
       }
@@ -757,8 +757,8 @@ case object ExtendedAclChangeStore extends ZkAclChangeStore {
 object ResourceZNode {
   def path(resource: ResourcePattern): String = ZkAclStore(resource.patternType).path(resource.resourceType, resource.name)
 
-  def encode(acls: Set[AclEntry]): Array[Byte] = Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls).asJava)
-  def decode(bytes: Array[Byte], stat: Stat): VersionedAcls = VersionedAcls(AclEntry.fromBytes(bytes), stat.getVersion)
+  def encode(acls: Set[AclEntry]): Array[Byte] = Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls.asJava))
+  def decode(bytes: Array[Byte], stat: Stat): VersionedAcls = VersionedAcls(AclEntry.fromBytes(bytes).asScala.toSet, stat.getVersion)
 }
 
 object ExtendedAclChangeEvent {
diff --git a/core/src/main/scala/kafka/zk/migration/ZkAclMigrationClient.scala b/core/src/main/scala/kafka/zk/migration/ZkAclMigrationClient.scala
index 482476ff4bbd1..7bb4d074dd9b7 100644
--- a/core/src/main/scala/kafka/zk/migration/ZkAclMigrationClient.scala
+++ b/core/src/main/scala/kafka/zk/migration/ZkAclMigrationClient.scala
@@ -18,7 +18,7 @@
 package kafka.zk.migration
 
 import kafka.security.authorizer.AclAuthorizer.{ResourceOrdering, VersionedAcls}
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.utils.Logging
 import kafka.zk.ZkMigrationClient.{logAndRethrow, wrapZkException}
 import kafka.zk.{KafkaZkClient, ResourceZNode, ZkAclStore, ZkVersion}
@@ -26,6 +26,7 @@ import kafka.zookeeper.{CreateRequest, DeleteRequest, SetDataRequest}
 import org.apache.kafka.common.acl.AccessControlEntry
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.metadata.migration.{AclMigrationClient, MigrationClientException, ZkMigrationLeadershipState}
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.zookeeper.CreateMode
 import org.apache.zookeeper.KeeperException.Code
 
diff --git a/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
index b3e1ba9a64d3b..9b47c82000e02 100644
--- a/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AbstractAuthorizerIntegrationTest.scala
@@ -13,8 +13,8 @@
 package kafka.api
 
 import kafka.security.authorizer.AclAuthorizer
-import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server.{BaseRequestTest, KafkaConfig}
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.apache.kafka.clients.consumer.ConsumerConfig
 import org.apache.kafka.clients.producer.ProducerConfig
 import org.apache.kafka.common.TopicPartition
@@ -123,7 +123,7 @@ class AbstractAuthorizerIntegrationTest extends BaseRequestTest {
     doSetup(testInfo, createOffsetsTopic = false)
 
     // Allow inter-broker communication
-    addAndVerifyAcls(Set(new AccessControlEntry(brokerPrincipal.toString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(brokerPrincipal.toString, WILDCARD_HOST, CLUSTER_ACTION, ALLOW)), clusterResource)
 
     createOffsetsTopic(listenerName = interBrokerListenerName)
   }
diff --git a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
index ded470eba80d4..edee6a1de2d49 100644
--- a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
@@ -18,8 +18,6 @@ import java.util
 import java.util.concurrent.ExecutionException
 import java.util.regex.Pattern
 import java.util.{Collections, Optional, Properties}
-import kafka.security.authorizer.AclEntry
-import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.utils.{TestInfoUtils, TestUtils}
 import kafka.utils.TestUtils.waitUntilTrue
 import org.apache.kafka.clients.admin.{Admin, AlterConfigOp, NewTopic}
@@ -54,6 +52,8 @@ import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{ElectionType, IsolationLevel, KafkaException, Node, TopicPartition, Uuid, requests}
 import org.apache.kafka.test.{TestUtils => JTestUtils}
+import org.apache.kafka.security.authorizer.AclEntry
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.{CsvSource, ValueSource}
@@ -67,25 +67,25 @@ import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
 class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
-  val groupReadAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)))
-  val groupDescribeAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
-  val groupDeleteAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)))
-  val clusterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CLUSTER_ACTION, ALLOW)))
-  val clusterCreateAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)))
-  val clusterAlterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)))
-  val clusterDescribeAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
-  val clusterAlterConfigsAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER_CONFIGS, ALLOW)))
-  val clusterIdempotentWriteAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)))
-  val topicCreateAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)))
-  val topicReadAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)))
-  val topicWriteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)))
-  val topicDescribeAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
-  val topicAlterAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)))
-  val topicDeleteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)))
-  val topicDescribeConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE_CONFIGS, ALLOW)))
-  val topicAlterConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER_CONFIGS, ALLOW)))
-  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)))
-  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
+  val groupReadAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)))
+  val groupDescribeAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)))
+  val groupDeleteAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)))
+  val clusterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CLUSTER_ACTION, ALLOW)))
+  val clusterCreateAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)))
+  val clusterAlterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER, ALLOW)))
+  val clusterDescribeAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)))
+  val clusterAlterConfigsAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER_CONFIGS, ALLOW)))
+  val clusterIdempotentWriteAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, IDEMPOTENT_WRITE, ALLOW)))
+  val topicCreateAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)))
+  val topicReadAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)))
+  val topicWriteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)))
+  val topicDescribeAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)))
+  val topicAlterAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER, ALLOW)))
+  val topicDeleteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)))
+  val topicDescribeConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE_CONFIGS, ALLOW)))
+  val topicAlterConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER_CONFIGS, ALLOW)))
+  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)))
+  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)))
 
   val numRecords = 1
 
@@ -878,7 +878,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testProduceWithTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val producer = createProducer()
     assertThrows(classOf[TopicAuthorizationException], () => sendRecords(producer, numRecords, tp))
   }
@@ -887,7 +887,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testProduceWithTopicRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val producer = createProducer()
     assertThrows(classOf[TopicAuthorizationException], () => sendRecords(producer, numRecords, tp))
   }
@@ -896,7 +896,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testProduceWithTopicWrite(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, numRecords, tp)
   }
@@ -915,13 +915,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
   private def testCreatePermissionNeededToWriteToNonExistentTopic(resType: ResourceType): Unit = {
     val newTopicResource = new ResourcePattern(TOPIC, topic, LITERAL)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), newTopicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), newTopicResource)
     val producer = createProducer()
     val e = assertThrows(classOf[TopicAuthorizationException], () => sendRecords(producer, numRecords, tp))
     assertEquals(Collections.singleton(tp.topic), e.unauthorizedTopics())
 
     val resource = if (resType == ResourceType.TOPIC) newTopicResource else clusterResource
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)), resource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)), resource)
 
     sendRecords(producer, numRecords, tp)
   }
@@ -931,7 +931,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testConsumeUsingAssignWithNoAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
@@ -946,12 +946,12 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testSimpleConsumeWithOffsetLookupAndNoGroupAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
 
     // note this still depends on group access because we haven't set offsets explicitly, which means
     // they will first be fetched from the consumer coordinator (which requires group access)
@@ -966,12 +966,12 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testSimpleConsumeWithExplicitSeekAndNoGroupAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
 
     // in this case, we do an explicit seek, so there should be no need to query the coordinator at all
     // remove the group.id config to avoid coordinator created
@@ -986,12 +986,12 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testConsumeWithoutTopicDescribeAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
 
@@ -1004,13 +1004,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testConsumeWithTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
@@ -1023,13 +1023,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testConsumeWithTopicWrite(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
@@ -1042,13 +1042,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testConsumeWithTopicAndGroupRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
@@ -1061,12 +1061,12 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionWithNoTopicAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern))
@@ -1079,13 +1079,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionWithTopicDescribeOnlyAndGroupRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern))
     val e = assertThrows(classOf[TopicAuthorizationException], () => consumeRecords(consumer))
@@ -1098,26 +1098,26 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionWithTopicAndGroupRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
 
     // create an unmatched topic
     val unmatchedTopic = "unmatched"
     createTopicWithBrokerPrincipal(unmatchedTopic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)),  new ResourcePattern(TOPIC, unmatchedTopic, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)),  new ResourcePattern(TOPIC, unmatchedTopic, LITERAL))
     sendRecords(producer, 1, new TopicPartition(unmatchedTopic, part))
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern))
     consumeRecords(consumer)
 
     // set the subscription pattern to an internal topic that the consumer has read permission to. Since
     // internal topics are not included, we should not be assigned any partitions from this topic
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)),  new ResourcePattern(TOPIC,
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)),  new ResourcePattern(TOPIC,
       GROUP_METADATA_TOPIC_NAME, LITERAL))
     consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     consumer.poll(0)
@@ -1131,13 +1131,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionMatchingInternalTopic(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -1147,7 +1147,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     assertEquals(Set(topic).asJava, consumer.subscription)
 
     // now authorize the user for the internal topic and verify that we can subscribe
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), new ResourcePattern(TOPIC,
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), new ResourcePattern(TOPIC,
       GROUP_METADATA_TOPIC_NAME, LITERAL))
     consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     TestUtils.retry(60000) {
@@ -1161,15 +1161,15 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionMatchingInternalTopicWithDescribeOnlyPermission(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val internalTopicResource = new ResourcePattern(TOPIC, GROUP_METADATA_TOPIC_NAME, LITERAL)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), internalTopicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), internalTopicResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -1187,13 +1187,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testPatternSubscriptionNotMatchingInternalTopic(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, 1, tp)
     removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
     val consumer = createConsumer()
@@ -1205,7 +1205,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testCreatePermissionOnTopicToReadFromNonExistentTopic(quorum: String): Unit = {
     testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
-      Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)),
+      Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)),
       TOPIC)
   }
 
@@ -1213,14 +1213,14 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testCreatePermissionOnClusterToReadFromNonExistentTopic(quorum: String): Unit = {
     testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
-      Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)),
+      Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)),
       CLUSTER)
   }
 
   private def testCreatePermissionNeededToReadFromNonExistentTopic(newTopic: String, acls: Set[AccessControlEntry], resType: ResourceType): Unit = {
     val topicPartition = new TopicPartition(newTopic, 0)
     val newTopicResource = new ResourcePattern(TOPIC, newTopic, LITERAL)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), newTopicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), newTopicResource)
     addAndVerifyAcls(groupReadAcl(groupResource), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(topicPartition).asJava)
@@ -1274,7 +1274,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testCommitWithNoTopicAccess(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     assertThrows(classOf[TopicAuthorizationException], () => consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava))
   }
@@ -1284,8 +1284,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testCommitWithTopicWrite(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val consumer = createConsumer()
     assertThrows(classOf[TopicAuthorizationException], () => consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava))
   }
@@ -1295,8 +1295,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testCommitWithTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     assertThrows(classOf[TopicAuthorizationException], () => consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava))
   }
@@ -1304,7 +1304,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testCommitWithNoGroupAccess(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     assertThrows(classOf[GroupAuthorizationException], () => consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava))
   }
@@ -1313,8 +1313,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testCommitWithTopicAndGroupRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
@@ -1331,7 +1331,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testOffsetFetchWithNoGroupAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     assertThrows(classOf[GroupAuthorizationException], () => consumer.position(tp))
@@ -1340,7 +1340,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testOffsetFetchWithNoTopicAccess(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     assertThrows(classOf[TopicAuthorizationException], () => consumer.position(tp))
@@ -1352,14 +1352,14 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
 
     val offset = 15L
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(offset)).asJava)
 
     removeAllClientAcls()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     // send offset fetch requests directly since the consumer does not expose an API to do so
     // note there's only one broker, so no need to lookup the group coordinator
@@ -1371,7 +1371,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     assertTrue(offsetFetchResponse.partitionDataMap(group).isEmpty)
 
     // now add describe permission on the topic and verify that the offset can be fetched
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     assertEquals(Errors.NONE, offsetFetchResponse.groupLevelError(group))
     assertTrue(offsetFetchResponse.partitionDataMap(group).containsKey(tp))
@@ -1411,10 +1411,10 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topics(1), numPartitions = 2)
     createTopicWithBrokerPrincipal(topics(2), numPartitions = 3)
     groupResources.foreach(r => {
-      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), r)
+      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), r)
     })
     topicResources.foreach(t => {
-      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), t)
+      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), t)
     })
 
     val offset = 15L
@@ -1457,10 +1457,10 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // test handling partial errors, where one group is fully authorized, some groups don't have
     // the right topic authorizations, and some groups have no authorization
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResources(0))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResources(1))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResources(3))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResources(0))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResources(0))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResources(1))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResources(3))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResources(0))
     val offsetFetchRequest = createOffsetFetchRequest(groupToPartitionMap)
     var offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     offsetFetchResponse.data().groups().forEach(g =>
@@ -1491,9 +1491,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // test that after adding some of the ACLs, we get no group level authorization errors, but
     // still get topic level authorization errors for topics we don't have ACLs for
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResources(2))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResources(4))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResources(1))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResources(2))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResources(4))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResources(1))
     offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     offsetFetchResponse.data().groups().forEach(g =>
       g.groupId() match {
@@ -1527,7 +1527,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // test that after adding all necessary ACLs, we get no partition level or group level errors
     // from the offsetFetch response
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResources(2))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResources(2))
     offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     offsetFetchResponse.data.groups.asScala.map(_.groupId).foreach( groupId =>
       verifyResponse(offsetFetchResponse.groupLevelError(groupId), offsetFetchResponse.partitionDataMap(groupId), partitionMap(groupId))
@@ -1538,8 +1538,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testOffsetFetchTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
@@ -1549,8 +1549,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testOffsetFetchWithTopicAndGroupRead(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.position(tp)
@@ -1567,7 +1567,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testMetadataWithTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.partitionsFor(topic)
   }
@@ -1583,7 +1583,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testListOffsetsWithTopicDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.endOffsets(Set(tp).asJava)
   }
@@ -1591,7 +1591,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testDescribeGroupApiWithNoGroupAcl(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val result = createAdminClient().describeConsumerGroups(Seq(group).asJava)
     TestUtils.assertFutureExceptionTypeEquals(result.describedGroups().get(group), classOf[GroupAuthorizationException])
   }
@@ -1600,8 +1600,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testDescribeGroupApiWithGroupDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     createAdminClient().describeConsumerGroups(Seq(group).asJava).describedGroups().get(group).get()
   }
 
@@ -1611,15 +1611,15 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
 
     // write some record to the topic
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = createProducer()
     sendRecords(producer, numRecords = 1, tp)
 
     // use two consumers to write to two different groups
     val group2 = "other group"
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), new ResourcePattern(GROUP, group2, LITERAL))
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), new ResourcePattern(GROUP, group2, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.subscribe(Collections.singleton(topic))
     consumeRecords(consumer)
@@ -1634,13 +1634,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // first use cluster describe permission
     removeAllClientAcls()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), clusterResource)
     // it should list both groups (due to cluster describe permission)
     assertEquals(Set(group, group2), adminClient.listConsumerGroups().all().get().asScala.map(_.groupId()).toSet)
 
     // now replace cluster describe with group read permission
     removeAllClientAcls()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     // it should list only one group now
     val groupList = adminClient.listConsumerGroups().all().get().asScala.toList
     assertEquals(1, groupList.length)
@@ -1659,9 +1659,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testDeleteGroupApiWithDeleteGroupAcl(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)), groupResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1673,8 +1673,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testDeleteGroupApiWithNoDeleteGroupAcl(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1694,9 +1694,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testDeleteGroupOffsetsWithAcl(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1710,8 +1710,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testDeleteGroupOffsetsWithoutDeleteAcl(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1725,8 +1725,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testDeleteGroupOffsetsWithDeleteAclWithoutTopicAcl(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
     // Create the consumer group
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
@@ -1734,8 +1734,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // Remove the topic ACL & Check that it does not work without it
     removeAllClientAcls()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
     val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
     TestUtils.assertFutureExceptionTypeEquals(result.all(), classOf[TopicAuthorizationException])
     TestUtils.assertFutureExceptionTypeEquals(result.partitionResult(tp), classOf[TopicAuthorizationException])
@@ -1759,7 +1759,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testUnauthorizedDeleteTopicsWithDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val deleteResponse = connectAndReceive[DeleteTopicsResponse](deleteTopicsRequest)
     assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteResponse.data.responses.find(topic).errorCode)
   }
@@ -1768,7 +1768,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testDeleteTopicsWithWildCardAuth(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val deleteResponse = connectAndReceive[DeleteTopicsResponse](deleteTopicsRequest)
     assertEquals(Errors.NONE.code, deleteResponse.data.responses.find(topic).errorCode)
   }
@@ -1785,7 +1785,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testUnauthorizedDeleteRecordsWithDescribe(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val deleteRecordsResponse = connectAndReceive[DeleteRecordsResponse](deleteRecordsRequest)
     assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteRecordsResponse.data.topics.asScala.head.
       partitions.asScala.head.errorCode)
@@ -1795,7 +1795,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testDeleteRecordsWithWildCardAuth(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val deleteRecordsResponse = connectAndReceive[DeleteRecordsResponse](deleteRecordsRequest)
     assertEquals(Errors.NONE.code, deleteRecordsResponse.data.topics.asScala.head.
       partitions.asScala.head.errorCode)
@@ -1812,7 +1812,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testCreatePartitionsWithWildCardAuth(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
     val createPartitionsResponse = connectAndReceive[CreatePartitionsResponse](createPartitionsRequest)
     assertEquals(Errors.NONE.code, createPartitionsResponse.data.results.asScala.head.errorCode)
   }
@@ -1820,7 +1820,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testTransactionalProducerInitTransactionsNoWriteTransactionalIdAcl(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     assertThrows(classOf[TransactionalIdAuthorizationException], () => producer.initTransactions())
   }
@@ -1837,9 +1837,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testSendOffsetsWithNoConsumerGroupDescribeAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CLUSTER_ACTION, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1853,8 +1853,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testSendOffsetsWithNoConsumerGroupWriteAccess(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1867,7 +1867,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testIdempotentProducerNoIdempotentWriteAclInInitProducerId(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), topicResource)
     assertIdempotentSendAuthorizationFailure()
   }
 
@@ -1906,8 +1906,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testIdempotentProducerNoIdempotentWriteAclInProduce(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
     idempotentProducerShouldFailInProduce(() => removeAllClientAcls())
   }
 
@@ -1919,7 +1919,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // revoke the IdempotentWrite permission
     removeAclIdempotenceRequired()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
 
     // the send should now fail with a cluster auth error
     var e = assertThrows(classOf[ExecutionException], () => producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, "hi".getBytes)).get())
@@ -1934,7 +1934,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def shouldInitTransactionsWhenAclSet(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
   }
@@ -1944,9 +1944,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testTransactionalProducerTopicAuthorizationExceptionInSendCallback(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1961,9 +1961,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def testTransactionalProducerTopicAuthorizationExceptionInCommit(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -1979,11 +1979,11 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessDuringSend(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     removeAllClientAcls()
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     producer.beginTransaction()
     val future = producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes))
     JTestUtils.assertFutureThrows(future, classOf[TransactionalIdAuthorizationException])
@@ -1994,8 +1994,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnEndTransaction(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
 
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -2008,8 +2008,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testListTransactionsAuthorization(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
 
     // Start a transaction and write to a topic.
     val producer = buildTransactionalProducer()
@@ -2030,11 +2030,11 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     assertListTransactionResult(expectedTransactionalIds = Set(transactionalId))
 
     // Now revoke authorization and verify that the transaction is no longer listable
-    removeAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    removeAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
     assertListTransactionResult(expectedTransactionalIds = Set())
 
     // The minimum permission needed is `Describe`
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), transactionalIdResource)
     assertListTransactionResult(expectedTransactionalIds = Set(transactionalId))
   }
 
@@ -2042,8 +2042,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def shouldNotIncludeUnauthorizedTopicsInDescribeTransactionsResponse(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
 
     // Start a transaction and write to a topic.
     val producer = buildTransactionalProducer()
@@ -2053,7 +2053,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     // Remove only topic authorization so that we can verify that the
     // topic does not get included in the response.
-    removeAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    removeAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val response = connectAndReceive[DescribeTransactionsResponse](describeTransactionsRequest)
     assertEquals(1, response.data.transactionStates.size)
     val transactionStateData = response.data.transactionStates.asScala.find(_.transactionalId == transactionalId).get
@@ -2065,9 +2065,9 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def shouldSuccessfullyAbortTransactionAfterTopicAuthorizationException(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)),
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)),
       new ResourcePattern(TOPIC, topic, LITERAL))
     val producer = buildTransactionalProducer()
     producer.initTransactions()
@@ -2084,8 +2084,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnSendOffsetsToTxn(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
@@ -2101,8 +2101,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def shouldSendSuccessfullyWhenIdempotentAndHasCorrectACL(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
     val producer = buildIdempotentProducer()
     producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
   }
@@ -2124,8 +2124,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     val wildcard = new ResourcePattern(TOPIC, ResourcePattern.WILDCARD_RESOURCE, LITERAL)
     val prefixed = new ResourcePattern(TOPIC, "t", PREFIXED)
     val literal = new ResourcePattern(TOPIC, topic, LITERAL)
-    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)
-    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, DENY)
+    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)
+    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, DENY)
     val producer = buildIdempotentProducer()
 
     addAndVerifyAcls(Set(denyWriteAce), wildcard)
@@ -2147,14 +2147,14 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
 
     for (_ <- 1 to 3) {
-      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
       assertIdempotentSendAuthorizationFailure()
 
-      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)), topicResource)
       assertIdempotentSendSuccess()
 
       removeAllClientAcls()
-      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+      addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)), topicResource)
       assertIdempotentSendAuthorizationFailure()
     }
   }
@@ -2170,11 +2170,11 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     val unrelatedTopicResource = new ResourcePattern(TOPIC, "topic-2", LITERAL)
     val unrelatedGroupResource = new ResourcePattern(GROUP, "to", PREFIXED)
 
-    val acl1 = new AccessControlEntry(clientPrincipalString, WildcardHost, READ, DENY)
-    val acl2 = new AccessControlEntry(unrelatedPrincipalString, WildcardHost, READ, DENY)
-    val acl3 = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, DENY)
-    val acl4 = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)
-    val acl5 = new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)
+    val acl1 = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, DENY)
+    val acl2 = new AccessControlEntry(unrelatedPrincipalString, WILDCARD_HOST, READ, DENY)
+    val acl3 = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, DENY)
+    val acl4 = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)
+    val acl5 = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW)
 
     addAndVerifyAcls(Set(acl1, acl4, acl5), topicResource)
     addAndVerifyAcls(Set(acl2, acl3), unrelatedTopicResource)
@@ -2186,11 +2186,11 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ValueSource(strings = Array("zk", "kraft"))
   def testAuthorizeByResourceTypeDenyTakesPrecedence(quorum: String): Unit = {
     createTopicWithBrokerPrincipal(topic)
-    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)
+    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)
     addAndVerifyAcls(Set(allowWriteAce), topicResource)
     assertIdempotentSendSuccess()
 
-    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, DENY)
+    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, DENY)
     addAndVerifyAcls(Set(denyWriteAce), topicResource)
     assertIdempotentSendAuthorizationFailure()
   }
@@ -2202,8 +2202,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     val wildcard = new ResourcePattern(TOPIC, ResourcePattern.WILDCARD_RESOURCE, LITERAL)
     val prefixed = new ResourcePattern(TOPIC, "t", PREFIXED)
     val literal = new ResourcePattern(TOPIC, topic, LITERAL)
-    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)
-    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, DENY)
+    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)
+    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, DENY)
 
     addAndVerifyAcls(Set(allowWriteAce), prefixed)
     addAndVerifyAcls(Set(allowWriteAce), literal)
@@ -2219,8 +2219,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
     val prefixed = new ResourcePattern(TOPIC, topic.substring(0, 1), PREFIXED)
     val literal = new ResourcePattern(TOPIC, topic, LITERAL)
-    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)
-    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, DENY)
+    val allowWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, ALLOW)
+    val denyWriteAce = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, WRITE, DENY)
 
     addAndVerifyAcls(Set(denyWriteAce), prefixed)
     addAndVerifyAcls(Set(allowWriteAce), literal)
@@ -2245,8 +2245,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     val clusterResource = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
     val acls = Set(
-      new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW),
-      new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)
+      new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW),
+      new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER, ALLOW)
     )
     addAndVerifyAcls(acls, clusterResource)
 
@@ -2265,7 +2265,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
     removeAllClientAcls()
 
-    val allowAllOpsAcl = new AccessControlEntry(clientPrincipalString, WildcardHost, ALL, ALLOW)
+    val allowAllOpsAcl = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALL, ALLOW)
     addAndVerifyAcls(Set(allowAllOpsAcl), topicResource)
 
     val metadataRequestTopic = new MetadataRequestTopic()
@@ -2290,7 +2290,7 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     createTopicWithBrokerPrincipal(topic)
     removeAllClientAcls()
 
-    val allowAllOpsAcl = new AccessControlEntry(clientPrincipalString, WildcardHost, ALL, ALLOW)
+    val allowAllOpsAcl = new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALL, ALLOW)
     addAndVerifyAcls(Set(allowAllOpsAcl), topicResource)
 
     val describeConfigsRequest = new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData()
@@ -2341,8 +2341,8 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
 
     val clusterResource = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
     val acls = Set(
-      new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW),
-      new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)
+      new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, DESCRIBE, ALLOW),
+      new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, ALTER, ALLOW)
     )
     addAndVerifyAcls(acls, clusterResource)
 
@@ -2527,11 +2527,11 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array("zk", "kraft"))
   def testPrefixAcls(quorum: String): Unit = {
-    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)),
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, CREATE, ALLOW)),
       new ResourcePattern(TOPIC, "f", PREFIXED))
-    addAndVerifyAcls(Set(new AccessControlEntry("User:otherPrincipal", WildcardHost, CREATE, DENY)),
+    addAndVerifyAcls(Set(new AccessControlEntry("User:otherPrincipal", WILDCARD_HOST, CREATE, DENY)),
       new ResourcePattern(TOPIC, "fooa", PREFIXED))
-    addAndVerifyAcls(Set(new AccessControlEntry("User:otherPrincipal", WildcardHost, CREATE, ALLOW)),
+    addAndVerifyAcls(Set(new AccessControlEntry("User:otherPrincipal", WILDCARD_HOST, CREATE, ALLOW)),
       new ResourcePattern(TOPIC, "foob", PREFIXED))
     createAdminClient().createTopics(Collections.
       singletonList(new NewTopic("foobar", 1, 1.toShort))).all().get()
diff --git a/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala
index 6fa8bc9b2bc55..e7b5089a0d2db 100644
--- a/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala
@@ -19,7 +19,6 @@ package kafka.api
 import java.util
 import java.util.Properties
 import java.util.concurrent.ExecutionException
-import kafka.security.authorizer.AclEntry
 import kafka.server.KafkaConfig
 import kafka.utils.Logging
 import kafka.utils.TestUtils._
@@ -29,6 +28,7 @@ import org.apache.kafka.common.acl.AclOperation
 import org.apache.kafka.common.errors.{TopicExistsException, UnknownTopicOrPartitionException}
 import org.apache.kafka.common.resource.ResourceType
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.security.authorizer.AclEntry
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo, Timeout}
 
@@ -99,7 +99,7 @@ abstract class BaseAdminIntegrationTest extends IntegrationTestHarness with Logg
       assertNotEquals(Uuid.ZERO_UUID, createResult.topicId(topic).get())
       assertEquals(topicIds(topic), createResult.topicId(topic).get())
     }
-    
+
 
     val failedCreateResult = client.createTopics(newTopics.asJava)
     val results = failedCreateResult.values()
@@ -183,12 +183,12 @@ abstract class BaseAdminIntegrationTest extends IntegrationTestHarness with Logg
 
     //with includeAuthorizedOperations flag
     topicResult = getTopicMetadata(client, topic, new DescribeTopicsOptions().includeAuthorizedOperations(true))
-    expectedOperations = AclEntry.supportedOperations(ResourceType.TOPIC).asJava
+    expectedOperations = AclEntry.supportedOperations(ResourceType.TOPIC)
     assertEquals(expectedOperations, topicResult.authorizedOperations)
   }
 
   def configuredClusterPermissions: Set[AclOperation] =
-    AclEntry.supportedOperations(ResourceType.CLUSTER)
+    AclEntry.supportedOperations(ResourceType.CLUSTER).asScala.toSet
 
   override def modifyConfigs(configs: Seq[Properties]): Unit = {
     super.modifyConfigs(configs)
diff --git a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
index 9469a2d4e4d18..0be5e306e2948 100644
--- a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
@@ -15,7 +15,7 @@ package kafka.api
 import java.util
 import java.util.Properties
 
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.server.KafkaConfig
 import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
 import org.apache.kafka.clients.admin._
@@ -25,6 +25,7 @@ import org.apache.kafka.common.acl._
 import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.config.ZkConfigs
 import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNull}
@@ -68,7 +69,7 @@ object DescribeAuthorizedOperationsTest {
     operation: AclOperation
   ): AccessControlEntry = {
     new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, userName).toString,
-      AclEntry.WildcardHost, operation, ALLOW)
+      AclEntry.WILDCARD_HOST, operation, ALLOW)
   }
 }
 
@@ -88,7 +89,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
   override def configureSecurityBeforeServersStart(testInfo: TestInfo): Unit = {
     val authorizer = CoreUtils.createObject[Authorizer](classOf[AclAuthorizer].getName)
     val clusterResource = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
-    val topicResource = new ResourcePattern(ResourceType.TOPIC, AclEntry.WildcardResource, PatternType.LITERAL)
+    val topicResource = new ResourcePattern(ResourceType.TOPIC, AclEntry.WILDCARD_RESOURCE, PatternType.LITERAL)
 
     try {
       authorizer.configure(this.configs.head.originals())
@@ -140,7 +141,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
     val describeConsumerGroupsResult = client.describeConsumerGroups(Seq(Group1, Group2, Group3).asJava,
       new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
     assertEquals(3, describeConsumerGroupsResult.describedGroups().size())
-    val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP).asJava
+    val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP)
 
     val group1Description = describeConsumerGroupsResult.describedGroups().get(Group1).get
     assertEquals(expectedOperations, group1Description.authorizedOperations())
@@ -170,7 +171,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
     assertEquals(Set(ClusterAllAcl), results.values.keySet.asScala)
     results.all.get
 
-    val expectedOperations = AclEntry.supportedOperations(ResourceType.CLUSTER).asJava
+    val expectedOperations = AclEntry.supportedOperations(ResourceType.CLUSTER)
 
     clusterDescribeResult = client.describeCluster(new DescribeClusterOptions().
       includeAuthorizedOperations(true))
@@ -198,7 +199,7 @@ class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslS
     assertEquals(Set(Topic1Acl, Topic2All), results.values.keySet.asScala)
     results.all.get
 
-    val expectedOperations = AclEntry.supportedOperations(ResourceType.TOPIC).asJava
+    val expectedOperations = AclEntry.supportedOperations(ResourceType.TOPIC)
 
     describeTopicsResult = client.describeTopics(Set(Topic1, Topic2).asJava,
       new DescribeTopicsOptions().includeAuthorizedOperations(true)).allTopicNames.get()
diff --git a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
index 0ac83d640a3d4..023b5c38825b1 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
@@ -22,7 +22,6 @@ import com.yammer.metrics.core.Gauge
 import java.util.{Collections, Properties}
 import java.util.concurrent.ExecutionException
 import kafka.security.authorizer.AclAuthorizer
-import kafka.security.authorizer.AclEntry.WildcardHost
 import org.apache.kafka.metadata.authorizer.StandardAuthorizer
 import kafka.server._
 import kafka.utils._
@@ -38,6 +37,7 @@ import org.apache.kafka.common.resource._
 import org.apache.kafka.common.resource.ResourceType._
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.security.auth._
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.server.metrics.KafkaYammerMetrics
 import org.junit.jupiter.api.Assertions._
@@ -93,11 +93,11 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
   def clientPrincipal: KafkaPrincipal
   def kafkaPrincipal: KafkaPrincipal
 
-  def GroupReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, READ, ALLOW))
-  def TopicReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, READ, ALLOW))
-  def TopicWriteAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, WRITE, ALLOW))
-  def TopicDescribeAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, DESCRIBE, ALLOW))
-  def TopicCreateAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, CREATE, ALLOW))
+  def GroupReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WILDCARD_HOST, READ, ALLOW))
+  def TopicReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WILDCARD_HOST, READ, ALLOW))
+  def TopicWriteAcl = Set(new AccessControlEntry(clientPrincipal.toString, WILDCARD_HOST, WRITE, ALLOW))
+  def TopicDescribeAcl = Set(new AccessControlEntry(clientPrincipal.toString, WILDCARD_HOST, DESCRIBE, ALLOW))
+  def TopicCreateAcl = Set(new AccessControlEntry(clientPrincipal.toString, WILDCARD_HOST, CREATE, ALLOW))
 
   def AclTopicWrite(topicResource : ResourcePattern = topicResource) = new AclBinding(topicResource,
     new AccessControlEntry(clientPrincipal.toString, "*", AclOperation.WRITE, AclPermissionType.ALLOW))
diff --git a/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
index 82e637ae00281..fd9a5c01ab29a 100644
--- a/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
@@ -16,7 +16,6 @@ import java.util.Properties
 import java.util.concurrent.ExecutionException
 import kafka.api.GroupAuthorizerIntegrationTest._
 import kafka.security.authorizer.AclAuthorizer
-import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server.{BaseRequestTest, KafkaConfig}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.consumer.ConsumerConfig
@@ -30,6 +29,7 @@ import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern,
 import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal}
 import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder
 import org.apache.kafka.metadata.authorizer.StandardAuthorizer
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.{BeforeEach, TestInfo}
 import org.junit.jupiter.params.ParameterizedTest
@@ -110,7 +110,7 @@ class GroupAuthorizerIntegrationTest extends BaseRequestTest {
   private def createAcl(aclOperation: AclOperation,
                         aclPermissionType: AclPermissionType,
                         principal: KafkaPrincipal = ClientPrincipal): AccessControlEntry = {
-    new AccessControlEntry(principal.toString, WildcardHost, aclOperation, aclPermissionType)
+    new AccessControlEntry(principal.toString, WILDCARD_HOST, aclOperation, aclPermissionType)
   }
 
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
index b8690d238dc05..d997956ac0de1 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -26,7 +26,6 @@ import java.util.concurrent.{CountDownLatch, ExecutionException, TimeUnit}
 import java.util.{Collections, Optional, Properties}
 import java.{time, util}
 import kafka.integration.KafkaServerTestHarness
-import kafka.security.authorizer.AclEntry
 import kafka.server.metadata.KRaftMetadataCache
 import kafka.server.{DynamicConfig, KafkaConfig}
 import kafka.utils.TestUtils._
@@ -44,6 +43,7 @@ import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceT
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{ConsumerGroupState, ElectionType, TopicCollection, TopicPartition, TopicPartitionInfo, TopicPartitionReplica, Uuid}
 import org.apache.kafka.controller.ControllerRequestContextUtil.ANONYMOUS_CONTEXT
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.config.{Defaults, ZkConfigs}
 import org.apache.kafka.storage.internals.log.{CleanerConfig, LogConfig}
 import org.junit.jupiter.api.Assertions._
@@ -1226,7 +1226,7 @@ class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
             assertEquals(testNumPartitions, topicPartitions.size)
           }
 
-          val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP).asJava
+          val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP)
           assertEquals(expectedOperations, testGroupDescription.authorizedOperations())
 
           // Test that the fake group is listed as dead.
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
index 22579d92eafca..47e8de08031b4 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
@@ -14,7 +14,6 @@ package kafka.api
 
 import java.util
 import kafka.security.authorizer.AclAuthorizer
-import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
 import kafka.server.KafkaConfig
 import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
 import kafka.utils.TestUtils._
@@ -29,6 +28,7 @@ import org.apache.kafka.common.resource.PatternType.LITERAL
 import org.apache.kafka.common.resource.ResourceType.{GROUP, TOPIC}
 import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourcePatternFilter, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.security.authorizer.AclEntry.{WILDCARD_HOST, WILDCARD_PRINCIPAL_STRING}
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.kafka.server.config.ZkConfigs
 import org.apache.kafka.storage.internals.log.LogConfig
@@ -483,7 +483,7 @@ class SaslSslAdminIntegrationTest extends BaseAdminIntegrationTest with SaslSetu
       val authorizer = CoreUtils.createObject[Authorizer](authorizerForInitClass.getName)
       try {
         authorizer.configure(configs.head.originals())
-        val ace = new AccessControlEntry(WildcardPrincipalString, WildcardHost, ALL, ALLOW)
+        val ace = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, ALL, ALLOW)
         authorizer.createAcls(null, List(new AclBinding(new ResourcePattern(TOPIC, "*", LITERAL), ace)).asJava)
         authorizer.createAcls(null, List(new AclBinding(new ResourcePattern(GROUP, "*", LITERAL), ace)).asJava)
 
@@ -521,7 +521,7 @@ class SaslSslAdminIntegrationTest extends BaseAdminIntegrationTest with SaslSetu
 
     private def clusterAcl(permissionType: AclPermissionType, operation: AclOperation): AccessControlEntry = {
       new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*").toString,
-        WildcardHost, operation, permissionType)
+        WILDCARD_HOST, operation, permissionType)
     }
   }
 }
diff --git a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
index e905e3a8a1098..07a7bee68450d 100644
--- a/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/zk/ZkMigrationIntegrationTest.scala
@@ -16,7 +16,6 @@
  */
 package kafka.zk
 
-import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
 import kafka.server.KafkaConfig
 import kafka.test.{ClusterConfig, ClusterGenerator, ClusterInstance}
 import kafka.test.annotation.{AutoStart, ClusterConfigProperty, ClusterTemplate, ClusterTest, Type}
@@ -45,6 +44,7 @@ import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
 import org.apache.kafka.metadata.authorizer.StandardAcl
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.raft.RaftConfig
+import org.apache.kafka.security.authorizer.AclEntry.{WILDCARD_HOST, WILDCARD_PRINCIPAL_STRING}
 import org.apache.kafka.security.PasswordEncoder
 import org.apache.kafka.server.ControllerRequestCompletionHandler
 import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
@@ -128,12 +128,12 @@ class ZkMigrationIntegrationTest {
     val prefixedResource = new ResourcePattern(TOPIC, "bar-", PREFIXED)
     val username = "alice"
     val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
-    val wildcardPrincipal = SecurityUtils.parseKafkaPrincipal(WildcardPrincipalString)
+    val wildcardPrincipal = SecurityUtils.parseKafkaPrincipal(WILDCARD_PRINCIPAL_STRING)
 
-    val acl1 = new AclBinding(resource1, new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW))
+    val acl1 = new AclBinding(resource1, new AccessControlEntry(principal.toString, WILDCARD_HOST, READ, ALLOW))
     val acl2 = new AclBinding(resource1, new AccessControlEntry(principal.toString, "192.168.0.1", WRITE, ALLOW))
-    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WildcardHost, DESCRIBE, ALLOW))
-    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WildcardHost, READ, ALLOW))
+    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WILDCARD_HOST, DESCRIBE, ALLOW))
+    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WILDCARD_HOST, READ, ALLOW))
 
     val result = admin.createAcls(List(acl1, acl2, acl3, acl4).asJava)
     result.all().get
diff --git a/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala b/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala
index bfddee2020ec1..5a0c3e714faa0 100644
--- a/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala
+++ b/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala
@@ -18,10 +18,10 @@
 package kafka.zk
 
 import java.nio.charset.StandardCharsets.UTF_8
-import kafka.security.authorizer.AclEntry
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.resource.ResourceType.{GROUP, TOPIC}
+import org.apache.kafka.security.authorizer.AclEntry
 import org.junit.jupiter.api.Assertions.{assertEquals, assertThrows}
 import org.junit.jupiter.api.Test
 
@@ -66,7 +66,7 @@ class LiteralAclStoreTest {
   @Test
   def shouldDecodeResourceUsingTwoPartLogic(): Unit = {
     val resource = new ResourcePattern(GROUP, "PREFIXED:this, including the PREFIXED part, is a valid two part group name", LITERAL)
-    val encoded = (resource.resourceType.toString + AclEntry.ResourceSeparator + resource.name).getBytes(UTF_8)
+    val encoded = (resource.resourceType.toString + AclEntry.RESOURCE_SEPARATOR + resource.name).getBytes(UTF_8)
 
     val actual = store.changeStore.decode(encoded)
 
diff --git a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
index 0bd452cf5cc28..19be739996186 100644
--- a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
@@ -20,7 +20,7 @@ import java.io.File
 import java.util.Properties
 import javax.management.InstanceAlreadyExistsException
 import kafka.admin.AclCommand.AclCommandOptions
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.utils.{Exit, LogCaptureAppender, Logging, TestUtils}
 import kafka.server.QuorumTestHarness
@@ -33,6 +33,7 @@ import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.{AppInfoParser, SecurityUtils}
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.authorizer.Authorizer
 import org.apache.log4j.Level
 import org.junit.jupiter.api.Assertions._
@@ -246,9 +247,9 @@ class AclCommandTest extends QuorumTestHarness with Logging {
     callMain(cmdArgs ++ cmd :+ "--add")
 
     withAuthorizer() { authorizer =>
-      val writeAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, WRITE, ALLOW)
-      val describeAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, DESCRIBE, ALLOW)
-      val createAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, CREATE, ALLOW)
+      val writeAcl = new AccessControlEntry(principal.toString, AclEntry.WILDCARD_HOST, WRITE, ALLOW)
+      val describeAcl = new AccessControlEntry(principal.toString, AclEntry.WILDCARD_HOST, DESCRIBE, ALLOW)
+      val createAcl = new AccessControlEntry(principal.toString, AclEntry.WILDCARD_HOST, CREATE, ALLOW)
       TestUtils.waitAndVerifyAcls(Set(writeAcl, describeAcl, createAcl), authorizer,
         new ResourcePattern(TOPIC, "Test-", PREFIXED))
     }
diff --git a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
index 8facbd9d4c5bc..af3f361479966 100644
--- a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
+++ b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
@@ -30,7 +30,6 @@ import com.yammer.metrics.core.{Gauge, Meter}
 
 import javax.net.ssl._
 import kafka.cluster.EndPoint
-import kafka.security.CredentialProvider
 import kafka.server.{ApiVersionManager, KafkaConfig, SimpleApiVersionManager, ThrottleCallback, ThrottledChannel}
 import kafka.utils.Implicits._
 import kafka.utils.TestUtils
@@ -46,6 +45,7 @@ import org.apache.kafka.common.requests._
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.utils.{AppInfoParser, LogContext, MockTime, Time, Utils}
+import org.apache.kafka.security.CredentialProvider
 import org.apache.kafka.server.common.{Features, MetadataVersion}
 import org.apache.kafka.test.{TestSslUtils, TestUtils => JTestUtils}
 import org.apache.log4j.Level
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala
index 74b1b757ecef1..da79a3c100a12 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala
@@ -24,7 +24,6 @@ import java.util.concurrent.{Executors, TimeUnit}
 import javax.security.auth.Subject
 import javax.security.auth.callback.CallbackHandler
 import kafka.api.SaslSetup
-import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server.{KafkaConfig, QuorumTestHarness}
 import kafka.utils.JaasTestUtils.{JaasModule, JaasSection}
 import kafka.utils.{JaasTestUtils, TestUtils}
@@ -41,6 +40,7 @@ import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.resource.ResourceType.TOPIC
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.test.{TestUtils => JTestUtils}
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.apache.zookeeper.server.auth.DigestLoginModule
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.{AfterEach, BeforeEach, Test, TestInfo}
@@ -116,8 +116,8 @@ class AclAuthorizerWithZkSaslTest extends QuorumTestHarness with SaslSetup {
   }
 
   private def verifyAclUpdate(): Unit = {
-    val allowReadAcl = new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW)
-    val allowWriteAcl = new AccessControlEntry(principal.toString, WildcardHost, WRITE, ALLOW)
+    val allowReadAcl = new AccessControlEntry(principal.toString, WILDCARD_HOST, READ, ALLOW)
+    val allowWriteAcl = new AccessControlEntry(principal.toString, WILDCARD_HOST, WRITE, ALLOW)
     val acls = Set(allowReadAcl, allowWriteAcl)
 
     TestUtils.retry(maxWaitMs = 15000) {
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala
index ddf00a564d6b1..f6867b6233b89 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala
@@ -17,15 +17,16 @@
 package kafka.security.authorizer
 
 import java.nio.charset.StandardCharsets.UTF_8
-
 import kafka.utils.Json
+import org.apache.kafka.common.acl.AccessControlEntry
 import org.apache.kafka.common.acl.AclOperation.READ
 import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.security.authorizer.AclEntry
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
 
-import scala.jdk.CollectionConverters._
+import java.util
 
 class AclEntryTest {
 
@@ -35,13 +36,13 @@ class AclEntryTest {
 
   @Test
   def testAclJsonConversion(): Unit = {
-    val acl1 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), DENY, "host1" , READ)
-    val acl2 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), ALLOW, "*", READ)
-    val acl3 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), DENY, "host1", READ)
+    val acl1 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice").toString, "host1", READ, DENY))
+    val acl2 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob").toString, "*", READ, ALLOW))
+    val acl3 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob").toString, "host1", READ, DENY))
 
-    val acls = Set[AclEntry](acl1, acl2, acl3)
+    val acls = new util.HashSet[AclEntry](util.Arrays.asList(acl1, acl2, acl3))
 
-    assertEquals(acls, AclEntry.fromBytes(Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls).asJava)))
+    assertEquals(acls, AclEntry.fromBytes(Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls))))
     assertEquals(acls, AclEntry.fromBytes(AclJson.getBytes(UTF_8)))
   }
 }
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
index 8b424d44944be..2f2fe4b20472d 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AuthorizerTest.scala
@@ -17,7 +17,6 @@
 package kafka.security.authorizer
 
 import kafka.Kafka
-import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
 import kafka.server.{KafkaConfig, QuorumTestHarness}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import kafka.zk.ZkAclStore
@@ -38,6 +37,7 @@ import org.apache.kafka.common.utils.{Time, SecurityUtils => JSecurityUtils}
 import org.apache.kafka.controller.MockAclMutator
 import org.apache.kafka.metadata.authorizer.StandardAuthorizerTest.AuthorizerTestServerInfo
 import org.apache.kafka.metadata.authorizer.StandardAuthorizer
+import org.apache.kafka.security.authorizer.AclEntry.{WILDCARD_HOST, WILDCARD_PRINCIPAL_STRING}
 import org.apache.kafka.server.authorizer._
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.common.MetadataVersion.{IBP_2_0_IV0, IBP_2_0_IV1}
@@ -64,14 +64,14 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
   private final val ZK = "zk"
 
 
-  private val allowReadAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, READ, ALLOW)
-  private val allowWriteAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, WRITE, ALLOW)
-  private val denyReadAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, READ, DENY)
+  private val allowReadAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, READ, ALLOW)
+  private val allowWriteAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, WRITE, ALLOW)
+  private val denyReadAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, READ, DENY)
 
   private val wildCardResource = new ResourcePattern(TOPIC, WILDCARD_RESOURCE, LITERAL)
   private val prefixedResource = new ResourcePattern(TOPIC, "foo", PREFIXED)
   private val clusterResource = new ResourcePattern(CLUSTER, CLUSTER_NAME, LITERAL)
-  private val wildcardPrincipal = JSecurityUtils.parseKafkaPrincipal(WildcardPrincipalString)
+  private val wildcardPrincipal = JSecurityUtils.parseKafkaPrincipal(WILDCARD_PRINCIPAL_STRING)
 
   private var authorizer1: Authorizer = _
   private var authorizer2: Authorizer = _
@@ -168,13 +168,13 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val acl4 = new AccessControlEntry(user1.toString, host1.getHostAddress, WRITE, ALLOW)
 
     //user1 has DESCRIBE access from all hosts.
-    val acl5 = new AccessControlEntry(user1.toString, WildcardHost, DESCRIBE, ALLOW)
+    val acl5 = new AccessControlEntry(user1.toString, WILDCARD_HOST, DESCRIBE, ALLOW)
 
     //user2 has READ access from all hosts.
-    val acl6 = new AccessControlEntry(user2.toString, WildcardHost, READ, ALLOW)
+    val acl6 = new AccessControlEntry(user2.toString, WILDCARD_HOST, READ, ALLOW)
 
     //user3 has WRITE access from all hosts.
-    val acl7 = new AccessControlEntry(user3.toString, WildcardHost, WRITE, ALLOW)
+    val acl7 = new AccessControlEntry(user3.toString, WILDCARD_HOST, WRITE, ALLOW)
 
     val acls = Set(acl1, acl2, acl3, acl4, acl5, acl6, acl7)
 
@@ -232,7 +232,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val host = InetAddress.getByName("192.168.2.1")
     val session = newRequestContext(user, host)
 
-    val allowAll = new AccessControlEntry(WildcardPrincipalString, WildcardHost, AclOperation.ALL, ALLOW)
+    val allowAll = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, AclOperation.ALL, ALLOW)
     val denyAcl = new AccessControlEntry(user.toString, host.getHostAddress, AclOperation.ALL, DENY)
     val acls = Set(allowAll, denyAcl)
 
@@ -244,7 +244,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array(KRAFT, ZK))
   def testAllowAllAccess(quorum: String): Unit = {
-    val allowAllAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, AclOperation.ALL, ALLOW)
+    val allowAllAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, AclOperation.ALL, ALLOW)
 
     changeAclAndVerify(Set.empty, Set(allowAllAcl), Set.empty)
 
@@ -255,7 +255,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array(KRAFT, ZK))
   def testSuperUserHasAccess(quorum: String): Unit = {
-    val denyAllAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, AclOperation.ALL, DENY)
+    val denyAllAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, AclOperation.ALL, DENY)
 
     changeAclAndVerify(Set.empty, Set(denyAllAcl), Set.empty)
 
@@ -272,7 +272,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array(KRAFT, ZK))
   def testSuperUserWithCustomPrincipalHasAccess(quorum: String): Unit = {
-    val denyAllAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, AclOperation.ALL, DENY)
+    val denyAllAcl = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, AclOperation.ALL, DENY)
     changeAclAndVerify(Set.empty, Set(denyAllAcl), Set.empty)
 
     val session = newRequestContext(new CustomPrincipal(KafkaPrincipal.USER_TYPE, "superuser1"), InetAddress.getByName("192.0.4.4"))
@@ -362,7 +362,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     var acls = changeAclAndVerify(Set.empty, Set(acl1, acl2, acl3, acl4), Set.empty)
 
     //test addAcl is additive
-    val acl5 = new AccessControlEntry(user2.toString, WildcardHost, READ, ALLOW)
+    val acl5 = new AccessControlEntry(user2.toString, WILDCARD_HOST, READ, ALLOW)
     acls = changeAclAndVerify(acls, Set(acl5), Set.empty)
 
     //test get by principal name.
@@ -372,7 +372,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
       "changes not propagated in timeout period")
 
     val resourceToAcls = Map[ResourcePattern, Set[AccessControlEntry]](
-      new ResourcePattern(TOPIC, WILDCARD_RESOURCE, LITERAL) -> Set(new AccessControlEntry(user2.toString, WildcardHost, READ, ALLOW)),
+      new ResourcePattern(TOPIC, WILDCARD_RESOURCE, LITERAL) -> Set(new AccessControlEntry(user2.toString, WILDCARD_HOST, READ, ALLOW)),
       new ResourcePattern(CLUSTER , WILDCARD_RESOURCE, LITERAL) -> Set(new AccessControlEntry(user2.toString, host1, READ, ALLOW)),
       new ResourcePattern(GROUP, WILDCARD_RESOURCE, LITERAL) -> acls,
       new ResourcePattern(GROUP, "test-ConsumerGroup", LITERAL) -> acls
@@ -466,10 +466,10 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val commonResource = new ResourcePattern(TOPIC, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
-    val acl1 = new AccessControlEntry(user1.toString, WildcardHost, READ, ALLOW)
+    val acl1 = new AccessControlEntry(user1.toString, WILDCARD_HOST, READ, ALLOW)
 
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
-    val acl2 = new AccessControlEntry(user2.toString, WildcardHost, READ, DENY)
+    val acl2 = new AccessControlEntry(user2.toString, WILDCARD_HOST, READ, DENY)
 
     addAcls(authorizer1, Set(acl1), commonResource)
     addAcls(authorizer1, Set(acl2), commonResource)
@@ -482,10 +482,10 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val commonResource = new ResourcePattern(TOPIC, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
-    val acl1 = new AccessControlEntry(user1.toString, WildcardHost, READ, ALLOW)
+    val acl1 = new AccessControlEntry(user1.toString, WILDCARD_HOST, READ, ALLOW)
 
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
-    val acl2 = new AccessControlEntry(user2.toString, WildcardHost, READ, DENY)
+    val acl2 = new AccessControlEntry(user2.toString, WILDCARD_HOST, READ, DENY)
 
     // Add on each instance
     addAcls(authorizer1, Set(acl1), commonResource)
@@ -495,7 +495,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     TestUtils.waitAndVerifyAcls(Set(acl1, acl2), authorizer2, commonResource)
 
     val user3 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "joe")
-    val acl3 = new AccessControlEntry(user3.toString, WildcardHost, READ, DENY)
+    val acl3 = new AccessControlEntry(user3.toString, WILDCARD_HOST, READ, DENY)
 
     // Add on one instance and delete on another
     addAcls(authorizer1, Set(acl3), commonResource)
@@ -513,7 +513,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
 
     val acls= (0 to 50).map { i =>
       val useri = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, i.toString)
-      (new AccessControlEntry(useri.toString, WildcardHost, READ, ALLOW), i)
+      (new AccessControlEntry(useri.toString, WILDCARD_HOST, READ, ALLOW), i)
     }
 
     // Alternate authorizer, Remove all acls that end in 0
@@ -563,7 +563,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host = InetAddress.getByName("192.168.3.1")
     val hostContext = newRequestContext(user, host)
-    val acl = new AccessControlEntry(user.toString, WildcardHost, parentOp, ALLOW)
+    val acl = new AccessControlEntry(user.toString, WILDCARD_HOST, parentOp, ALLOW)
     addAcls(authorizer1, Set(acl), clusterResource)
     AclOperation.values.filter(validOp).foreach { op =>
       val authorized = authorize(authorizer1, hostContext, op, clusterResource)
@@ -579,8 +579,8 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host1 = InetAddress.getByName("192.168.3.1")
     val host1Context = newRequestContext(user1, host1)
-    val acls = Set(new AccessControlEntry(user1.toString, WildcardHost, parentOp, DENY),
-      new AccessControlEntry(user1.toString, WildcardHost, AclOperation.ALL, ALLOW))
+    val acls = Set(new AccessControlEntry(user1.toString, WILDCARD_HOST, parentOp, DENY),
+      new AccessControlEntry(user1.toString, WILDCARD_HOST, AclOperation.ALL, ALLOW))
     addAcls(authorizer1, acls, clusterResource)
     AclOperation.values.filter(validOp).foreach { op =>
       val authorized = authorize(authorizer1, host1Context, op, clusterResource)
@@ -594,7 +594,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
 
   @Test
   def testHighConcurrencyDeletionOfResourceAcls(): Unit = {
-    val acl = new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username).toString, WildcardHost, AclOperation.ALL, ALLOW)
+    val acl = new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username).toString, WILDCARD_HOST, AclOperation.ALL, ALLOW)
 
     // Alternate authorizer to keep adding and removing ZooKeeper path
     val concurrentFunctions = (0 to 50).map { _ =>
@@ -736,7 +736,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
   @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedQuorumName)
   @ValueSource(strings = Array(KRAFT, ZK))
   def testGetAclsPrincipal(quorum: String): Unit = {
-    val aclOnSpecificPrincipal = new AccessControlEntry(principal.toString, WildcardHost, WRITE, ALLOW)
+    val aclOnSpecificPrincipal = new AccessControlEntry(principal.toString, WILDCARD_HOST, WRITE, ALLOW)
     addAcls(authorizer1, Set(aclOnSpecificPrincipal), resource)
 
     assertEquals(0,
@@ -747,7 +747,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
       getAcls(authorizer1, new KafkaPrincipal(principal.getPrincipalType, principal.getName)).size, "acl on specific should be returned for different principal instance")
 
     removeAcls(authorizer1, Set.empty, resource)
-    val aclOnWildcardPrincipal = new AccessControlEntry(WildcardPrincipalString, WildcardHost, WRITE, ALLOW)
+    val aclOnWildcardPrincipal = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, WRITE, ALLOW)
     addAcls(authorizer1, Set(aclOnWildcardPrincipal), resource)
 
     assertEquals(1, getAcls(authorizer1, wildcardPrincipal).size, "acl on wildcard should be returned for wildcard request")
@@ -761,10 +761,10 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val resource2 = new ResourcePattern(TOPIC, "bar-" + UUID.randomUUID(), LITERAL)
     val prefixedResource = new ResourcePattern(TOPIC, "bar-", PREFIXED)
 
-    val acl1 = new AclBinding(resource1, new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW))
+    val acl1 = new AclBinding(resource1, new AccessControlEntry(principal.toString, WILDCARD_HOST, READ, ALLOW))
     val acl2 = new AclBinding(resource1, new AccessControlEntry(principal.toString, "192.168.0.1", WRITE, ALLOW))
-    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WildcardHost, DESCRIBE, ALLOW))
-    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WildcardHost, READ, ALLOW))
+    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WILDCARD_HOST, DESCRIBE, ALLOW))
+    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WILDCARD_HOST, READ, ALLOW))
 
     authorizer1.createAcls(requestContext, List(acl1, acl2, acl3, acl4).asJava)
     assertEquals(Set(acl1, acl2, acl3, acl4), authorizer1.acls(AclBindingFilter.ANY).asScala.toSet)
@@ -993,7 +993,7 @@ class AuthorizerTest extends QuorumTestHarness with BaseAuthorizerTest {
     val literalResource = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
     val prefixedResource = new ResourcePattern(TOPIC, "bar-", PREFIXED)
     val wildcardResource = new ResourcePattern(TOPIC, "*", LITERAL)
-    val ace = new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW)
+    val ace = new AccessControlEntry(principal.toString, WILDCARD_HOST, READ, ALLOW)
     val updateSemaphore = new Semaphore(1)
 
     def createAcl(createAuthorizer: Authorizer, resource: ResourcePattern): AclBinding = {
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/BaseAuthorizerTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/BaseAuthorizerTest.scala
index c502b48cec4f3..4dcb6c04d1313 100644
--- a/core/src/test/scala/unit/kafka/security/authorizer/BaseAuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/authorizer/BaseAuthorizerTest.scala
@@ -19,8 +19,6 @@ package kafka.security.authorizer
 
 import java.net.InetAddress
 import java.util.UUID
-
-import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
 import kafka.server.KafkaConfig
 import kafka.zookeeper.ZooKeeperClient
 import org.apache.kafka.common.acl.AclOperation.{ALL, READ, WRITE}
@@ -34,6 +32,8 @@ import org.apache.kafka.common.resource.ResourcePattern.WILDCARD_RESOURCE
 import org.apache.kafka.common.resource.ResourceType.{CLUSTER, GROUP, TOPIC, TRANSACTIONAL_ID}
 import org.apache.kafka.common.resource.{ResourcePattern, ResourceType}
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.security.authorizer.AclEntry
+import org.apache.kafka.security.authorizer.AclEntry.{WILDCARD_HOST, WILDCARD_PRINCIPAL_STRING}
 import org.apache.kafka.server.authorizer.{AuthorizationResult, Authorizer}
 import org.junit.jupiter.api.Assertions.{assertFalse, assertTrue}
 import org.junit.jupiter.api.Test
@@ -234,7 +234,7 @@ trait BaseAuthorizerTest {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user1")
     val host1 = InetAddress.getByName("192.168.1.1")
     val host2 = InetAddress.getByName("192.168.1.2")
-    val allHost = AclEntry.WildcardHost
+    val allHost = AclEntry.WILDCARD_HOST;
     val resource1 = new ResourcePattern(TOPIC, "sb1" + UUID.randomUUID(), LITERAL)
     val resource2 = new ResourcePattern(TOPIC, "sb2" + UUID.randomUUID(), LITERAL)
     val allowHost1 = new AccessControlEntry(user1.toString, host1.getHostAddress, READ, ALLOW)
@@ -274,7 +274,7 @@ trait BaseAuthorizerTest {
   def testAuthorizeByResourceTypeWithAllPrincipalAce(): Unit = {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user1")
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user2")
-    val allUser = AclEntry.WildcardPrincipalString
+    val allUser = AclEntry.WILDCARD_PRINCIPAL_STRING
     val host1 = InetAddress.getByName("192.168.1.1")
     val resource1 = new ResourcePattern(TOPIC, "sb1" + UUID.randomUUID(), LITERAL)
     val resource2 = new ResourcePattern(TOPIC, "sb2" + UUID.randomUUID(), LITERAL)
@@ -313,7 +313,7 @@ trait BaseAuthorizerTest {
 
   @Test
   def testAuthorzeByResourceTypeSuperUserHasAccess(): Unit = {
-    val denyAllAce = new AccessControlEntry(WildcardPrincipalString, WildcardHost, AclOperation.ALL, DENY)
+    val denyAllAce = new AccessControlEntry(WILDCARD_PRINCIPAL_STRING, WILDCARD_HOST, AclOperation.ALL, DENY)
     val superUser1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, superUserName)
     val host1 = InetAddress.getByName("192.0.4.4")
     val allTopicsResource = new ResourcePattern(TOPIC, WILDCARD_RESOURCE, LITERAL)
diff --git a/core/src/test/scala/unit/kafka/security/token/delegation/DelegationTokenManagerTest.scala b/core/src/test/scala/unit/kafka/security/token/delegation/DelegationTokenManagerTest.scala
index 47f6783f2b9dd..6879fa31c0b38 100644
--- a/core/src/test/scala/unit/kafka/security/token/delegation/DelegationTokenManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/token/delegation/DelegationTokenManagerTest.scala
@@ -21,7 +21,6 @@ import java.net.InetAddress
 import java.nio.ByteBuffer
 import java.util.{Base64, Properties}
 import kafka.security.authorizer.AclAuthorizer
-import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server.{CreateTokenResult, DelegationTokenManager, DelegationTokenManagerZk, KafkaConfig, QuorumTestHarness}
 import kafka.utils.TestUtils
 import kafka.zk.KafkaZkClient
@@ -38,6 +37,7 @@ import org.apache.kafka.common.security.token.delegation.internals.DelegationTok
 import org.apache.kafka.common.security.token.delegation.{DelegationToken, TokenInformation}
 import org.apache.kafka.common.utils.{MockTime, SecurityUtils, Time}
 import org.apache.kafka.network.Session
+import org.apache.kafka.security.authorizer.AclEntry.WILDCARD_HOST
 import org.apache.kafka.security.authorizer.AuthorizerUtils
 import org.apache.kafka.server.authorizer._
 import org.apache.kafka.server.config.Defaults
@@ -304,7 +304,7 @@ class DelegationTokenManagerTest extends QuorumTestHarness  {
 
     //get all tokens for multiple owners (owner1, renewer4) and with permission
     createAcl(new AclBinding(new ResourcePattern(DELEGATION_TOKEN, tokenId3, LITERAL),
-      new AccessControlEntry(owner1.toString, WildcardHost, DESCRIBE, ALLOW)))
+      new AccessControlEntry(owner1.toString, WILDCARD_HOST, DESCRIBE, ALLOW)))
     tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(owner1, renewer4))
     assertEquals(3, tokens.size)
 
@@ -319,7 +319,7 @@ class DelegationTokenManagerTest extends QuorumTestHarness  {
     //get all tokens for multiple owners (renewer2, renewer3) which are token renewers principals and with permissions
     hostSession = new Session(renewer2, InetAddress.getByName("192.168.1.1"))
     createAcl(new AclBinding(new ResourcePattern(DELEGATION_TOKEN, tokenId2, LITERAL),
-      new AccessControlEntry(renewer2.toString, WildcardHost, DESCRIBE, ALLOW)))
+      new AccessControlEntry(renewer2.toString, WILDCARD_HOST, DESCRIBE, ALLOW)))
     tokens = getTokens(tokenManager, aclAuthorizer, hostSession,  renewer2, List(renewer2, renewer3))
     assertEquals(2, tokens.size)
 
diff --git a/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala
index acb0a215b19e5..04d03c7b98b06 100644
--- a/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala
@@ -17,21 +17,21 @@
 
 package kafka.server
 
-import java.lang.{Byte => JByte}
-import java.util.Properties
 import kafka.network.SocketServer
-import kafka.security.authorizer.AclEntry
 import kafka.utils.TestInfoUtils
 import org.apache.kafka.common.message.{DescribeClusterRequestData, DescribeClusterResponseData}
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.requests.{DescribeClusterRequest, DescribeClusterResponse}
 import org.apache.kafka.common.resource.ResourceType
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.security.authorizer.AclEntry
 import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue}
 import org.junit.jupiter.api.{BeforeEach, TestInfo}
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
 
+import java.lang.{Byte => JByte}
+import java.util.Properties
 import scala.jdk.CollectionConverters._
 
 class DescribeClusterRequestTest extends BaseRequestTest {
@@ -77,7 +77,7 @@ class DescribeClusterRequestTest extends BaseRequestTest {
 
     val expectedClusterAuthorizedOperations = if (includeClusterAuthorizedOperations) {
       Utils.to32BitField(
-        AclEntry.supportedOperations(ResourceType.CLUSTER)
+        AclEntry.supportedOperations(ResourceType.CLUSTER).asScala
           .map(_.code.asInstanceOf[JByte]).asJava)
     } else {
       Int.MinValue
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index a08dc797af491..09be6a918d6fd 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -67,7 +67,7 @@ import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, Enve
 import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, KafkaPrincipalSerde, SecurityProtocol}
 import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, Deserializer, IntegerSerializer, Serializer}
-import org.apache.kafka.common.utils.Utils._
+import org.apache.kafka.common.utils.Utils.formatAddress
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.controller.QuorumController
 import org.apache.kafka.metadata.properties.MetaProperties
diff --git a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
index 27720860ee411..9e770d0141b82 100644
--- a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
@@ -22,11 +22,11 @@ import java.util.{Collections, Properties}
 import kafka.api.LeaderAndIsr
 import kafka.cluster.{Broker, EndPoint}
 import kafka.controller.{LeaderIsrAndControllerEpoch, ReplicaAssignment}
-import kafka.security.authorizer.AclEntry
 import kafka.server.{KafkaConfig, QuorumTestHarness}
 import kafka.utils.CoreUtils
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
 import kafka.zookeeper._
+import org.apache.kafka.common.acl.AccessControlEntry
 import org.apache.kafka.common.acl.AclOperation.READ
 import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
 import org.apache.kafka.common.config.TopicConfig
@@ -43,6 +43,7 @@ import org.apache.kafka.common.utils.{SecurityUtils, Time}
 import org.apache.kafka.common.{TopicPartition, Uuid}
 import org.apache.kafka.metadata.LeaderRecoveryState
 import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
+import org.apache.kafka.security.authorizer.AclEntry
 import org.apache.kafka.server.common.MetadataVersion
 import org.apache.kafka.server.config.{ConfigType, ZkConfigs}
 import org.apache.kafka.storage.internals.log.LogConfig
@@ -599,7 +600,7 @@ class KafkaZkClientTest extends QuorumTestHarness {
     ZkAclStore.stores.foreach(store => {
       assertFalse(zkClient.pathExists(store.aclPath))
       assertFalse(zkClient.pathExists(store.changeStore.aclChangePath))
-      AclEntry.ResourceTypes.foreach(resource => assertFalse(zkClient.pathExists(store.path(resource))))
+      AclEntry.RESOURCE_TYPES.forEach(resource => assertFalse(zkClient.pathExists(store.path(resource))))
     })
 
     // create acl paths
@@ -608,7 +609,7 @@ class KafkaZkClientTest extends QuorumTestHarness {
     ZkAclStore.stores.foreach(store => {
       assertTrue(zkClient.pathExists(store.aclPath))
       assertTrue(zkClient.pathExists(store.changeStore.aclChangePath))
-      AclEntry.ResourceTypes.foreach(resource => assertTrue(zkClient.pathExists(store.path(resource))))
+      AclEntry.RESOURCE_TYPES.forEach(resource => assertTrue(zkClient.pathExists(store.path(resource))))
 
       val resource1 = new ResourcePattern(TOPIC, Uuid.randomUuid().toString, store.patternType)
       val resource2 = new ResourcePattern(TOPIC, Uuid.randomUuid().toString, store.patternType)
@@ -620,9 +621,9 @@ class KafkaZkClientTest extends QuorumTestHarness {
       assertFalse(zkClient.resourceExists(resource1))
 
 
-      val acl1 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), DENY, "host1" , READ)
-      val acl2 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), ALLOW, "*", READ)
-      val acl3 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), DENY, "host1", READ)
+      val acl1 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice").toString, "host1" , READ, DENY))
+      val acl2 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob").toString, "*", READ, ALLOW))
+      val acl3 = new AclEntry(new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob").toString, "host1", READ, DENY))
 
       // Conditional set should fail if path not created
       assertFalse(zkClient.conditionalSetAclsForResource(resource1, Set(acl1, acl3), 0)._1)
@@ -647,7 +648,7 @@ class KafkaZkClientTest extends QuorumTestHarness {
       assertEquals(1, versionedAcls.zkVersion)
 
       //get resource Types
-      assertEquals(AclEntry.ResourceTypes.map(SecurityUtils.resourceTypeName), zkClient.getResourceTypes(store.patternType).toSet)
+      assertEquals(AclEntry.RESOURCE_TYPES.asScala.map(SecurityUtils.resourceTypeName), zkClient.getResourceTypes(store.patternType).toSet)
 
       //get resource name
       val resourceNames = zkClient.getResourceNames(store.patternType, TOPIC)
diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
index a4045f5794871..b8360f64e7ee9 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkAclMigrationClientTest.scala
@@ -16,8 +16,7 @@
  */
 package kafka.zk.migration
 
-import kafka.security.authorizer.{AclAuthorizer, AclEntry}
-import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
+import kafka.security.authorizer.AclAuthorizer
 import kafka.utils.TestUtils
 import org.apache.kafka.common.Uuid
 import org.apache.kafka.common.acl._
@@ -27,6 +26,8 @@ import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.SecurityUtils
 import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
 import org.apache.kafka.metadata.migration.KRaftMigrationZkWriter
+import org.apache.kafka.security.authorizer.AclEntry
+import org.apache.kafka.security.authorizer.AclEntry.{WILDCARD_HOST, WILDCARD_PRINCIPAL_STRING}
 import org.apache.kafka.server.common.ApiMessageAndVersion
 import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue, fail}
 import org.junit.jupiter.api.Test
@@ -80,14 +81,14 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
     val prefixedResource = new ResourcePattern(ResourceType.TOPIC, "bar-" + Uuid.randomUuid(), PatternType.PREFIXED)
     val username = "alice"
     val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
-    val wildcardPrincipal = SecurityUtils.parseKafkaPrincipal(WildcardPrincipalString)
+    val wildcardPrincipal = SecurityUtils.parseKafkaPrincipal(WILDCARD_PRINCIPAL_STRING)
 
-    val ace1 = new AccessControlEntry(principal.toString, WildcardHost, AclOperation.READ, AclPermissionType.ALLOW)
+    val ace1 = new AccessControlEntry(principal.toString, WILDCARD_HOST, AclOperation.READ, AclPermissionType.ALLOW)
     val acl1 = new AclBinding(resource1, ace1)
     val ace2 = new AccessControlEntry(principal.toString, "192.168.0.1", AclOperation.WRITE, AclPermissionType.ALLOW)
     val acl2 = new AclBinding(resource1, ace2)
-    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WildcardHost, AclOperation.DESCRIBE, AclPermissionType.ALLOW))
-    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WildcardHost, AclOperation.READ, AclPermissionType.ALLOW))
+    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WILDCARD_HOST, AclOperation.DESCRIBE, AclPermissionType.ALLOW))
+    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WILDCARD_HOST, AclOperation.READ, AclPermissionType.ALLOW))
 
     val authorizer = new AclAuthorizer()
     try {
@@ -128,8 +129,8 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
     val username2 = "blah"
     val principal1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username1)
     val principal2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username2)
-    val acl1Resource1 = new AclEntry(new AccessControlEntry(principal1.toString, WildcardHost, AclOperation.WRITE, AclPermissionType.ALLOW))
-    val acl1Resource2 = new AclEntry(new AccessControlEntry(principal2.toString, WildcardHost, AclOperation.READ, AclPermissionType.ALLOW))
+    val acl1Resource1 = new AclEntry(new AccessControlEntry(principal1.toString, WILDCARD_HOST, AclOperation.WRITE, AclPermissionType.ALLOW))
+    val acl1Resource2 = new AclEntry(new AccessControlEntry(principal2.toString, WILDCARD_HOST, AclOperation.READ, AclPermissionType.ALLOW))
 
     zkClient.createAclPaths()
     zkClient.createAclsForResourceIfNotExists(resource1, Set(acl1Resource1))
@@ -144,7 +145,7 @@ class ZkAclMigrationClientTest extends ZkMigrationTestHarness {
       .setId(Uuid.randomUuid())
       .setHost("192.168.10.1")
       .setOperation(AclOperation.READ.code())
-      .setPrincipal(WildcardPrincipalString)
+      .setPrincipal(WILDCARD_PRINCIPAL_STRING)
       .setPermissionType(AclPermissionType.ALLOW.code())
       .setPatternType(resource3.patternType().code())
       .setResourceName(resource3.name())
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/acl/AuthorizerBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/acl/AuthorizerBenchmark.java
index 01f93924127bd..6b7d4285876cf 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/acl/AuthorizerBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/acl/AuthorizerBenchmark.java
@@ -19,7 +19,6 @@
 
 import kafka.security.authorizer.AclAuthorizer;
 import kafka.security.authorizer.AclAuthorizer.VersionedAcls;
-import kafka.security.authorizer.AclEntry;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.apache.kafka.common.acl.AclBinding;
@@ -36,6 +35,7 @@
 import org.apache.kafka.common.resource.ResourceType;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.security.authorizer.AclEntry;
 import org.apache.kafka.metadata.authorizer.StandardAcl;
 import org.apache.kafka.metadata.authorizer.StandardAuthorizer;
 import org.apache.kafka.server.authorizer.Action;
diff --git a/server/src/main/java/org/apache/kafka/security/CredentialProvider.java b/server/src/main/java/org/apache/kafka/security/CredentialProvider.java
new file mode 100644
index 0000000000000..a8c4d378a1b83
--- /dev/null
+++ b/server/src/main/java/org/apache/kafka/security/CredentialProvider.java
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.security;
+
+import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+
+import java.util.Collection;
+import java.util.Properties;
+
+public class CredentialProvider {
+    public final DelegationTokenCache tokenCache;
+    public final CredentialCache credentialCache = new CredentialCache();
+
+    public CredentialProvider(Collection<String> scramMechanisms, DelegationTokenCache tokenCache) {
+        this.tokenCache = tokenCache;
+        ScramCredentialUtils.createCache(credentialCache, scramMechanisms);
+    }
+
+    public void updateCredentials(String username, Properties config) {
+        for (ScramMechanism mechanism : ScramMechanism.values()) {
+            CredentialCache.Cache<ScramCredential> cache = credentialCache.cache(mechanism.mechanismName(), ScramCredential.class);
+            if (cache != null) {
+                String c = config.getProperty(mechanism.mechanismName());
+                if (c == null) {
+                    cache.remove(username);
+                } else {
+                    cache.put(username, ScramCredentialUtils.credentialFromString(c));
+                }
+            }
+        }
+    }
+
+    public void updateCredential(
+        org.apache.kafka.clients.admin.ScramMechanism mechanism,
+        String name,
+        ScramCredential credential
+    ) {
+        CredentialCache.Cache<ScramCredential> cache = credentialCache.cache(mechanism.mechanismName(), ScramCredential.class);
+        cache.put(name, credential);
+    }
+
+    public void removeCredentials(
+        org.apache.kafka.clients.admin.ScramMechanism mechanism,
+        String name
+    ) {
+        CredentialCache.Cache<ScramCredential> cache = credentialCache.cache(mechanism.mechanismName(), ScramCredential.class);
+        if (cache != null) {
+            cache.remove(name);
+        }
+    }
+}
diff --git a/server/src/main/java/org/apache/kafka/security/authorizer/AclEntry.java b/server/src/main/java/org/apache/kafka/security/authorizer/AclEntry.java
new file mode 100644
index 0000000000000..fbb4c65d99981
--- /dev/null
+++ b/server/src/main/java/org/apache/kafka/security/authorizer/AclEntry.java
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.security.authorizer;
+
+import org.apache.kafka.common.acl.AccessControlEntry;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.utils.SecurityUtils;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.server.util.Json;
+import org.apache.kafka.server.util.json.DecodeJson;
+import org.apache.kafka.server.util.json.JsonObject;
+import org.apache.kafka.server.util.json.JsonValue;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.common.acl.AclOperation.ALTER;
+import static org.apache.kafka.common.acl.AclOperation.ALTER_CONFIGS;
+import static org.apache.kafka.common.acl.AclOperation.CLUSTER_ACTION;
+import static org.apache.kafka.common.acl.AclOperation.CREATE;
+import static org.apache.kafka.common.acl.AclOperation.CREATE_TOKENS;
+import static org.apache.kafka.common.acl.AclOperation.DELETE;
+import static org.apache.kafka.common.acl.AclOperation.DESCRIBE;
+import static org.apache.kafka.common.acl.AclOperation.DESCRIBE_CONFIGS;
+import static org.apache.kafka.common.acl.AclOperation.DESCRIBE_TOKENS;
+import static org.apache.kafka.common.acl.AclOperation.IDEMPOTENT_WRITE;
+import static org.apache.kafka.common.acl.AclOperation.READ;
+import static org.apache.kafka.common.acl.AclOperation.WRITE;
+
+public class AclEntry extends AccessControlEntry {
+    private static final DecodeJson.DecodeInteger INT = new DecodeJson.DecodeInteger();
+    private static final DecodeJson.DecodeString STRING = new DecodeJson.DecodeString();
+
+    public static final KafkaPrincipal WILDCARD_PRINCIPAL = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*");
+    public static final String WILDCARD_PRINCIPAL_STRING = WILDCARD_PRINCIPAL.toString();
+    public static final String WILDCARD_HOST = "*";
+    public static final String WILDCARD_RESOURCE = ResourcePattern.WILDCARD_RESOURCE;
+    public static final String RESOURCE_SEPARATOR = ":";
+    public static final Set<ResourceType> RESOURCE_TYPES = Arrays.stream(ResourceType.values())
+        .filter(t -> !(t == ResourceType.UNKNOWN || t == ResourceType.ANY))
+        .collect(Collectors.toSet());
+    public static final Set<AclOperation> ACL_OPERATIONS = Arrays.stream(AclOperation.values())
+        .filter(t -> !(t == AclOperation.UNKNOWN || t == AclOperation.ANY))
+        .collect(Collectors.toSet());
+
+    private static final String PRINCIPAL_KEY = "principal";
+    private static final String PERMISSION_TYPE_KEY = "permissionType";
+    private static final String OPERATION_KEY = "operation";
+    private static final String HOSTS_KEY = "host";
+    public static final String VERSION_KEY = "version";
+    public static final int CURRENT_VERSION = 1;
+    private static final String ACLS_KEY = "acls";
+
+    public final AccessControlEntry ace;
+    public final KafkaPrincipal kafkaPrincipal;
+
+    public AclEntry(AccessControlEntry ace) {
+        super(ace.principal(), ace.host(), ace.operation(), ace.permissionType());
+        this.ace = ace;
+
+        kafkaPrincipal = ace.principal() == null
+            ? null
+            : SecurityUtils.parseKafkaPrincipal(ace.principal());
+    }
+
+    /**
+     * Parse JSON representation of ACLs
+     * @param bytes of acls json string
+     *
+     * <p>
+        {
+            "version": 1,
+            "acls": [
+                {
+                    "host":"host1",
+                    "permissionType": "Deny",
+                    "operation": "Read",
+                    "principal": "User:alice"
+                }
+            ]
+        }
+     * </p>
+     *
+     * @return set of AclEntry objects from the JSON string
+     */
+    public static Set<AclEntry> fromBytes(byte[] bytes) throws IOException {
+        if (bytes == null || bytes.length == 0)
+            return Collections.emptySet();
+
+        Optional<JsonValue> jsonValue = Json.parseBytes(bytes);
+        if (!jsonValue.isPresent())
+            return Collections.emptySet();
+
+        JsonObject js = jsonValue.get().asJsonObject();
+
+        //the acl json version.
+        Utils.require(js.apply(VERSION_KEY).to(INT) == CURRENT_VERSION);
+
+        Set<AclEntry> res = new HashSet<>();
+
+        Iterator<JsonValue> aclsIter = js.apply(ACLS_KEY).asJsonArray().iterator();
+        while (aclsIter.hasNext()) {
+            JsonObject itemJs = aclsIter.next().asJsonObject();
+            KafkaPrincipal principal = SecurityUtils.parseKafkaPrincipal(itemJs.apply(PRINCIPAL_KEY).to(STRING));
+            AclPermissionType permissionType = SecurityUtils.permissionType(itemJs.apply(PERMISSION_TYPE_KEY).to(STRING));
+            String host = itemJs.apply(HOSTS_KEY).to(STRING);
+            AclOperation operation = SecurityUtils.operation(itemJs.apply(OPERATION_KEY).to(STRING));
+
+            res.add(new AclEntry(new AccessControlEntry(principal.toString(),
+                host, operation, permissionType)));
+        }
+
+        return res;
+    }
+
+    public static Map<String, Object> toJsonCompatibleMap(Set<AclEntry> acls) {
+        Map<String, Object> res = new HashMap<>();
+        res.put(AclEntry.VERSION_KEY, AclEntry.CURRENT_VERSION);
+        res.put(AclEntry.ACLS_KEY, acls.stream().map(AclEntry::toMap).collect(Collectors.toList()));
+        return res;
+    }
+
+    public static Set<AclOperation> supportedOperations(ResourceType resourceType) {
+        switch (resourceType) {
+            case TOPIC:
+                return new HashSet<>(Arrays.asList(READ, WRITE, CREATE, DESCRIBE, DELETE, ALTER, DESCRIBE_CONFIGS, ALTER_CONFIGS));
+            case GROUP:
+                return new HashSet<>(Arrays.asList(READ, DESCRIBE, DELETE));
+            case CLUSTER:
+                return new HashSet<>(Arrays.asList(CREATE, CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE, ALTER, DESCRIBE));
+            case TRANSACTIONAL_ID:
+                return new HashSet<>(Arrays.asList(DESCRIBE, WRITE));
+            case DELEGATION_TOKEN:
+                return Collections.singleton(DESCRIBE);
+            case USER:
+                return new HashSet<>(Arrays.asList(CREATE_TOKENS, DESCRIBE_TOKENS));
+            default:
+                throw new IllegalArgumentException("Not a concrete resource type");
+        }
+    }
+
+    public static Errors authorizationError(ResourceType resourceType) {
+        switch (resourceType) {
+            case TOPIC:
+                return Errors.TOPIC_AUTHORIZATION_FAILED;
+            case GROUP:
+                return Errors.GROUP_AUTHORIZATION_FAILED;
+            case CLUSTER:
+                return Errors.CLUSTER_AUTHORIZATION_FAILED;
+            case TRANSACTIONAL_ID:
+                return Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED;
+            case DELEGATION_TOKEN:
+                return Errors.DELEGATION_TOKEN_AUTHORIZATION_FAILED;
+            default:
+                throw new IllegalArgumentException("Authorization error type not known");
+        }
+    }
+
+    public Map<String, Object> toMap() {
+        Map<String, Object> res = new HashMap<>();
+        res.put(AclEntry.PRINCIPAL_KEY, principal());
+        res.put(AclEntry.PERMISSION_TYPE_KEY, SecurityUtils.permissionTypeName(permissionType()));
+        res.put(AclEntry.OPERATION_KEY, SecurityUtils.operationName(operation()));
+        res.put(AclEntry.HOSTS_KEY, host());
+        return res;
+    }
+
+    @Override
+    public int hashCode() {
+        return ace.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return super.equals(o); // to keep spotbugs happy
+    }
+
+    @Override
+    public String toString() {
+        return String.format("%s has %s permission for operations: %s from hosts: %s", principal(), permissionType().name(), operation(), host());
+    }
+}
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
index a0ad01bad955d..32a00324df6d4 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/AuthorizerIntegrationTest.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.tools.consumer.group;
 
 import kafka.api.AbstractAuthorizerIntegrationTest;
-import kafka.security.authorizer.AclEntry;
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -34,7 +33,7 @@ public class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest
     @ParameterizedTest(name = TEST_WITH_PARAMETERIZED_QUORUM_NAME)
     @ValueSource(strings = {"zk", "kraft"})
     public void testDescribeGroupCliWithGroupDescribe(String quorum) throws Exception {
-        addAndVerifyAcls(JavaConverters.asScalaSet(Collections.singleton(new AccessControlEntry(ClientPrincipal().toString(), AclEntry.WildcardHost(), DESCRIBE, ALLOW))).toSet(), groupResource());
+        addAndVerifyAcls(JavaConverters.asScalaSet(Collections.singleton(new AccessControlEntry(ClientPrincipal().toString(), "*", DESCRIBE, ALLOW))).toSet(), groupResource());
 
         String[] cgcArgs = new String[]{"--bootstrap-server", bootstrapServers(listenerName()), "--describe", "--group", group()};
         ConsumerGroupCommandOptions opts = ConsumerGroupCommandOptions.fromArgs(cgcArgs);

From 9a9b532d5d5beeecf7a4b769731ee609625429e1 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Sat, 30 Mar 2024 19:18:41 +0800
Subject: [PATCH 233/258] KAFKA-16447: Fix failed ReplicaManagerTest (#15630)

The change in https://github.com/apache/kafka/pull/15373/files#r1544335647 updated broker port from 9093 to 9094. Some of test cases check broker endpoint with fixed string 9093. I change test cases to get endpoint by broker id, so these cases will not fail if someone change the port again in the future.

Reviewers: Luke Chen <showuon@gmail.com>
---
 .../kafka/server/ReplicaManagerTest.scala     | 28 ++++++++++++-------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index ed64d986c269e..35a5d9cc7fb73 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -5161,7 +5161,8 @@ class ReplicaManagerTest {
       }
 
       val fetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), fetcher.map(_.leader.brokerEndPoint()))
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), fetcher.map(_.leader.brokerEndPoint()))
     } finally {
       replicaManager.shutdown(checkpointHW = false)
     }
@@ -5193,7 +5194,8 @@ class ReplicaManagerTest {
       }
 
       val fetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), fetcher.map(_.leader.brokerEndPoint()))
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), fetcher.map(_.leader.brokerEndPoint()))
 
       // Append on a follower should fail
       val followerResponse = sendProducerAppend(replicaManager, topicPartition, numOfRecords)
@@ -5255,7 +5257,8 @@ class ReplicaManagerTest {
       }
 
       val fetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), fetcher.map(_.leader.brokerEndPoint()))
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), fetcher.map(_.leader.brokerEndPoint()))
 
       // Apply the same delta again
       replicaManager.applyDelta(followerTopicsDelta, followerMetadataImage)
@@ -5270,7 +5273,7 @@ class ReplicaManagerTest {
       }
 
       val noChangeFetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), noChangeFetcher.map(_.leader.brokerEndPoint()))
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), noChangeFetcher.map(_.leader.brokerEndPoint()))
     } finally {
       replicaManager.shutdown(checkpointHW = false)
     }
@@ -5301,7 +5304,8 @@ class ReplicaManagerTest {
       }
 
       val fetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), fetcher.map(_.leader.brokerEndPoint()))
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), fetcher.map(_.leader.brokerEndPoint()))
 
       // Apply changes that remove replica
       val notReplicaTopicsDelta = topicsChangeDelta(followerMetadataImage.topics(), otherId, true)
@@ -5348,7 +5352,8 @@ class ReplicaManagerTest {
       }
 
       val fetcher = replicaManager.replicaFetcherManager.getFetcher(topicPartition)
-      assertEquals(Some(BrokerEndPoint(otherId, "localhost", 9093)), fetcher.map(_.leader.brokerEndPoint()))
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
+      assertEquals(Some(BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port())), fetcher.map(_.leader.brokerEndPoint()))
 
       // Apply changes that remove topic and replica
       val removeTopicsDelta = topicsDeleteDelta(followerMetadataImage.topics())
@@ -5652,11 +5657,12 @@ class ReplicaManagerTest {
 
       // Verify that addFetcherForPartitions was called with the correct
       // init offset.
+      val otherEndpoint = ClusterImageTest.IMAGE1.broker(otherId).listeners().get("PLAINTEXT")
       verify(mockReplicaFetcherManager)
         .addFetcherForPartitions(
           Map(topicPartition -> InitialFetchState(
             topicId = Some(FOO_UUID),
-            leader = BrokerEndPoint(otherId, "localhost", 9093),
+            leader = BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port()),
             currentLeaderEpoch = 0,
             initOffset = 0
           ))
@@ -5700,7 +5706,7 @@ class ReplicaManagerTest {
         .addFetcherForPartitions(
           Map(topicPartition -> InitialFetchState(
             topicId = Some(FOO_UUID),
-            leader = BrokerEndPoint(otherId, "localhost", 9093),
+            leader = BrokerEndPoint(otherId, otherEndpoint.host(), otherEndpoint.port()),
             currentLeaderEpoch = 1,
             initOffset = 1
           ))
@@ -5857,10 +5863,11 @@ class ReplicaManagerTest {
       }
 
       // Verify that the partition was removed and added back.
+      val localIdPlus1Endpoint = ClusterImageTest.IMAGE1.broker(localId + 1).listeners().get("PLAINTEXT")
       verify(mockReplicaFetcherManager).removeFetcherForPartitions(Set(topicPartition))
       verify(mockReplicaFetcherManager).addFetcherForPartitions(Map(topicPartition -> InitialFetchState(
         topicId = Some(FOO_UUID),
-        leader = BrokerEndPoint(localId + 1, "localhost", 9093),
+        leader = BrokerEndPoint(localId + 1, localIdPlus1Endpoint.host(), localIdPlus1Endpoint.port()),
         currentLeaderEpoch = 0,
         initOffset = 0
       )))
@@ -5916,10 +5923,11 @@ class ReplicaManagerTest {
       }
 
       // Verify that the partition was removed and added back.
+      val localIdPlus2Endpoint = ClusterImageTest.IMAGE1.broker(localId + 2).listeners().get("PLAINTEXT")
       verify(mockReplicaFetcherManager).removeFetcherForPartitions(Set(topicPartition))
       verify(mockReplicaFetcherManager).addFetcherForPartitions(Map(topicPartition -> InitialFetchState(
         topicId = Some(FOO_UUID),
-        leader = BrokerEndPoint(localId + 2, "localhost", 9093),
+        leader = BrokerEndPoint(localId + 2, localIdPlus2Endpoint.host(), localIdPlus2Endpoint.port()),
         currentLeaderEpoch = 1,
         initOffset = 0
       )))

From d4caa1c10ec81b9c87eaaf52b73c83d5579b68d3 Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Sun, 31 Mar 2024 11:22:16 +0800
Subject: [PATCH 234/258] MINOR: Remove redundant
 ApiVersionsResponse#filterApis (#15611)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../apache/kafka/common/requests/ApiVersionsResponse.java  | 7 -------
 .../apache/kafka/clients/admin/KafkaAdminClientTest.java   | 2 +-
 2 files changed, 1 insertion(+), 8 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsResponse.java
index f8f364f06f794..04ee2014d1d44 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ApiVersionsResponse.java
@@ -166,13 +166,6 @@ public static ApiVersionsResponse createApiVersionsResponse(
         );
     }
 
-    public static ApiVersionCollection filterApis(
-        RecordVersion minRecordVersion,
-        ApiMessageType.ListenerType listenerType
-    ) {
-        return filterApis(minRecordVersion, listenerType, false, false);
-    }
-
     public static ApiVersionCollection filterApis(
         RecordVersion minRecordVersion,
         ApiMessageType.ListenerType listenerType,
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
index 71e802365fd0a..3f9f3707479d7 100644
--- a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java
@@ -712,7 +712,7 @@ private static ApiVersionsResponse prepareApiVersionsResponseForDescribeFeatures
         if (error == Errors.NONE) {
             return ApiVersionsResponse.createApiVersionsResponse(
                 0,
-                ApiVersionsResponse.filterApis(RecordVersion.current(), ApiMessageType.ListenerType.ZK_BROKER),
+                ApiVersionsResponse.filterApis(RecordVersion.current(), ApiMessageType.ListenerType.ZK_BROKER, false, false),
                 convertSupportedFeaturesMap(defaultFeatureMetadata().supportedFeatures()),
                 Collections.singletonMap("test_feature_1", (short) 2),
                 defaultFeatureMetadata().finalizedFeaturesEpoch().get(),

From a640a81040f6ef6f85819b60194f0394f5f2194e Mon Sep 17 00:00:00 2001
From: Johnny Hsu <44309740+johnnychhsu@users.noreply.github.com>
Date: Mon, 1 Apr 2024 10:52:53 +0800
Subject: [PATCH 235/258] KAFKA-16323: fix testRemoteFetchExpiresPerSecMetric
 (#15463)

The variable of metrics item (kafka.server:type=DelayedRemoteFetchMetrics,name=ExpiresPerSec) is singleton object and it could be removed by other tests which are running in same JVM (and it is not recreated). Hence, verifying the metrics value is not stable to this test case.

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>, Kamal Chandraprakash <kamal.chandraprakash@gmail.com>
---
 .../kafka/server/ReplicaManagerTest.scala     | 23 +++----------------
 1 file changed, 3 insertions(+), 20 deletions(-)

diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index 35a5d9cc7fb73..c7587e4d61552 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -4185,16 +4185,13 @@ class ReplicaManagerTest {
         mock(classOf[FetchDataInfo])
       }).when(spyRLM).read(any())
 
-      // Get the current type=DelayedRemoteFetchMetrics,name=ExpiresPerSec metric value before fetching
-      val curExpiresPerSec = safeYammerMetricValue("type=DelayedRemoteFetchMetrics,name=ExpiresPerSec").asInstanceOf[Long]
+      val curExpiresPerSec = DelayedRemoteFetchMetrics.expiredRequestMeter.count()
       replicaManager.fetchMessages(params, Seq(tidp0 -> new PartitionData(topicId, fetchOffset, 0, 100000, Optional.of[Integer](leaderEpoch), Optional.of[Integer](leaderEpoch))), UnboundedQuota, fetchCallback)
       // advancing the clock to expire the delayed remote fetch
       timer.advanceClock(2000L)
 
-      // verify the metric value is incremented since the delayed remote fetch is expired
-      TestUtils.waitUntilTrue(() => curExpiresPerSec + 1 == safeYammerMetricValue("type=DelayedRemoteFetchMetrics,name=ExpiresPerSec").asInstanceOf[Long],
-        "The ExpiresPerSec value is not incremented. Current value is: " +
-          safeYammerMetricValue("type=DelayedRemoteFetchMetrics,name=ExpiresPerSec").asInstanceOf[Long])
+      // verify the DelayedRemoteFetchMetrics.expiredRequestMeter.mark is called since the delayed remote fetch is expired
+      TestUtils.waitUntilTrue(() => (curExpiresPerSec + 1) == DelayedRemoteFetchMetrics.expiredRequestMeter.count(), "DelayedRemoteFetchMetrics.expiredRequestMeter.count() should be 1, but got: " + DelayedRemoteFetchMetrics.expiredRequestMeter.count(), 10000L)
       latch.countDown()
     } finally {
       Utils.tryAll(util.Arrays.asList[Callable[Void]](
@@ -4210,20 +4207,6 @@ class ReplicaManagerTest {
     }
   }
 
-  private def safeYammerMetricValue(name: String): Any = {
-    val allMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
-    val opt = allMetrics.find { case (n, _) => n.getMBeanName.endsWith(name) }
-    if (opt.isEmpty)
-      0L
-    else {
-      opt.get._2 match {
-        case m: Gauge[_] => m.value
-        case m: Meter => m.count()
-        case m => fail(s"Unexpected broker metric of class ${m.getClass}")
-      }
-    }
-  }
-
   private def yammerMetricValue(name: String): Any = {
     val allMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
     val (_, metric) = allMetrics.find { case (n, _) => n.getMBeanName.endsWith(name) }

From 40e87ae35beb389d6419d32130174d7c68fa4d19 Mon Sep 17 00:00:00 2001
From: Gaurav Narula <gaurav_narula2@apple.com>
Date: Mon, 1 Apr 2024 04:55:08 +0100
Subject: [PATCH 236/258] KAFKA-15823: disconnect from controller on
 AuthenticationException (#14760)

This PR changes the handling of authenticationException on a request from the node to the controller.

We disconnect controller connection and invalidate the cache so that the next run of the thread will establish a connection with the (potentially) updated controller.

Reviewers: Luke Chen <showuon@gmail.com>, Igor Soarez <soarez@apple.com>, Omnia Ibrahim <o.g.h.ibrahim@gmail.com>
---
 .../NodeToControllerChannelManager.scala      | 29 +++++++++++--------
 .../NodeToControllerRequestThreadTest.scala   |  1 +
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/core/src/main/scala/kafka/server/NodeToControllerChannelManager.scala b/core/src/main/scala/kafka/server/NodeToControllerChannelManager.scala
index 19d19c87bb1a4..e0f0118ea383d 100644
--- a/core/src/main/scala/kafka/server/NodeToControllerChannelManager.scala
+++ b/core/src/main/scala/kafka/server/NodeToControllerChannelManager.scala
@@ -345,8 +345,10 @@ class NodeToControllerRequestThread(
   private[server] def handleResponse(queueItem: NodeToControllerQueueItem)(response: ClientResponse): Unit = {
     debug(s"Request ${queueItem.request} received $response")
     if (response.authenticationException != null) {
-      error(s"Request ${queueItem.request} failed due to authentication error with controller",
+      error(s"Request ${queueItem.request} failed due to authentication error with controller. Disconnecting the " +
+        s"connection to the stale controller ${activeControllerAddress().map(_.idString).getOrElse("null")}",
         response.authenticationException)
+      maybeDisconnectAndUpdateController()
       queueItem.callback.onComplete(response)
     } else if (response.versionMismatch != null) {
       error(s"Request ${queueItem.request} failed due to unsupported version error",
@@ -358,23 +360,26 @@ class NodeToControllerRequestThread(
     } else if (response.responseBody().errorCounts().containsKey(Errors.NOT_CONTROLLER)) {
       debug(s"Request ${queueItem.request} received NOT_CONTROLLER exception. Disconnecting the " +
         s"connection to the stale controller ${activeControllerAddress().map(_.idString).getOrElse("null")}")
-      // just close the controller connection and wait for metadata cache update in doWork
-      activeControllerAddress().foreach { controllerAddress =>
-        try {
-          // We don't care if disconnect has an error, just log it and get a new network client
-          networkClient.disconnect(controllerAddress.idString)
-        } catch {
-          case t: Throwable => error("Had an error while disconnecting from NetworkClient.", t)
-        }
-        updateControllerAddress(null)
-      }
-
+      maybeDisconnectAndUpdateController()
       requestQueue.putFirst(queueItem)
     } else {
       queueItem.callback.onComplete(response)
     }
   }
 
+  private def maybeDisconnectAndUpdateController(): Unit = {
+    // just close the controller connection and wait for metadata cache update in doWork
+    activeControllerAddress().foreach { controllerAddress =>
+      try {
+        // We don't care if disconnect has an error, just log it and get a new network client
+        networkClient.disconnect(controllerAddress.idString)
+      } catch {
+        case t: Throwable => error("Had an error while disconnecting from NetworkClient.", t)
+      }
+      updateControllerAddress(null)
+    }
+  }
+
   override def doWork(): Unit = {
     val controllerInformation = controllerNodeProvider.getControllerInfo()
     maybeResetNetworkClient(controllerInformation)
diff --git a/core/src/test/scala/kafka/server/NodeToControllerRequestThreadTest.scala b/core/src/test/scala/kafka/server/NodeToControllerRequestThreadTest.scala
index 5b98fe4d9c853..c15c8bf78a025 100644
--- a/core/src/test/scala/kafka/server/NodeToControllerRequestThreadTest.scala
+++ b/core/src/test/scala/kafka/server/NodeToControllerRequestThreadTest.scala
@@ -427,6 +427,7 @@ class NodeToControllerRequestThreadTest {
     testRequestThread.enqueue(queueItem)
     pollUntil(testRequestThread, () => callbackResponse.get != null)
     assertNotNull(callbackResponse.get.authenticationException)
+    assertEquals(None, testRequestThread.activeControllerAddress())
   }
 
   @Test

From a524b6217b83f1ad327ae5c39ea0345ef6cb90da Mon Sep 17 00:00:00 2001
From: Ori Hoch <ori@uumpa.com>
Date: Mon, 1 Apr 2024 13:55:33 +0300
Subject: [PATCH 237/258] MINOR: Fix doc of Admin#createDelegationToken
 (#15632)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 clients/src/main/java/org/apache/kafka/clients/admin/Admin.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
index e9f0a58e41b98..d936ec80ffe81 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java
@@ -754,7 +754,7 @@ default CreateDelegationTokenResult createDelegationToken() {
      * </ul>
      *
      * @param options The options to use when creating delegation token.
-     * @return The DeleteRecordsResult.
+     * @return The CreateDelegationTokenResult.
      */
     CreateDelegationTokenResult createDelegationToken(CreateDelegationTokenOptions options);
 

From cc6b919212ae62d75850214ae2c93379b78ff325 Mon Sep 17 00:00:00 2001
From: "Kuan-Po (Cooper) Tseng" <brandboat@gmail.com>
Date: Tue, 2 Apr 2024 04:21:10 +0800
Subject: [PATCH 238/258] KAFKA-16435 Add test for KAFKA-16428 (#15635)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../zk/migration/ZkConfigMigrationClientTest.scala  | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
index 971fdecbb8333..755817f368702 100644
--- a/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/migration/ZkConfigMigrationClientTest.scala
@@ -40,9 +40,10 @@ import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
 import org.apache.kafka.server.common.ApiMessageAndVersion
 import org.apache.kafka.server.config.ConfigType
 import org.apache.kafka.server.util.MockRandom
-import org.junit.jupiter.api.Assertions.{assertEquals, assertTrue, fail}
+import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertTrue, fail}
 import org.junit.jupiter.api.Test
 
+import java.nio.charset.StandardCharsets.UTF_8
 import java.util
 import java.util.Properties
 import scala.collection.Map
@@ -112,10 +113,20 @@ class ZkConfigMigrationClientTest extends ZkMigrationTestHarness {
         assertEquals(newProps.get(key), value)
       }
     }
+    assertPathExistenceAndData("/config/changes/config_change_0000000000", """{"version":2,"entity_path":"brokers/1"}""")
 
     migrationState = migrationClient.configClient().deleteConfigs(
       new ConfigResource(ConfigResource.Type.BROKER, "1"), migrationState)
     assertEquals(0, zkClient.getEntityConfigs(ConfigType.BROKER, "1").size())
+    assertPathExistenceAndData("/config/changes/config_change_0000000001", """{"version":2,"entity_path":"brokers/1"}""")
+
+    // make sure there is no more config change notification in znode
+    assertFalse(zkClient.pathExists("/config/changes/config_change_0000000002"))
+  }
+
+  private def assertPathExistenceAndData(expectedPath: String, data: String): Unit = {
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some(data), zkClient.getDataAndStat(expectedPath)._1.map(new String(_, UTF_8)))
   }
 
   @Test

From 2f733ac58386a1669bd021e673fae581172e2b56 Mon Sep 17 00:00:00 2001
From: Kamal Chandraprakash <kchandraprakash@uber.com>
Date: Tue, 2 Apr 2024 07:37:54 +0530
Subject: [PATCH 239/258] KAFKA-16161: Avoid empty remote metadata snapshot
 file in partition dir (#15636)

Avoid empty remote metadata snapshot file in partition dir

Reviewers: Luke Chen <showuon@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>, Satish Duggana <satishd@apache.org>
---
 .../FileBasedRemoteLogMetadataCache.java      | 111 -------
 .../RemoteLogMetadataSnapshotFile.java        | 271 ------------------
 .../RemoteLogSegmentMetadataSnapshot.java     |   6 +-
 .../RemotePartitionMetadataEventHandler.java  |  10 +-
 .../storage/RemotePartitionMetadataStore.java |  40 +--
 .../TopicBasedRemoteLogMetadataManager.java   |   3 +-
 .../metadata/storage/ConsumerTaskTest.java    |   4 -
 .../FileBasedRemoteLogMetadataCacheTest.java  |  91 ------
 .../RemoteLogMetadataSnapshotFileTest.java    |  86 ------
 9 files changed, 14 insertions(+), 608 deletions(-)
 delete mode 100644 storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCache.java
 delete mode 100644 storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFile.java
 delete mode 100644 storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCacheTest.java
 delete mode 100644 storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFileTest.java

diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCache.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCache.java
deleted file mode 100644
index cc992667ce867..0000000000000
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCache.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.server.log.remote.metadata.storage;
-
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.TopicIdPartition;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentId;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.nio.file.Path;
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * This is a wrapper around {@link RemoteLogMetadataCache} providing a file based snapshot of
- * {@link RemoteLogMetadataCache} for the given {@code topicIdPartition}. Snapshot is stored in the given
- * {@code partitionDir}.
- */
-public class FileBasedRemoteLogMetadataCache extends RemoteLogMetadataCache {
-    private static final Logger log = LoggerFactory.getLogger(FileBasedRemoteLogMetadataCache.class);
-    private final RemoteLogMetadataSnapshotFile snapshotFile;
-    private final TopicIdPartition topicIdPartition;
-
-    @SuppressWarnings("this-escape")
-    public FileBasedRemoteLogMetadataCache(TopicIdPartition topicIdPartition,
-                                           Path partitionDir) {
-        if (!partitionDir.toFile().exists() || !partitionDir.toFile().isDirectory()) {
-            throw new KafkaException("Given partition directory:" + partitionDir + " must be an existing directory.");
-        }
-
-        this.topicIdPartition = topicIdPartition;
-        snapshotFile = new RemoteLogMetadataSnapshotFile(partitionDir);
-
-        try {
-            snapshotFile.read().ifPresent(snapshot -> loadRemoteLogSegmentMetadata(snapshot));
-        } catch (IOException e) {
-            throw new KafkaException(e);
-        }
-    }
-
-    protected final void loadRemoteLogSegmentMetadata(RemoteLogMetadataSnapshotFile.Snapshot snapshot) {
-        log.info("Loading snapshot for partition {} is: {}", topicIdPartition, snapshot);
-        for (RemoteLogSegmentMetadataSnapshot metadataSnapshot : snapshot.remoteLogSegmentMetadataSnapshots()) {
-            switch (metadataSnapshot.state()) {
-                case COPY_SEGMENT_STARTED:
-                    addCopyInProgressSegment(createRemoteLogSegmentMetadata(metadataSnapshot));
-                    break;
-                case COPY_SEGMENT_FINISHED:
-                    handleSegmentWithCopySegmentFinishedState(createRemoteLogSegmentMetadata(metadataSnapshot));
-                    break;
-                case DELETE_SEGMENT_STARTED:
-                    handleSegmentWithDeleteSegmentStartedState(createRemoteLogSegmentMetadata(metadataSnapshot));
-                    break;
-                case DELETE_SEGMENT_FINISHED:
-                default:
-                    throw new IllegalArgumentException("Given remoteLogSegmentMetadata has invalid state: " + metadataSnapshot);
-            }
-        }
-    }
-
-    private RemoteLogSegmentMetadata createRemoteLogSegmentMetadata(RemoteLogSegmentMetadataSnapshot snapshot) {
-        return new RemoteLogSegmentMetadata(new RemoteLogSegmentId(topicIdPartition, snapshot.segmentId()), snapshot.startOffset(),
-                                            snapshot.endOffset(), snapshot.maxTimestampMs(), snapshot.brokerId(), snapshot.eventTimestampMs(),
-                                            snapshot.segmentSizeInBytes(), snapshot.customMetadata(), snapshot.state(), snapshot.segmentLeaderEpochs()
-        );
-    }
-
-    /**
-     * Flushes the in-memory state to the snapshot file.
-     *
-     * @param metadataPartition       remote log metadata partition from which the messages have been consumed for the given
-     *                                user topic partition.
-     * @param metadataPartitionOffset remote log metadata partition offset up to which the messages have been consumed.
-     * @throws IOException if any errors occurred while writing the snapshot to the file.
-     */
-    public void flushToFile(int metadataPartition,
-                            Long metadataPartitionOffset) throws IOException {
-        List<RemoteLogSegmentMetadataSnapshot> snapshots = new ArrayList<>(idToSegmentMetadata.size());
-        for (RemoteLogLeaderEpochState state : leaderEpochEntries.values()) {
-            // Add unreferenced segments first, as to maintain the order when these segments are again read from
-            // the snapshot to build RemoteLogMetadataCache.
-            for (RemoteLogSegmentId id : state.unreferencedSegmentIds()) {
-                snapshots.add(RemoteLogSegmentMetadataSnapshot.create(idToSegmentMetadata.get(id)));
-            }
-            
-            // Add referenced segments.
-            for (RemoteLogSegmentId id : state.referencedSegmentIds()) {
-                snapshots.add(RemoteLogSegmentMetadataSnapshot.create(idToSegmentMetadata.get(id)));
-            }
-        }
-
-        snapshotFile.write(new RemoteLogMetadataSnapshotFile.Snapshot(metadataPartition, metadataPartitionOffset, snapshots));
-    }
-}
diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFile.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFile.java
deleted file mode 100644
index db49bb9cb8a8f..0000000000000
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFile.java
+++ /dev/null
@@ -1,271 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.server.log.remote.metadata.storage;
-
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.server.log.remote.metadata.storage.serialization.RemoteLogMetadataSerde;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.nio.channels.Channels;
-import java.nio.channels.FileChannel;
-import java.nio.channels.ReadableByteChannel;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.StandardOpenOption;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.Objects;
-import java.util.Optional;
-
-/**
- * This class represents the remote log data snapshot stored in a file for a specific topic partition. This is used by
- * {@link TopicBasedRemoteLogMetadataManager} to store the remote log metadata received for a specific partition from
- * remote log metadata topic. This will avoid reading the remote log metadata messages from the topic again when a
- * broker restarts.
- */
-public class RemoteLogMetadataSnapshotFile {
-    private static final Logger log = LoggerFactory.getLogger(RemoteLogMetadataSnapshotFile.class);
-
-    public static final String COMMITTED_LOG_METADATA_SNAPSHOT_FILE_NAME = "remote_log_snapshot";
-
-    // File format:
-    // <header>[<entry>...]
-    // header: <version:short><metadata-partition:int><metadata-partition-offset:long><entries-size:int>
-    // entry: <entry-length><entry-bytes>
-
-    // header size: 2 (version) + 4 (partition num) + 8 (offset) + 4 (entries size) = 18
-    private static final int HEADER_SIZE = 18;
-
-    private final File metadataStoreFile;
-    private final RemoteLogMetadataSerde serde = new RemoteLogMetadataSerde();
-
-    /**
-     * Creates a CommittedLogMetadataSnapshotFile instance backed by a file with the name `remote_log_snapshot` in
-     * the given {@code metadataStoreDir}. It creates the file if it does not exist.
-     *
-     * @param metadataStoreDir directory in which the snapshot file to be created.
-     */
-    RemoteLogMetadataSnapshotFile(Path metadataStoreDir) {
-        this.metadataStoreFile = new File(metadataStoreDir.toFile(), COMMITTED_LOG_METADATA_SNAPSHOT_FILE_NAME);
-
-        // Create an empty file if it does not exist.
-        try {
-            final boolean fileExists = Files.exists(metadataStoreFile.toPath());
-            if (!fileExists) {
-                Files.createFile(metadataStoreFile.toPath());
-            }
-            log.info("Remote log metadata snapshot file: [{}], newFileCreated: [{}]", metadataStoreFile, !fileExists);
-        } catch (IOException e) {
-            throw new KafkaException(e);
-        }
-    }
-
-    /**
-     * Writes the given snapshot replacing the earlier snapshot data.
-     *
-     * @param snapshot Snapshot to be stored.
-     * @throws IOException if there4 is any error in writing the given snapshot to the file.
-     */
-    public synchronized void write(Snapshot snapshot) throws IOException {
-        Path newMetadataSnapshotFilePath = new File(metadataStoreFile.getAbsolutePath() + ".tmp").toPath();
-        try (FileChannel fileChannel = FileChannel.open(newMetadataSnapshotFilePath,
-                                                        StandardOpenOption.CREATE, StandardOpenOption.READ, StandardOpenOption.WRITE)) {
-
-            // header: <version:short><metadata-partition:int><metadata-partition-offset:long>
-            ByteBuffer headerBuffer = ByteBuffer.allocate(HEADER_SIZE);
-
-            // Write version
-            headerBuffer.putShort(snapshot.version());
-
-            // Write metadata partition and metadata partition offset
-            headerBuffer.putInt(snapshot.metadataPartition());
-
-            // Write metadata partition offset
-            headerBuffer.putLong(snapshot.metadataPartitionOffset());
-
-            // Write entries size
-            Collection<RemoteLogSegmentMetadataSnapshot> metadataSnapshots = snapshot.remoteLogSegmentMetadataSnapshots();
-            headerBuffer.putInt(metadataSnapshots.size());
-
-            // Write header
-            headerBuffer.flip();
-            fileChannel.write(headerBuffer);
-
-            // Write each entry
-            ByteBuffer lenBuffer = ByteBuffer.allocate(4);
-            for (RemoteLogSegmentMetadataSnapshot metadataSnapshot : metadataSnapshots) {
-                final byte[] serializedBytes = serde.serialize(metadataSnapshot);
-                // entry format: <entry-length><entry-bytes>
-
-                // Write entry length
-                lenBuffer.putInt(serializedBytes.length);
-                lenBuffer.flip();
-                fileChannel.write(lenBuffer);
-                lenBuffer.rewind();
-
-                // Write entry bytes
-                fileChannel.write(ByteBuffer.wrap(serializedBytes));
-            }
-
-            fileChannel.force(true);
-        }
-
-        Utils.atomicMoveWithFallback(newMetadataSnapshotFilePath, metadataStoreFile.toPath());
-    }
-
-    /**
-     * @return the Snapshot if it exists.
-     * @throws IOException if there is any error in reading the stored snapshot.
-     */
-    public synchronized Optional<Snapshot> read() throws IOException {
-
-        // Checking for empty files.
-        if (metadataStoreFile.length() == 0) {
-            return Optional.empty();
-        }
-
-        try (ReadableByteChannel channel = Channels.newChannel(new FileInputStream(metadataStoreFile))) {
-
-            // header: <version:short><metadata-partition:int><metadata-partition-offset:long>
-            // Read header
-            ByteBuffer headerBuffer = ByteBuffer.allocate(HEADER_SIZE);
-            channel.read(headerBuffer);
-            headerBuffer.rewind();
-            short version = headerBuffer.getShort();
-            int metadataPartition = headerBuffer.getInt();
-            long metadataPartitionOffset = headerBuffer.getLong();
-            int metadataSnapshotsSize = headerBuffer.getInt();
-
-            List<RemoteLogSegmentMetadataSnapshot> result = new ArrayList<>(metadataSnapshotsSize);
-            ByteBuffer lenBuffer = ByteBuffer.allocate(4);
-            int lenBufferReadCt;
-            while ((lenBufferReadCt = channel.read(lenBuffer)) > 0) {
-                lenBuffer.rewind();
-
-                if (lenBufferReadCt != lenBuffer.capacity()) {
-                    throw new IOException("Invalid amount of data read for the length of an entry, file may have been corrupted.");
-                }
-
-                // entry format: <entry-length><entry-bytes>
-
-                // Read the length of each entry
-                final int len = lenBuffer.getInt();
-                lenBuffer.rewind();
-
-                // Read the entry
-                ByteBuffer data = ByteBuffer.allocate(len);
-                final int read = channel.read(data);
-                if (read != len) {
-                    throw new IOException("Invalid amount of data read, file may have been corrupted.");
-                }
-
-                // We are always adding RemoteLogSegmentMetadata only as you can see in #write() method.
-                // Did not add a specific serde for RemoteLogSegmentMetadata and reusing RemoteLogMetadataSerde
-                final RemoteLogSegmentMetadataSnapshot remoteLogSegmentMetadata =
-                        (RemoteLogSegmentMetadataSnapshot) serde.deserialize(data.array());
-                result.add(remoteLogSegmentMetadata);
-            }
-
-            if (metadataSnapshotsSize != result.size()) {
-                throw new IOException("Unexpected entries in the snapshot file. Expected size: " + metadataSnapshotsSize
-                                              + ", but found: " + result.size());
-            }
-
-            return Optional.of(new Snapshot(version, metadataPartition, metadataPartitionOffset, result));
-        }
-    }
-
-    /**
-     * This class represents the collection of remote log metadata for a specific topic partition.
-     */
-    public static final class Snapshot {
-        private static final short CURRENT_VERSION = 0;
-
-        private final short version;
-        private final int metadataPartition;
-        private final long metadataPartitionOffset;
-        private final Collection<RemoteLogSegmentMetadataSnapshot> remoteLogSegmentMetadataSnapshots;
-
-        public Snapshot(int metadataPartition,
-                        long metadataPartitionOffset,
-                        Collection<RemoteLogSegmentMetadataSnapshot> remoteLogSegmentMetadataSnapshots) {
-            this(CURRENT_VERSION, metadataPartition, metadataPartitionOffset, remoteLogSegmentMetadataSnapshots);
-        }
-
-        public Snapshot(short version,
-                        int metadataPartition,
-                        long metadataPartitionOffset,
-                        Collection<RemoteLogSegmentMetadataSnapshot> remoteLogSegmentMetadataSnapshots) {
-            // We will add multiple version support in future if needed. For now, the only supported version is CURRENT_VERSION viz 0.
-            if (version != CURRENT_VERSION) {
-                throw new IllegalArgumentException("Unexpected version received: " + version);
-            }
-            this.version = version;
-            this.metadataPartition = metadataPartition;
-            this.metadataPartitionOffset = metadataPartitionOffset;
-            this.remoteLogSegmentMetadataSnapshots = remoteLogSegmentMetadataSnapshots;
-        }
-
-        public short version() {
-            return version;
-        }
-
-        public int metadataPartition() {
-            return metadataPartition;
-        }
-
-        public long metadataPartitionOffset() {
-            return metadataPartitionOffset;
-        }
-
-        public Collection<RemoteLogSegmentMetadataSnapshot> remoteLogSegmentMetadataSnapshots() {
-            return remoteLogSegmentMetadataSnapshots;
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if (!(o instanceof Snapshot)) return false;
-            Snapshot snapshot = (Snapshot) o;
-            return version == snapshot.version && metadataPartition == snapshot.metadataPartition
-                    && metadataPartitionOffset == snapshot.metadataPartitionOffset
-                    && Objects.equals(remoteLogSegmentMetadataSnapshots, snapshot.remoteLogSegmentMetadataSnapshots);
-        }
-
-        @Override
-        public int hashCode() {
-            return Objects.hash(version, metadataPartition, metadataPartitionOffset, remoteLogSegmentMetadataSnapshots);
-        }
-
-        @Override
-        public String toString() {
-            return "Snapshot{" +
-                    "version=" + version +
-                    ", metadataPartition=" + metadataPartition +
-                    ", metadataPartitionOffset=" + metadataPartitionOffset +
-                    ", remoteLogSegmentMetadataSnapshotsSize" + remoteLogSegmentMetadataSnapshots.size() +
-                    '}';
-        }
-    }
-}
diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogSegmentMetadataSnapshot.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogSegmentMetadataSnapshot.java
index ec1ed6a66d15c..de33a05441390 100644
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogSegmentMetadataSnapshot.java
+++ b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogSegmentMetadataSnapshot.java
@@ -33,10 +33,8 @@
 /**
  * This class represents the entry containing the metadata about a remote log segment. This is similar to
  * {@link RemoteLogSegmentMetadata} but it does not contain topic partition information. This class keeps
- * only remote log segment ID but not the topic partition.
- *
- * This class is used in storing the snapshot of remote log metadata for a specific topic partition as mentioned
- * in {@link RemoteLogMetadataSnapshotFile.Snapshot}.
+ * only remote log segment ID but not the topic partition. This class is used in storing the snapshot of
+ * remote log metadata for a specific topic partition.
  */
 public class RemoteLogSegmentMetadataSnapshot extends RemoteLogMetadata {
 
diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataEventHandler.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataEventHandler.java
index f4f43b0d88377..075cab58817c5 100644
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataEventHandler.java
+++ b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataEventHandler.java
@@ -22,8 +22,6 @@
 import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadataUpdate;
 import org.apache.kafka.server.log.remote.storage.RemotePartitionDeleteMetadata;
 
-import java.io.IOException;
-
 public abstract class RemotePartitionMetadataEventHandler {
 
     public void handleRemoteLogMetadata(RemoteLogMetadata remoteLogMetadata) {
@@ -44,9 +42,11 @@ public void handleRemoteLogMetadata(RemoteLogMetadata remoteLogMetadata) {
 
     protected abstract void handleRemotePartitionDeleteMetadata(RemotePartitionDeleteMetadata remotePartitionDeleteMetadata);
 
-    public abstract void syncLogMetadataSnapshot(TopicIdPartition topicIdPartition,
-                                                 int metadataPartition,
-                                                 Long metadataPartitionOffset) throws IOException;
+    public void syncLogMetadataSnapshot(TopicIdPartition topicIdPartition,
+                                        int metadataPartition,
+                                        Long metadataPartitionOffset) {
+        // no-op by default
+    }
 
     public abstract void clearTopicPartition(TopicIdPartition topicIdPartition);
 
diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataStore.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataStore.java
index f9394eee99f36..14f5b3dd7177e 100644
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataStore.java
+++ b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/RemotePartitionMetadataStore.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.server.log.remote.metadata.storage;
 
 import org.apache.kafka.common.TopicIdPartition;
-import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.ReplicaNotAvailableException;
 import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentId;
 import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata;
@@ -30,9 +29,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
-import java.io.File;
 import java.io.IOException;
-import java.nio.file.Path;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.Map;
@@ -46,16 +43,13 @@
 public class RemotePartitionMetadataStore extends RemotePartitionMetadataEventHandler implements Closeable {
     private static final Logger log = LoggerFactory.getLogger(RemotePartitionMetadataStore.class);
 
-    private final Path logDir;
-
     private Map<TopicIdPartition, RemotePartitionDeleteMetadata> idToPartitionDeleteMetadata =
             new ConcurrentHashMap<>();
 
-    private Map<TopicIdPartition, FileBasedRemoteLogMetadataCache> idToRemoteLogMetadataCache =
+    private Map<TopicIdPartition, RemoteLogMetadataCache> idToRemoteLogMetadataCache =
             new ConcurrentHashMap<>();
 
-    public RemotePartitionMetadataStore(Path logDir) {
-        this.logDir = logDir;
+    public RemotePartitionMetadataStore() {
     }
 
     @Override
@@ -74,10 +68,6 @@ public void handleRemoteLogSegmentMetadata(RemoteLogSegmentMetadata remoteLogSeg
         }
     }
 
-    private Path partitionLogDirectory(TopicPartition topicPartition) {
-        return new File(logDir.toFile(), topicPartition.topic() + "-" + topicPartition.partition()).toPath();
-    }
-
     @Override
     public void handleRemoteLogSegmentMetadataUpdate(RemoteLogSegmentMetadataUpdate rlsmUpdate) {
         log.debug("Updating remote log segment: [{}]", rlsmUpdate);
@@ -110,22 +100,6 @@ public void handleRemotePartitionDeleteMetadata(RemotePartitionDeleteMetadata re
         }
     }
 
-    @Override
-    public void syncLogMetadataSnapshot(TopicIdPartition topicIdPartition,
-                                        int metadataPartition,
-                                        Long metadataPartitionOffset) throws IOException {
-        RemotePartitionDeleteMetadata partitionDeleteMetadata = idToPartitionDeleteMetadata.get(topicIdPartition);
-        if (partitionDeleteMetadata != null) {
-            log.info("Skipping syncing of metadata snapshot as remote partition [{}] is with state: [{}] ", topicIdPartition,
-                     partitionDeleteMetadata);
-        } else {
-            FileBasedRemoteLogMetadataCache remoteLogMetadataCache = idToRemoteLogMetadataCache.get(topicIdPartition);
-            if (remoteLogMetadataCache != null) {
-                remoteLogMetadataCache.flushToFile(metadataPartition, metadataPartitionOffset);
-            }
-        }
-    }
-
     @Override
     public void clearTopicPartition(TopicIdPartition topicIdPartition) {
         idToRemoteLogMetadataCache.remove(topicIdPartition);
@@ -145,16 +119,15 @@ public Iterator<RemoteLogSegmentMetadata> listRemoteLogSegments(TopicIdPartition
         return getRemoteLogMetadataCache(topicIdPartition).listRemoteLogSegments(leaderEpoch);
     }
 
-    private FileBasedRemoteLogMetadataCache getRemoteLogMetadataCache(TopicIdPartition topicIdPartition)
+    private RemoteLogMetadataCache getRemoteLogMetadataCache(TopicIdPartition topicIdPartition)
             throws RemoteResourceNotFoundException {
-        FileBasedRemoteLogMetadataCache remoteLogMetadataCache = idToRemoteLogMetadataCache.get(topicIdPartition);
+        RemoteLogMetadataCache remoteLogMetadataCache = idToRemoteLogMetadataCache.get(topicIdPartition);
         if (remoteLogMetadataCache == null) {
             throw new RemoteResourceNotFoundException("No resource found for partition: " + topicIdPartition);
         }
 
         if (!remoteLogMetadataCache.isInitialized()) {
-            // Throwing a retriable ReplicaNotAvailableException here for clients retry. We can introduce a new more
-            // appropriate exception with a KIP in the future.
+            // Throwing a retriable ReplicaNotAvailableException here for clients retry.
             throw new ReplicaNotAvailableException("Remote log metadata cache is not initialized for partition: " + topicIdPartition);
         }
 
@@ -189,8 +162,7 @@ public void close() throws IOException {
 
     @Override
     public void maybeLoadPartition(TopicIdPartition partition) {
-        idToRemoteLogMetadataCache.computeIfAbsent(partition,
-            topicIdPartition -> new FileBasedRemoteLogMetadataCache(topicIdPartition, partitionLogDirectory(topicIdPartition.topicPartition())));
+        idToRemoteLogMetadataCache.computeIfAbsent(partition, idPartition -> new RemoteLogMetadataCache());
     }
 
     @Override
diff --git a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/TopicBasedRemoteLogMetadataManager.java b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/TopicBasedRemoteLogMetadataManager.java
index a3fc1c6e9b87e..b8e3d106664e7 100644
--- a/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/TopicBasedRemoteLogMetadataManager.java
+++ b/storage/src/main/java/org/apache/kafka/server/log/remote/metadata/storage/TopicBasedRemoteLogMetadataManager.java
@@ -40,7 +40,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.File;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
@@ -359,7 +358,7 @@ public void configure(Map<String, ?> configs) {
 
             rlmmConfig = new TopicBasedRemoteLogMetadataManagerConfig(configs);
             rlmTopicPartitioner = new RemoteLogMetadataTopicPartitioner(rlmmConfig.metadataTopicPartitionsCount());
-            remotePartitionMetadataStore = new RemotePartitionMetadataStore(new File(rlmmConfig.logDir()).toPath());
+            remotePartitionMetadataStore = new RemotePartitionMetadataStore();
             configured = true;
             log.info("Successfully configured topic-based RLMM with config: {}", rlmmConfig);
 
diff --git a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/ConsumerTaskTest.java b/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/ConsumerTaskTest.java
index 2b36c4bb03968..424c86b6df652 100644
--- a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/ConsumerTaskTest.java
+++ b/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/ConsumerTaskTest.java
@@ -390,10 +390,6 @@ protected void handleRemoteLogSegmentMetadataUpdate(RemoteLogSegmentMetadataUpda
         protected void handleRemotePartitionDeleteMetadata(RemotePartitionDeleteMetadata remotePartitionDeleteMetadata) {
         }
 
-        @Override
-        public void syncLogMetadataSnapshot(TopicIdPartition topicIdPartition, int metadataPartition, Long metadataPartitionOffset) {
-        }
-
         @Override
         public void clearTopicPartition(TopicIdPartition topicIdPartition) {
             isPartitionCleared.put(topicIdPartition, true);
diff --git a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCacheTest.java b/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCacheTest.java
deleted file mode 100644
index d5341e07b07ad..0000000000000
--- a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/FileBasedRemoteLogMetadataCacheTest.java
+++ /dev/null
@@ -1,91 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.server.log.remote.metadata.storage;
-
-import org.apache.kafka.common.TopicIdPartition;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.Uuid;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentId;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata.CustomMetadata;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadataUpdate;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentState;
-import org.apache.kafka.test.TestUtils;
-import org.junit.jupiter.api.Test;
-
-import java.nio.file.Path;
-import java.util.Collections;
-import java.util.Optional;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-
-public class FileBasedRemoteLogMetadataCacheTest {
-
-    @Test
-    public void testFileBasedRemoteLogMetadataCacheWithUnreferencedSegments() throws Exception {
-        TopicIdPartition partition = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("test", 0));
-        int brokerId = 0;
-        Path path = TestUtils.tempDirectory().toPath();
-
-        // Create file based metadata cache.
-        FileBasedRemoteLogMetadataCache cache = new FileBasedRemoteLogMetadataCache(partition, path);
-
-        // Add a segment with start offset as 0 for leader epoch 0.
-        RemoteLogSegmentId segmentId1 = new RemoteLogSegmentId(partition, Uuid.randomUuid());
-        RemoteLogSegmentMetadata metadata1 = new RemoteLogSegmentMetadata(segmentId1,
-                                                                          0, 100, System.currentTimeMillis(), brokerId, System.currentTimeMillis(),
-                                                                          1024 * 1024, Collections.singletonMap(0, 0L));
-        cache.addCopyInProgressSegment(metadata1);
-        RemoteLogSegmentMetadataUpdate metadataUpdate1 = new RemoteLogSegmentMetadataUpdate(
-                segmentId1, System.currentTimeMillis(),
-                Optional.of(new CustomMetadata(new byte[]{0, 1, 2, 3})),
-                RemoteLogSegmentState.COPY_SEGMENT_FINISHED, brokerId);
-        cache.updateRemoteLogSegmentMetadata(metadataUpdate1);
-        Optional<RemoteLogSegmentMetadata> receivedMetadata = cache.remoteLogSegmentMetadata(0, 0L);
-        assertTrue(receivedMetadata.isPresent());
-        assertEquals(metadata1.createWithUpdates(metadataUpdate1), receivedMetadata.get());
-
-        // Add a new segment with start offset as 0 for leader epoch 0, which should replace the earlier segment.
-        RemoteLogSegmentId segmentId2 = new RemoteLogSegmentId(partition, Uuid.randomUuid());
-        RemoteLogSegmentMetadata metadata2 = new RemoteLogSegmentMetadata(segmentId2,
-                                                                          0, 900, System.currentTimeMillis(), brokerId, System.currentTimeMillis(),
-                                                                          1024 * 1024, Collections.singletonMap(0, 0L));
-        cache.addCopyInProgressSegment(metadata2);
-        RemoteLogSegmentMetadataUpdate metadataUpdate2 = new RemoteLogSegmentMetadataUpdate(
-                segmentId2, System.currentTimeMillis(),
-                Optional.of(new CustomMetadata(new byte[]{4, 5, 6, 7})),
-                RemoteLogSegmentState.COPY_SEGMENT_FINISHED, brokerId);
-        cache.updateRemoteLogSegmentMetadata(metadataUpdate2);
-
-        // Fetch segment for leader epoch:0 and start offset:0, it should be the newly added segment.
-        Optional<RemoteLogSegmentMetadata> receivedMetadata2 = cache.remoteLogSegmentMetadata(0, 0L);
-        assertTrue(receivedMetadata2.isPresent());
-        assertEquals(metadata2.createWithUpdates(metadataUpdate2), receivedMetadata2.get());
-        // Flush the cache to the file.
-        cache.flushToFile(0, 0L);
-
-        // Create a new cache with loading from the stored path.
-        FileBasedRemoteLogMetadataCache loadedCache = new FileBasedRemoteLogMetadataCache(partition, path);
-
-        // Fetch segment for leader epoch:0 and start offset:0, it should be metadata2.
-        // This ensures that the ordering of metadata is taken care after loading from the stored snapshots.
-        Optional<RemoteLogSegmentMetadata> receivedMetadataAfterLoad = loadedCache.remoteLogSegmentMetadata(0, 0L);
-        assertTrue(receivedMetadataAfterLoad.isPresent());
-        assertEquals(metadata2.createWithUpdates(metadataUpdate2), receivedMetadataAfterLoad.get());
-    }
-}
diff --git a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFileTest.java b/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFileTest.java
deleted file mode 100644
index dbfbbf3b0440e..0000000000000
--- a/storage/src/test/java/org/apache/kafka/server/log/remote/metadata/storage/RemoteLogMetadataSnapshotFileTest.java
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.server.log.remote.metadata.storage;
-
-import org.apache.kafka.common.Uuid;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata.CustomMetadata;
-import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentState;
-import org.apache.kafka.test.TestUtils;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-
-import java.io.File;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Optional;
-
-public class RemoteLogMetadataSnapshotFileTest {
-
-    @Test
-    public void testEmptyCommittedLogMetadataFile() throws Exception {
-        File metadataStoreDir = TestUtils.tempDirectory("_rlmm_committed");
-        RemoteLogMetadataSnapshotFile snapshotFile = new RemoteLogMetadataSnapshotFile(metadataStoreDir.toPath());
-
-        // There should be an empty snapshot as nothing is written into it.
-        Assertions.assertFalse(snapshotFile.read().isPresent());
-    }
-
-    @Test
-    public void testEmptySnapshotWithCommittedLogMetadataFile() throws Exception {
-        File metadataStoreDir = TestUtils.tempDirectory("_rlmm_committed");
-        RemoteLogMetadataSnapshotFile snapshotFile = new RemoteLogMetadataSnapshotFile(metadataStoreDir.toPath());
-
-        snapshotFile.write(new RemoteLogMetadataSnapshotFile.Snapshot(0, 0L, Collections.emptyList()));
-
-        // There should be an empty snapshot as the written snapshot did not have any remote log segment metadata.
-        Assertions.assertTrue(snapshotFile.read().isPresent());
-        Assertions.assertTrue(snapshotFile.read().get().remoteLogSegmentMetadataSnapshots().isEmpty());
-    }
-
-    @Test
-    public void testWriteReadCommittedLogMetadataFile() throws Exception {
-        File metadataStoreDir = TestUtils.tempDirectory("_rlmm_committed");
-        RemoteLogMetadataSnapshotFile snapshotFile = new RemoteLogMetadataSnapshotFile(metadataStoreDir.toPath());
-
-        List<RemoteLogSegmentMetadataSnapshot> remoteLogSegmentMetadatas = new ArrayList<>();
-        long startOffset = 0;
-        for (int i = 0; i < 100; i++) {
-            long endOffset = startOffset + 100L;
-            CustomMetadata customMetadata = new CustomMetadata(new byte[]{(byte) i});
-            remoteLogSegmentMetadatas.add(
-                    new RemoteLogSegmentMetadataSnapshot(Uuid.randomUuid(), startOffset, endOffset,
-                                                         System.currentTimeMillis(), 1, 100, 1024,
-                                                         Optional.of(customMetadata),
-                                                         RemoteLogSegmentState.COPY_SEGMENT_FINISHED, Collections.singletonMap(i, startOffset)
-                    ));
-            startOffset = endOffset + 1;
-        }
-
-        RemoteLogMetadataSnapshotFile.Snapshot snapshot = new RemoteLogMetadataSnapshotFile.Snapshot(0, 120,
-                                                                                                     remoteLogSegmentMetadatas);
-        snapshotFile.write(snapshot);
-
-        Optional<RemoteLogMetadataSnapshotFile.Snapshot> maybeReadSnapshot = snapshotFile.read();
-        Assertions.assertTrue(maybeReadSnapshot.isPresent());
-
-        Assertions.assertEquals(snapshot, maybeReadSnapshot.get());
-        Assertions.assertEquals(new HashSet<>(snapshot.remoteLogSegmentMetadataSnapshots()),
-                                new HashSet<>(maybeReadSnapshot.get().remoteLogSegmentMetadataSnapshots()));
-    }
-}

From 2932eb2b4c5cd848c7821932de6c17c85763b936 Mon Sep 17 00:00:00 2001
From: Igor Soarez <i@soarez.me>
Date: Tue, 2 Apr 2024 04:11:32 +0100
Subject: [PATCH 240/258] KAFKA-16365: AssignmentsManager callback handling
 issues (#15521)

When moving replicas between directories in the same broker, future replica promotion hinges on acknowledgment from the controller of a change in the directory assignment.

ReplicaAlterLogDirsThread relies on AssignmentsManager for a completion notification of the directory assignment change.

In its current form, under certain assignment scheduling, AssignmentsManager both miss completion notifications, or prematurely trigger them.

Reviewers: Luke Chen <showuon@gmail.com>, Omnia Ibrahim <o.g.h.ibrahim@gmail.com>, Gaurav Narula <gaurav_narula2@apple.com>
---
 .../kafka/server/AssignmentsManager.java      | 31 ++++----
 .../kafka/server/AssignmentsManagerTest.java  | 76 +++++++++++++++----
 2 files changed, 76 insertions(+), 31 deletions(-)

diff --git a/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java b/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
index 241e2d3aae99f..62b18568d8fab 100644
--- a/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
+++ b/server/src/main/java/org/apache/kafka/server/AssignmentsManager.java
@@ -179,9 +179,6 @@ void merge(AssignmentEvent other) {
             if (!partition.equals(other.partition)) {
                 throw new IllegalArgumentException("Cannot merge events for different partitions");
             }
-            if (!dirId.equals(other.dirId)) {
-                throw new IllegalArgumentException("Cannot merge events for different directories");
-            }
             completionHandlers.addAll(other.completionHandlers);
         }
         void onComplete() {
@@ -191,25 +188,29 @@ void onComplete() {
         }
         @Override
         public void run() throws Exception {
+            log.trace("Received assignment {}", this);
             AssignmentEvent existing = pending.getOrDefault(partition, null);
+            boolean existingIsInFlight = false;
             if (existing == null && inflight != null) {
                 existing = inflight.getOrDefault(partition, null);
+                existingIsInFlight = true;
             }
             if (existing != null) {
                 if (existing.dirId.equals(dirId)) {
                     existing.merge(this);
-                    if (log.isDebugEnabled()) log.debug("Ignoring duplicate assignment {}", this);
+                    log.debug("Ignoring duplicate assignment {}", this);
                     return;
                 }
                 if (existing.timestampNs > timestampNs) {
-                    existing.onComplete();
-                    if (log.isDebugEnabled()) log.debug("Dropping assignment {} because it's older than {}", this, existing);
+                    existing.merge(this);
+                    log.debug("Dropping assignment {} because it's older than existing {}", this, existing);
                     return;
+                } else if (!existingIsInFlight) {
+                    this.merge(existing);
+                    log.debug("Dropping existing assignment {} because it's older than {}", existing, this);
                 }
             }
-            if (log.isDebugEnabled()) {
-                log.debug("Received new assignment {}", this);
-            }
+            log.debug("Queueing new assignment {}", this);
             pending.put(partition, this);
 
             if (inflight == null || inflight.isEmpty()) {
@@ -270,9 +271,7 @@ public void run() throws Exception {
             }
             Map<TopicIdPartition, Uuid> assignment = inflight.entrySet().stream()
                     .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().dirId));
-            if (log.isDebugEnabled()) {
-                log.debug("Dispatching {} assignments:  {}", assignment.size(), assignment);
-            }
+            log.debug("Dispatching {} assignments:  {}", assignment.size(), assignment);
             channelManager.sendRequest(new AssignReplicasToDirsRequest.Builder(
                     buildRequestData(brokerId, brokerEpochSupplier.get(), assignment)),
                     new AssignReplicasToDirsRequestCompletionHandler());
@@ -329,9 +328,7 @@ public void onTimeout() {
         }
         @Override
         public void onComplete(ClientResponse response) {
-            if (log.isDebugEnabled()) {
-                log.debug("Received controller response: {}", response);
-            }
+            log.debug("Received controller response: {}", response);
             appendResponseEvent(response);
         }
         void appendResponseEvent(ClientResponse response) {
@@ -350,9 +347,7 @@ private void scheduleDispatch() {
     }
 
     private void scheduleDispatch(long delayNs) {
-        if (log.isDebugEnabled()) {
-            log.debug("Scheduling dispatch in {}ns", delayNs);
-        }
+        log.debug("Scheduling dispatch in {}ns", delayNs);
         eventQueue.enqueue(EventQueue.EventInsertionType.DEFERRED, DispatchEvent.TAG,
                 new EventQueue.LatestDeadlineFunction(time.nanoseconds() + delayNs), new DispatchEvent());
     }
diff --git a/server/src/test/java/org/apache/kafka/server/AssignmentsManagerTest.java b/server/src/test/java/org/apache/kafka/server/AssignmentsManagerTest.java
index 043afde938684..25d58f4339dcd 100644
--- a/server/src/test/java/org/apache/kafka/server/AssignmentsManagerTest.java
+++ b/server/src/test/java/org/apache/kafka/server/AssignmentsManagerTest.java
@@ -45,11 +45,14 @@
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 
 import static org.apache.kafka.metadata.AssignmentsHelper.buildRequestData;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.atMostOnce;
 import static org.mockito.Mockito.doAnswer;
@@ -282,19 +285,7 @@ void testOnCompletion() throws Exception {
         doAnswer(invocation -> {
             AssignReplicasToDirsRequestData request = invocation.getArgument(0, AssignReplicasToDirsRequest.Builder.class).build().data();
             ControllerRequestCompletionHandler completionHandler = invocation.getArgument(1, ControllerRequestCompletionHandler.class);
-            Map<Uuid, Map<TopicIdPartition, Errors>> errors = new HashMap<>();
-            for (AssignReplicasToDirsRequestData.DirectoryData directory : request.directories()) {
-                for (AssignReplicasToDirsRequestData.TopicData topic : directory.topics()) {
-                    for (AssignReplicasToDirsRequestData.PartitionData partition : topic.partitions()) {
-                        TopicIdPartition topicIdPartition = new TopicIdPartition(topic.topicId(), partition.partitionIndex());
-                        errors.computeIfAbsent(directory.id(), d -> new HashMap<>()).put(topicIdPartition, Errors.NONE);
-                    }
-                }
-            }
-            AssignReplicasToDirsResponseData responseData = AssignmentsHelper.buildResponseData(Errors.NONE.code(), 0, errors);
-            completionHandler.onComplete(new ClientResponse(null, null, null,
-                    0L, 0L, false, false, null, null,
-                            new AssignReplicasToDirsResponse(responseData)));
+            completionHandler.onComplete(buildSuccessfulResponse(request));
 
             return null;
         }).when(channelManager).sendRequest(any(AssignReplicasToDirsRequest.Builder.class),
@@ -310,6 +301,65 @@ void testOnCompletion() throws Exception {
         }
     }
 
+    private static ClientResponse buildSuccessfulResponse(AssignReplicasToDirsRequestData request) {
+        Map<Uuid, Map<TopicIdPartition, Errors>> errors = new HashMap<>();
+        for (AssignReplicasToDirsRequestData.DirectoryData directory : request.directories()) {
+            for (AssignReplicasToDirsRequestData.TopicData topic : directory.topics()) {
+                for (AssignReplicasToDirsRequestData.PartitionData partition : topic.partitions()) {
+                    TopicIdPartition topicIdPartition = new TopicIdPartition(topic.topicId(), partition.partitionIndex());
+                    errors.computeIfAbsent(directory.id(), d -> new HashMap<>()).put(topicIdPartition, Errors.NONE);
+                }
+            }
+        }
+        AssignReplicasToDirsResponseData responseData = AssignmentsHelper.buildResponseData(Errors.NONE.code(), 0, errors);
+        ClientResponse response = new ClientResponse(null, null, null,
+                0L, 0L, false, false, null, null,
+                new AssignReplicasToDirsResponse(responseData));
+        return response;
+    }
+
+    @Test
+    public void testAssignmentCompaction() throws Exception {
+        // Delay the first controller response to force assignment compaction logic
+        CompletableFuture<Runnable> completionFuture = new CompletableFuture<>();
+        doAnswer(invocation -> {
+            AssignReplicasToDirsRequestData request = invocation.getArgument(0, AssignReplicasToDirsRequest.Builder.class).build().data();
+            ControllerRequestCompletionHandler completionHandler = invocation.getArgument(1, ControllerRequestCompletionHandler.class);
+            ClientResponse response = buildSuccessfulResponse(request);
+            Runnable completion = () -> completionHandler.onComplete(response);
+            if (completionFuture.isDone()) completion.run();
+            else completionFuture.complete(completion);
+            return null;
+        }).when(channelManager).sendRequest(any(AssignReplicasToDirsRequest.Builder.class),
+                any(ControllerRequestCompletionHandler.class));
+
+        CountDownLatch remainingInvocations = new CountDownLatch(20);
+        Runnable onComplete = () -> {
+            assertTrue(completionFuture.isDone(), "Premature invocation");
+            assertTrue(remainingInvocations.getCount() > 0, "Extra invocation");
+            remainingInvocations.countDown();
+        };
+        Uuid[] dirs = {DIR_1, DIR_2, DIR_3};
+        for (int i = 0; i < remainingInvocations.getCount(); i++) {
+            time.sleep(100);
+            manager.onAssignment(new TopicIdPartition(TOPIC_1, 0), dirs[i % 3], onComplete);
+        }
+        activeWait(completionFuture::isDone);
+        completionFuture.get().run();
+        activeWait(() -> remainingInvocations.getCount() == 0);
+    }
+
+    void activeWait(Supplier<Boolean> predicate) throws InterruptedException {
+        TestUtils.waitForCondition(() -> {
+            boolean conditionSatisfied = predicate.get();
+            if (!conditionSatisfied) {
+                time.sleep(100);
+                manager.wakeup();
+            }
+            return conditionSatisfied;
+        }, TestUtils.DEFAULT_MAX_WAIT_MS, 50, null);
+    }
+
     static Metric findMetric(String name) {
         for (Map.Entry<MetricName, Metric> entry : KafkaYammerMetrics.defaultRegistry().allMetrics().entrySet()) {
             MetricName metricName = entry.getKey();

From 7a10f4a17ef0ade273ad561cbf02ebbc776e83d1 Mon Sep 17 00:00:00 2001
From: "Cheng-Kai, Zhang" <kai821104@gmail.com>
Date: Tue, 2 Apr 2024 14:54:01 +0800
Subject: [PATCH 241/258] MINOR: enhance kafka-reassign-partitions command
 output (#15610)

Currently, when we using kafka-reassign-partitions to move the log directory, the output only indicates which replica's movement has successfully started.

This PR propose to show more detailed information, helping end users understand that the operation is proceeding as expected.

Reviewers: Luke Chen <showuon@gmail.com>, Andrew Schofield <aschofield@confluent.io>
---
 .../tools/reassign/ReassignPartitionsCommand.java      | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
index 47c4a8234a177..a23a25e6f57b5 100644
--- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
@@ -806,12 +806,10 @@ private static void executeMoves(Admin adminClient,
         do {
             Set<TopicPartitionReplica> completed = alterReplicaLogDirs(adminClient, pendingReplicas);
             if (!completed.isEmpty()) {
-                System.out.printf("Successfully started log directory move%s for: %s%n",
-                    completed.size() == 1 ? "" : "s",
-                    completed.stream()
-                        .sorted(ReassignPartitionsCommand::compareTopicPartitionReplicas)
-                        .map(Object::toString)
-                        .collect(Collectors.joining(",")));
+                completed.stream().sorted(ReassignPartitionsCommand::compareTopicPartitionReplicas).forEach(replica -> {
+                    System.out.printf("Successfully started moving log directory to %s for replica %s-%s with broker %s %n",
+                        pendingReplicas.get(replica), replica.topic(), replica.partition(), replica.brokerId());
+                });
             }
             completed.forEach(pendingReplicas::remove);
             if (pendingReplicas.isEmpty()) {

From b3116f4f76ebc8a074e0d7ce38bf46981da44723 Mon Sep 17 00:00:00 2001
From: Jeff Kim <kimkb2011@gmail.com>
Date: Tue, 2 Apr 2024 06:16:02 -0400
Subject: [PATCH 242/258] KAFKA-16148: Implement
 GroupMetadataManager#onUnloaded (#15446)

This patch completes all awaiting futures when a group is unloaded.

Reviewers: David Jacot <djacot@confluent.io>
---
 .../group/GroupCoordinatorShard.java          |   1 +
 .../group/GroupMetadataManager.java           |  44 +++++++-
 .../group/classic/ClassicGroup.java           |  11 --
 .../group/GroupCoordinatorShardTest.java      |  22 ++++
 .../group/GroupMetadataManagerTest.java       | 103 ++++++++++++++++++
 .../GroupMetadataManagerTestContext.java      |   4 +
 6 files changed, 172 insertions(+), 13 deletions(-)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
index 12c194c331b1f..be4a9bf7d0ade 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorShard.java
@@ -639,6 +639,7 @@ public void onLoaded(MetadataImage newImage) {
     public void onUnloaded() {
         timer.cancel(GROUP_EXPIRATION_KEY);
         coordinatorMetrics.deactivateMetricsShard(metricsShard);
+        groupMetadataManager.onUnloaded();
     }
 
     /**
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 9068ad17efc8f..9cfe8f617e630 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -275,6 +275,7 @@ GroupMetadataManager build() {
             );
         }
     }
+
     /**
      * The log context.
      */
@@ -1920,6 +1921,47 @@ public void onLoaded() {
         });
     }
 
+    /**
+     * Called when the partition is unloaded.
+     * ClassicGroup: Complete all awaiting join and sync futures. Transition group to Dead.
+     */
+    public void onUnloaded() {
+        groups.values().forEach(group -> {
+            switch (group.type()) {
+                case CONSUMER:
+                    ConsumerGroup consumerGroup = (ConsumerGroup) group;
+                    log.info("[GroupId={}] Unloaded group metadata for group epoch {}.",
+                        consumerGroup.groupId(), consumerGroup.groupEpoch());
+                    break;
+                case CLASSIC:
+                    ClassicGroup classicGroup = (ClassicGroup) group;
+                    log.info("[GroupId={}] Unloading group metadata for generation {}.",
+                        classicGroup.groupId(), classicGroup.generationId());
+
+                    classicGroup.transitionTo(DEAD);
+                    switch (classicGroup.previousState()) {
+                        case EMPTY:
+                        case DEAD:
+                            break;
+                        case PREPARING_REBALANCE:
+                            classicGroup.allMembers().forEach(member -> {
+                                classicGroup.completeJoinFuture(member, new JoinGroupResponseData()
+                                    .setMemberId(member.memberId())
+                                    .setErrorCode(NOT_COORDINATOR.code()));
+                            });
+
+                            break;
+                        case COMPLETING_REBALANCE:
+                        case STABLE:
+                            classicGroup.allMembers().forEach(member -> {
+                                classicGroup.completeSyncFuture(member, new SyncGroupResponseData()
+                                    .setErrorCode(NOT_COORDINATOR.code()));
+                            });
+                    }
+            }
+        });
+    }
+
     public static String consumerGroupSessionTimeoutKey(String groupId, String memberId) {
         return "session-timeout-" + groupId + "-" + memberId;
     }
@@ -3088,7 +3130,6 @@ private CoordinatorResult<Void, Record> updateStaticMemberThenRebalanceOrComplet
 
                         responseFuture.complete(
                             new JoinGroupResponseData()
-                                .setMembers(Collections.emptyList())
                                 .setMemberId(UNKNOWN_MEMBER_ID)
                                 .setGenerationId(group.generationId())
                                 .setProtocolName(group.protocolName().orElse(null))
@@ -3111,7 +3152,6 @@ private CoordinatorResult<Void, Record> updateStaticMemberThenRebalanceOrComplet
                         );
                     } else {
                         group.completeJoinFuture(newMember, new JoinGroupResponseData()
-                            .setMembers(Collections.emptyList())
                             .setMemberId(newMemberId)
                             .setGenerationId(group.generationId())
                             .setProtocolName(group.protocolName().orElse(null))
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/classic/ClassicGroup.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/classic/ClassicGroup.java
index 49b087ad3df49..3ba31d5d8551b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/classic/ClassicGroup.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/classic/ClassicGroup.java
@@ -71,16 +71,6 @@
  */
 public class ClassicGroup implements Group {
 
-    /**
-     * Empty generation.
-     */
-    public static final int NO_GENERATION = -1;
-
-    /**
-     * Protocol with empty name.
-     */
-    public static final String NO_PROTOCOL_NAME = "";
-
     /**
      * No leader.
      */
@@ -545,7 +535,6 @@ public ClassicGroupMember replaceStaticMember(
         JoinGroupResponseData joinGroupResponse = new JoinGroupResponseData()
             .setMembers(Collections.emptyList())
             .setMemberId(oldMemberId)
-            .setGenerationId(NO_GENERATION)
             .setProtocolName(null)
             .setProtocolType(null)
             .setLeader(NO_LEADER)
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
index 59868f36f10bd..19c4b366a924b 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupCoordinatorShardTest.java
@@ -1056,4 +1056,26 @@ public void testOnPartitionsDeleted() {
         assertEquals(records, result.records());
         assertNull(result.response());
     }
+
+    @Test
+    public void testOnUnloaded() {
+        GroupMetadataManager groupMetadataManager = mock(GroupMetadataManager.class);
+        OffsetMetadataManager offsetMetadataManager = mock(OffsetMetadataManager.class);
+        Time mockTime = new MockTime();
+        MockCoordinatorTimer<Void, Record> timer = new MockCoordinatorTimer<>(mockTime);
+        GroupCoordinatorShard coordinator = new GroupCoordinatorShard(
+            new LogContext(),
+            groupMetadataManager,
+            offsetMetadataManager,
+            mockTime,
+            timer,
+            GroupCoordinatorConfigTest.createGroupCoordinatorConfig(4096, 1000L, 24 * 60),
+            mock(CoordinatorMetrics.class),
+            mock(CoordinatorMetricsShard.class)
+        );
+
+        coordinator.onUnloaded();
+        assertEquals(0, timer.size());
+        verify(groupMetadataManager, times(1)).onUnloaded();
+    }
 }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index dc21a2140d2de..81c582ed4d4b0 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -94,6 +94,7 @@
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
+import static org.apache.kafka.common.protocol.Errors.NOT_COORDINATOR;
 import static org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH;
 import static org.apache.kafka.common.utils.Utils.mkSet;
 import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
@@ -113,6 +114,7 @@
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupSyncKey;
 import static org.apache.kafka.coordinator.group.RecordHelpersTest.mkMapOfPartitionRacks;
+import static org.apache.kafka.coordinator.group.classic.ClassicGroupMember.EMPTY_ASSIGNMENT;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.COMPLETING_REBALANCE;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.DEAD;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.EMPTY;
@@ -9520,6 +9522,107 @@ public void testClassicGroupJoinWithEmptyConsumerGroup() throws Exception {
         );
     }
 
+    @Test
+    public void testClassicGroupOnUnloadedEmptyAndPreparingRebalance() throws Exception {
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .build();
+
+        ClassicGroup emptyGroup = context.createClassicGroup("empty-group");
+        assertTrue(emptyGroup.isInState(EMPTY));
+
+        ClassicGroup preparingGroup = context.createClassicGroup("preparing-group");
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
+            .withGroupId("preparing-group")
+            .withMemberId(UNKNOWN_MEMBER_ID)
+            .withDefaultProtocolTypeAndProtocols()
+            .build();
+
+        // preparing-group should have 2 members.
+        GroupMetadataManagerTestContext.JoinResult joinResult1 = context.sendClassicGroupJoin(request);
+        GroupMetadataManagerTestContext.JoinResult joinResult2 = context.sendClassicGroupJoin(request);
+
+        assertFalse(joinResult1.joinFuture.isDone());
+        assertFalse(joinResult2.joinFuture.isDone());
+        assertTrue(preparingGroup.isInState(PREPARING_REBALANCE));
+        assertEquals(2, preparingGroup.size());
+
+        context.onUnloaded();
+
+        assertTrue(emptyGroup.isInState(DEAD));
+        assertTrue(preparingGroup.isInState(DEAD));
+        assertTrue(joinResult1.joinFuture.isDone());
+        assertTrue(joinResult2.joinFuture.isDone());
+        assertEquals(new JoinGroupResponseData()
+            .setMemberId(joinResult1.joinFuture.get().memberId())
+            .setMembers(Collections.emptyList())
+            .setErrorCode(NOT_COORDINATOR.code()), joinResult1.joinFuture.get());
+
+        assertEquals(new JoinGroupResponseData()
+            .setMemberId(joinResult2.joinFuture.get().memberId())
+            .setMembers(Collections.emptyList())
+            .setErrorCode(NOT_COORDINATOR.code()), joinResult2.joinFuture.get());
+    }
+
+    @Test
+    public void testClassicGroupOnUnloadedCompletingRebalance() throws Exception {
+        GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
+            .build();
+        ClassicGroup group = context.createClassicGroup("group-id");
+
+        // Set up a group in with a leader, follower, and a pending member.
+        // Have the pending member join the group and both the pending member
+        // and the follower sync. We should have 2 members awaiting sync.
+        GroupMetadataManagerTestContext.PendingMemberGroupResult pendingGroupResult = context.setupGroupWithPendingMember(group);
+        String pendingMemberId = pendingGroupResult.pendingMemberResponse.memberId();
+
+        // Compete join group for the pending member
+        JoinGroupRequestData request = new GroupMetadataManagerTestContext.JoinGroupRequestBuilder()
+            .withGroupId("group-id")
+            .withMemberId(pendingMemberId)
+            .withDefaultProtocolTypeAndProtocols()
+            .build();
+
+        GroupMetadataManagerTestContext.JoinResult joinResult = context.sendClassicGroupJoin(request);
+
+        assertTrue(joinResult.records.isEmpty());
+        assertTrue(joinResult.joinFuture.isDone());
+        assertEquals(Errors.NONE.code(), joinResult.joinFuture.get().errorCode());
+        assertEquals(3, group.allMembers().size());
+        assertEquals(0, group.numPendingJoinMembers());
+
+        // Follower and pending send SyncGroup request.
+        // Follower and pending member should be awaiting sync while the leader is pending sync.
+        GroupMetadataManagerTestContext.SyncResult followerSyncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
+                .withGroupId("group-id")
+                .withMemberId(pendingGroupResult.followerId)
+                .withGenerationId(joinResult.joinFuture.get().generationId())
+                .build());
+
+        GroupMetadataManagerTestContext.SyncResult pendingMemberSyncResult = context.sendClassicGroupSync(
+            new GroupMetadataManagerTestContext.SyncGroupRequestBuilder()
+                .withGroupId("group-id")
+                .withMemberId(pendingMemberId)
+                .withGenerationId(joinResult.joinFuture.get().generationId())
+                .build());
+
+        assertFalse(followerSyncResult.syncFuture.isDone());
+        assertFalse(pendingMemberSyncResult.syncFuture.isDone());
+        assertTrue(group.isInState(COMPLETING_REBALANCE));
+
+        context.onUnloaded();
+
+        assertTrue(group.isInState(DEAD));
+        assertTrue(followerSyncResult.syncFuture.isDone());
+        assertTrue(pendingMemberSyncResult.syncFuture.isDone());
+        assertEquals(new SyncGroupResponseData()
+            .setAssignment(EMPTY_ASSIGNMENT)
+            .setErrorCode(NOT_COORDINATOR.code()), followerSyncResult.syncFuture.get());
+        assertEquals(new SyncGroupResponseData()
+            .setAssignment(EMPTY_ASSIGNMENT)
+            .setErrorCode(NOT_COORDINATOR.code()), pendingMemberSyncResult.syncFuture.get());
+    }
+
     private static void checkJoinGroupResponse(
         JoinGroupResponseData expectedResponse,
         JoinGroupResponseData actualResponse,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
index 4a7cd8cae9ed0..86b0b12d998c2 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -1274,4 +1274,8 @@ public void replay(
         lastWrittenOffset++;
         snapshotRegistry.getOrCreateSnapshot(lastWrittenOffset);
     }
+
+    void onUnloaded() {
+        groupMetadataManager.onUnloaded();
+    }
 }

From ee61bb721eecb0404929f125fe43392f3d024453 Mon Sep 17 00:00:00 2001
From: Victor van den Hoven <victor.vanden.hoven@alliander.com>
Date: Tue, 2 Apr 2024 15:46:54 +0200
Subject: [PATCH 243/258] KAFKA-15417: move outerJoinBreak-flags out of the
 loop (#15510)

Follow up PR for https://github.com/apache/kafka/pull/14426 to fix a bug introduced by the previous PR.

Cf https://github.com/apache/kafka/pull/14426#discussion_r1518681146

Reviewers: Matthias J. Sax <matthias@confluent.io>
---
 .../kstream/internals/KStreamKStreamJoin.java |  65 +++++-----
 .../KStreamKStreamOuterJoinTest.java          | 111 ++++++++++++++++--
 2 files changed, 140 insertions(+), 36 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
index 124386b9bc3ae..b8b48ff2c4df6 100644
--- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
+++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamJoin.java
@@ -125,9 +125,11 @@ public void init(final ProcessorContext<K, VOut> context) {
         @SuppressWarnings("unchecked")
         @Override
         public void process(final Record<K, V1> record) {
+
             final long inputRecordTimestamp = record.timestamp();
             final long timeFrom = Math.max(0L, inputRecordTimestamp - joinBeforeMs);
             final long timeTo = Math.max(0L, inputRecordTimestamp + joinAfterMs);
+
             sharedTimeTracker.advanceStreamTime(inputRecordTimestamp);
 
             if (outer && record.key() == null && record.value() != null) {
@@ -193,7 +195,6 @@ public void process(final Record<K, V1> record) {
             }
         }
 
-        @SuppressWarnings("unchecked")
         private void emitNonJoinedOuterRecords(
             final KeyValueStore<TimestampedKeyAndJoinSide<K>, LeftOrRightValue<V1, V2>> store,
             final Record<K, V1> record) {
@@ -223,43 +224,35 @@ private void emitNonJoinedOuterRecords(
             try (final KeyValueIterator<TimestampedKeyAndJoinSide<K>, LeftOrRightValue<V1, V2>> it = store.all()) {
                 TimestampedKeyAndJoinSide<K> prevKey = null;
 
+                boolean outerJoinLeftWindowOpen = false;
+                boolean outerJoinRightWindowOpen = false;
                 while (it.hasNext()) {
-                    boolean outerJoinLeftBreak = false;
-                    boolean outerJoinRightBreak = false;
+                    if (outerJoinLeftWindowOpen && outerJoinRightWindowOpen) {
+                        // if windows are open for both joinSides we can break since there are no more candidates to emit
+                        break;
+                    }
                     final KeyValue<TimestampedKeyAndJoinSide<K>, LeftOrRightValue<V1, V2>> next = it.next();
                     final TimestampedKeyAndJoinSide<K> timestampedKeyAndJoinSide = next.key;
-                    final LeftOrRightValue<V1, V2> value = next.value;
-                    final K key = timestampedKeyAndJoinSide.getKey();
                     final long timestamp = timestampedKeyAndJoinSide.getTimestamp();
                     sharedTimeTracker.minTime = timestamp;
 
-                    // Skip next records if window has not closed
+                    // Continue with the next outer record if window for this joinSide has not closed yet
+                    // There might be an outer record for the other joinSide which window has not closed yet
+                    // We rely on the <timestamp><left/right-boolean><key> ordering of KeyValueIterator
                     final long outerJoinLookBackTimeMs = getOuterJoinLookBackTimeMs(timestampedKeyAndJoinSide);
                     if (sharedTimeTracker.minTime + outerJoinLookBackTimeMs + joinGraceMs >= sharedTimeTracker.streamTime) {
                         if (timestampedKeyAndJoinSide.isLeftSide()) {
-                            outerJoinLeftBreak = true; // there are no more candidates to emit on left-outerJoin-side
+                            outerJoinLeftWindowOpen = true; // there are no more candidates to emit on left-outerJoin-side
                         } else {
-                            outerJoinRightBreak = true; // there are no more candidates to emit on right-outerJoin-side
+                            outerJoinRightWindowOpen = true; // there are no more candidates to emit on right-outerJoin-side
                         }
-                        if (outerJoinLeftBreak && outerJoinRightBreak) {
-                            break; // there are no more candidates to emit on left-outerJoin-side and
-                                    // right-outerJoin-side
-                        } else {
-                            continue; // there are possibly candidates left on the other outerJoin-side
-                        }
-                    }
-
-                    final VOut nullJoinedValue;
-                    if (isLeftSide) {
-                        nullJoinedValue = joiner.apply(key,
-                                value.getLeftValue(),
-                                value.getRightValue());
-                    } else {
-                        nullJoinedValue = joiner.apply(key,
-                                (V1) value.getRightValue(),
-                                (V2) value.getLeftValue());
+                        // We continue with the next outer record
+                        continue;
                     }
-
+                    
+                    final K key = timestampedKeyAndJoinSide.getKey();
+                    final LeftOrRightValue<V1, V2> leftOrRightValue = next.value;
+                    final VOut nullJoinedValue = getNullJoinedValue(key, leftOrRightValue);
                     context().forward(
                         record.withKey(key).withValue(nullJoinedValue).withTimestamp(timestamp)
                     );
@@ -272,7 +265,6 @@ private void emitNonJoinedOuterRecords(
                         // we do not use delete() calls since it would incur extra get()
                         store.put(prevKey, null);
                     }
-
                     prevKey = timestampedKeyAndJoinSide;
                 }
 
@@ -283,7 +275,24 @@ private void emitNonJoinedOuterRecords(
             }
         }
 
-        private long getOuterJoinLookBackTimeMs(final TimestampedKeyAndJoinSide<K> timestampedKeyAndJoinSide) {
+        @SuppressWarnings("unchecked")
+        private VOut getNullJoinedValue(
+            final K key, 
+            final LeftOrRightValue<V1, V2> leftOrRightValue) {
+            // depending on the JoinSide fill in the joiner key and joiner values
+            if (isLeftSide) {
+                return joiner.apply(key,
+                        leftOrRightValue.getLeftValue(),
+                        leftOrRightValue.getRightValue());
+            } else {
+                return joiner.apply(key,
+                        (V1) leftOrRightValue.getRightValue(),
+                        (V2) leftOrRightValue.getLeftValue());
+            }
+        }
+
+        private long getOuterJoinLookBackTimeMs(
+            final TimestampedKeyAndJoinSide<K> timestampedKeyAndJoinSide) {
             // depending on the JoinSide we fill in the outerJoinLookBackTimeMs
             if (timestampedKeyAndJoinSide.isLeftSide()) {
                 return windowsAfterMs; // On the left-JoinSide we look back in time
diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
index 28a5f1488fbce..279c21ef61a40 100644
--- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamKStreamOuterJoinTest.java
@@ -499,10 +499,10 @@ public void testGracePeriod() {
             // joined records because the window has ended, but will not produce non-joined records because the window has not closed.
             // w1 = { 0:A0 (ts: 0) }
             // w2 = { 1:a1 (ts: 0) }
-            // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 0) }
-            // --> w2 = { 0:a0 (ts: 101), 1:a1 (ts: 101) }
-            inputTopic2.pipeInput(0, "a0", 101L);
+            // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 101) }
+            // --> w2 = { 1:a1 (ts: 0), 0:a0 (ts: 101) }
             inputTopic1.pipeInput(1, "A1", 101L);
+            inputTopic2.pipeInput(0, "a0", 101L);
             processor.checkAndClearProcessResult();
 
             // push a dummy item to the any stream after the window is closed; this should produced all expired non-joined records because
@@ -511,7 +511,7 @@ public void testGracePeriod() {
             // w2 = { 0:a0 (ts: 101), 1:a1 (ts: 101) }
             // --> w1 = { 0:A0 (ts: 0), 1:A1 (ts: 0) }
             // --> w2 = { 0:a0 (ts: 101), 1:a1 (ts: 101), 0:dummy (ts: 112) }
-            inputTopic2.pipeInput(0, "dummy", 211);
+            inputTopic2.pipeInput(0, "dummy", 112);
             processor.checkAndClearProcessResult(
                 new KeyValueTimestamp<>(1, "null+a1", 0L),
                 new KeyValueTimestamp<>(0, "A0+null", 0L)
@@ -519,6 +519,101 @@ public void testGracePeriod() {
         }
     }
 
+    @Test
+    public void testEmitAllNonJoinedResultsForAsymmetricWindow() {
+        final StreamsBuilder builder = new StreamsBuilder();
+
+        final KStream<Integer, String> stream1;
+        final KStream<Integer, String> stream2;
+        final KStream<Integer, String> joined;
+        final MockApiProcessorSupplier<Integer, String, Void, Void> supplier = new MockApiProcessorSupplier<>();
+        stream1 = builder.stream(topic1, consumed);
+        stream2 = builder.stream(topic2, consumed);
+
+        joined = stream1.outerJoin(
+            stream2,
+            MockValueJoiner.TOSTRING_JOINER,
+            JoinWindows.ofTimeDifferenceWithNoGrace(ofMillis(5)).after(ofMillis(20)),
+            StreamJoined.with(Serdes.Integer(), Serdes.String(), Serdes.String())
+        );
+        joined.process(supplier);
+
+        final Collection<Set<String>> copartitionGroups =
+            TopologyWrapper.getInternalTopologyBuilder(builder.build()).copartitionGroups();
+
+        assertEquals(1, copartitionGroups.size());
+        assertEquals(new HashSet<>(Arrays.asList(topic1, topic2)), copartitionGroups.iterator().next());
+
+        try (final TopologyTestDriver driver = new TopologyTestDriver(builder.build(), PROPS)) {
+            final TestInputTopic<Integer, String> inputTopic1 =
+                driver.createInputTopic(topic1, new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+            final TestInputTopic<Integer, String> inputTopic2 =
+                driver.createInputTopic(topic2, new IntegerSerializer(), new StringSerializer(), Instant.ofEpochMilli(0L), Duration.ZERO);
+            final MockApiProcessor<Integer, String, Void, Void> processor = supplier.theCapturedProcessor();
+
+            // push one item to the primary stream; this should not produce any items because there are no matching keys
+            // and window has not ended
+            // w1 = {}
+            // w2 = {}
+            // --> w1 = { 0:A0 (ts: 29) }
+            // --> w2 = {}
+            inputTopic1.pipeInput(0, "A0", 29L);
+            processor.checkAndClearProcessResult();
+
+            // push another item to the primary stream; this should not produce any items because there are no matching keys
+            // and window has not ended
+            // w1 = { 0:A0 (ts: 29) }
+            // w2 = {}
+            // --> w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // --> w2 = {}
+            inputTopic1.pipeInput(1, "A1", 30L);
+            processor.checkAndClearProcessResult();
+
+            // push one item to the other stream; this should not produce any items because there are no matching keys
+            // and window has not ended
+            // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 30) }
+            // w2 = {}
+            // --> w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // --> w2 = { 2:a2 (ts: 31) }
+            inputTopic2.pipeInput(2, "a2", 31L);
+            processor.checkAndClearProcessResult();
+
+            // push another item to the other stream; this should produce no inner joined-items because there are no matching keys 
+            // and window has not ended
+            // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 30) }
+            // w2 = { 2:a2 (ts: 31) }
+            // --> w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // --> w2 = { 2:a2 (ts: 31), 3:a3 (ts: 36) }
+            inputTopic2.pipeInput(3, "a3", 36L);
+            processor.checkAndClearProcessResult();
+
+            // push another item to the other stream; this should produce no inner joined-items because there are no matching keys 
+            // and should produce a right-join-item because before window has ended
+            // w1 = { 0:A0 (ts: 0), 1:A1 (ts: 30) }
+            // w2 = { 2:a2 (ts: 31), 3:a3 (ts: 36) }
+            // --> w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // --> w2 = { 2:a2 (ts: 31), 3:a3 (ts: 36), 4:a4 (ts: 37) }
+            inputTopic2.pipeInput(4, "a4", 37L);
+            processor.checkAndClearProcessResult(
+                    new KeyValueTimestamp<>(2, "null+a2", 31L)
+            );
+
+            // push another item to the other stream; this should produce no inner joined-items because there are no matching keys 
+            // and should produce a left-join-item because after window has ended
+            // and should produce two right-join-items because before window has ended
+            // w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // w2 = { 2:a0 (ts: 31), 3:a3 (ts: 36), 4:a4 (ts: 37) }
+            // --> w1 = { 0:A0 (ts: 29), 1:A1 (ts: 30) }
+            // --> w2 = { 2:a2 (ts: 31), 3:a3 (ts: 36), 4:a4 (ts: 37), 5:a5 (ts: 50) }
+            inputTopic2.pipeInput(5, "a5", 50L);
+            processor.checkAndClearProcessResult(
+                new KeyValueTimestamp<>(0, "A0+null", 29L),
+                new KeyValueTimestamp<>(3, "null+a3", 36L),
+                new KeyValueTimestamp<>(4, "null+a4", 37L)
+            );
+        }
+    }
+    
     @Test
     public void testOuterJoinWithInMemoryCustomSuppliers() {
         final JoinWindows joinWindows = JoinWindows.ofTimeDifferenceWithNoGrace(ofMillis(100L));
@@ -727,7 +822,7 @@ public void testWindowing() {
     }
 
     @Test
-    public void shouldNotEmitLeftJoinResultForAsymmetricBeforeWindow() {
+    public void testShouldNotEmitLeftJoinResultForAsymmetricBeforeWindow() {
         final StreamsBuilder builder = new StreamsBuilder();
         final int[] expectedKeys = new int[] {0, 1, 2, 3};
 
@@ -814,7 +909,7 @@ public void shouldNotEmitLeftJoinResultForAsymmetricBeforeWindow() {
     }
 
     @Test
-    public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
+    public void testShouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
         final StreamsBuilder builder = new StreamsBuilder();
         final int[] expectedKeys = new int[] {0, 1, 2, 3};
 
@@ -906,7 +1001,7 @@ public void shouldNotEmitLeftJoinResultForAsymmetricAfterWindow() {
      * behavior so that we can make decisions about defining it in the future.
      */
     @Test
-    public void shouldForwardCurrentHeaders() {
+    public void testShouldForwardCurrentHeaders() {
         final StreamsBuilder builder = new StreamsBuilder();
 
         final KStream<Integer, String> stream1;
@@ -1331,7 +1426,7 @@ public KeyValueBytesStoreSupplier keyValueStore(final DslKeyValueParams params)
     }
 
     @Test
-    public void shouldJoinWithNonTimestampedStore() {
+    public void testShouldJoinWithNonTimestampedStore() {
         final CapturingStoreSuppliers suppliers = new CapturingStoreSuppliers();
         final StreamJoined<Integer, String, String> streamJoined =
                 StreamJoined.with(Serdes.Integer(), Serdes.String(), Serdes.String())

From acecd370cc3b25f12926e7a4664a2648f08c6c9a Mon Sep 17 00:00:00 2001
From: rykovsi <45871243+rykovsi@users.noreply.github.com>
Date: Wed, 3 Apr 2024 14:14:47 +0300
Subject: [PATCH 244/258] KAFKA-16457 Useless import class
 `org.apache.kafka.common.config.ConfigDef.Type` (#15646)

Reviewers: Andrew Schofield <aschofield@confluent.io>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../main/java/org/apache/kafka/common/config/SslConfigs.java   | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java b/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java
index 5c3bbcdfb9660..bdf8bc1da78a3 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.common.config;
 
-import org.apache.kafka.common.config.ConfigDef.Type;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.utils.Java;
 import org.apache.kafka.common.utils.Utils;
@@ -144,7 +143,7 @@ public static void addClientSslSupport(ConfigDef config) {
                 .define(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, ConfigDef.Type.STRING, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_LOCATION_DOC)
                 .define(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_PASSWORD_DOC)
                 .define(SslConfigs.SSL_KEY_PASSWORD_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_KEY_PASSWORD_DOC)
-                .define(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_KEY_DOC)
+                .define(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, ConfigDef.Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_KEY_DOC)
                 .define(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, ConfigDef.Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_DOC)
                 .define(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, ConfigDef.Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_DOC)
                 .define(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_TRUSTSTORE_TYPE_DOC)

From 3208c5f487b6c7279499167580a8e175a790bcf6 Mon Sep 17 00:00:00 2001
From: Greg Harris <greg.harris@aiven.io>
Date: Wed, 3 Apr 2024 05:05:16 -0700
Subject: [PATCH 245/258] MINOR: AbstractConfig cleanup Part 2 (#15639)

Reviewers:  Manikumar Reddy <anikumar.reddy@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
---
 .../main/scala/kafka/controller/PartitionStateMachine.scala   | 2 +-
 core/src/main/scala/kafka/server/DynamicBrokerConfig.scala    | 4 ++--
 .../unit/kafka/controller/PartitionStateMachineTest.scala     | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
index 51634291b6d56..5dedad426406b 100755
--- a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
@@ -486,7 +486,7 @@ class ZkPartitionStateMachine(config: KafkaConfig,
     } else {
       val (logConfigs, failed) = zkClient.getLogConfigs(
         partitionsWithNoLiveInSyncReplicas.iterator.map { case (partition, _) => partition.topic }.toSet,
-        config.extractLogConfigMap
+        config.originals()
       )
 
       partitionsWithNoLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
index cae70f795321b..db35fed07c1fc 100755
--- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
@@ -740,13 +740,13 @@ class DynamicLogConfig(logManager: LogManager, server: KafkaBroker) extends Brok
     val originalLogConfig = logManager.currentDefaultConfig
     val originalUncleanLeaderElectionEnable = originalLogConfig.uncleanLeaderElectionEnable
     val newBrokerDefaults = new util.HashMap[String, Object](originalLogConfig.originals)
-    newConfig.extractLogConfigMap.forEach { (k, v) =>
+    newConfig.valuesFromThisConfig.forEach { (k, v) =>
       if (DynamicLogConfig.ReconfigurableConfigs.contains(k)) {
         DynamicLogConfig.KafkaConfigToLogConfigName.get(k).foreach { configName =>
           if (v == null)
              newBrokerDefaults.remove(configName)
           else
-            newBrokerDefaults.put(configName, v)
+            newBrokerDefaults.put(configName, v.asInstanceOf[AnyRef])
         }
       }
     }
diff --git a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
index 183e8657e0d44..10cbe58904564 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
@@ -258,7 +258,7 @@ class PartitionStateMachineTest {
       .thenReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
         TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
-    when(mockZkClient.getLogConfigs(Set.empty, config.extractLogConfigMap))
+    when(mockZkClient.getLogConfigs(Set.empty, config.originals()))
       .thenReturn((Map(partition.topic -> new LogConfig(new Properties)), Map.empty[String, Exception]))
     val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withPartitionEpoch(2)
@@ -434,7 +434,7 @@ class PartitionStateMachineTest {
     }
     prepareMockToGetTopicPartitionsStatesRaw()
     def prepareMockToGetLogConfigs(): Unit = {
-      when(mockZkClient.getLogConfigs(Set.empty, config.extractLogConfigMap)).thenReturn((Map.empty[String, LogConfig], Map.empty[String, Exception]))
+      when(mockZkClient.getLogConfigs(Set.empty, config.originals())).thenReturn((Map.empty[String, LogConfig], Map.empty[String, Exception]))
     }
     prepareMockToGetLogConfigs()
 

From 21479a31bdff0e15cfe7ee0a4e509232ed064b41 Mon Sep 17 00:00:00 2001
From: PoAn Yang <payang@apache.org>
Date: Wed, 3 Apr 2024 21:13:45 +0800
Subject: [PATCH 246/258] KAFKA-16413 add FileLockTest (#15624)

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>
---
 .../test/java/kafka/utils/FileLockTest.java   | 97 +++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 core/src/test/java/kafka/utils/FileLockTest.java

diff --git a/core/src/test/java/kafka/utils/FileLockTest.java b/core/src/test/java/kafka/utils/FileLockTest.java
new file mode 100644
index 0000000000000..a58c223e29444
--- /dev/null
+++ b/core/src/test/java/kafka/utils/FileLockTest.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.utils;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.nio.channels.OverlappingFileLockException;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class FileLockTest {
+    @Test
+    void testLock() {
+        File tempFile = TestUtils.tempFile();
+        FileLock lock1 = new FileLock(tempFile);
+        try {
+            lock1.lock();
+            assertThrows(OverlappingFileLockException.class, lock1::lock);
+
+            FileLock lock2 = new FileLock(tempFile);
+            assertThrows(OverlappingFileLockException.class, lock2::lock);
+            assertFalse(lock2.tryLock());
+            lock1.unlock();
+        } finally {
+            lock1.destroy();
+        }
+    }
+
+    @Test
+    void testTryLock() {
+        File tempFile = TestUtils.tempFile();
+        FileLock lock1 = new FileLock(tempFile);
+        try {
+            assertTrue(lock1.tryLock());
+            assertFalse(lock1.tryLock());
+
+            FileLock lock2 = new FileLock(tempFile);
+            assertFalse(lock2.tryLock());
+            assertThrows(OverlappingFileLockException.class, lock2::lock);
+            lock1.unlock();
+        } finally {
+            lock1.destroy();
+        }
+    }
+
+    @Test
+    void testUnlock() {
+        File tempFile = TestUtils.tempFile();
+        FileLock lock1 = new FileLock(tempFile);
+        try {
+            lock1.lock();
+            lock1.unlock();
+            lock1.lock();
+            lock1.unlock();
+
+            assertTrue(lock1.tryLock());
+            lock1.unlock();
+            assertTrue(lock1.tryLock());
+            lock1.unlock();
+
+            FileLock lock2 = new FileLock(tempFile);
+            assertTrue(lock2.tryLock());
+            lock2.unlock();
+            assertDoesNotThrow(lock2::lock);
+            lock2.unlock();
+        } finally {
+            lock1.destroy();
+        }
+    }
+
+    @Test
+    void testDestroy() {
+        File tempFile = TestUtils.tempFile();
+        FileLock lock1 = new FileLock(tempFile);
+        lock1.destroy();
+        assertFalse(tempFile.exists());
+        assertDoesNotThrow(lock1::destroy);
+    }
+}

From e95e91a0623b3f84438e4ebc8e77d1e37979ef62 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Wed, 3 Apr 2024 10:12:51 -0700
Subject: [PATCH 247/258] =?UTF-8?q?KAFKA-16275:=20Update=20kraft=5Fupgrade?=
 =?UTF-8?q?=5Ftest.py=20to=20support=20KIP-848=E2=80=99s=20group=20protoco?=
 =?UTF-8?q?l=20config=20(#15626)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the methods involved.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Walker Carlson <wcarlson@apache.org>
---
 .../tests/core/kraft_upgrade_test.py          | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/tests/kafkatest/tests/core/kraft_upgrade_test.py b/tests/kafkatest/tests/core/kraft_upgrade_test.py
index 3f3c4a81b103a..2d271b9e061bb 100644
--- a/tests/kafkatest/tests/core/kraft_upgrade_test.py
+++ b/tests/kafkatest/tests/core/kraft_upgrade_test.py
@@ -17,12 +17,12 @@
 from ducktape.mark.resource import cluster
 from ducktape.utils.util import wait_until
 from kafkatest.services.console_consumer import ConsoleConsumer
-from kafkatest.services.kafka import KafkaService
+from kafkatest.services.kafka import KafkaService, consumer_group
 from kafkatest.services.kafka.quorum import isolated_kraft, combined_kraft
 from kafkatest.services.verifiable_producer import VerifiableProducer
 from kafkatest.tests.produce_consume_validate import ProduceConsumeValidateTest
 from kafkatest.utils import is_int
-from kafkatest.version import LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, \
+from kafkatest.version import LATEST_3_1, LATEST_3_2, LATEST_3_3, LATEST_3_4, LATEST_3_5, LATEST_3_7, \
     DEV_BRANCH, KafkaVersion, LATEST_STABLE_METADATA_VERSION
 
 #
@@ -74,7 +74,7 @@ def perform_version_change(self, from_kafka_version):
         self.logger.info("Changing metadata.version to %s" % LATEST_STABLE_METADATA_VERSION)
         self.kafka.upgrade_metadata_version(LATEST_STABLE_METADATA_VERSION)
 
-    def run_upgrade(self, from_kafka_version):
+    def run_upgrade(self, from_kafka_version, group_protocol):
         """Test upgrade of Kafka broker cluster from various versions to the current version
 
         from_kafka_version is a Kafka version to upgrade from.
@@ -101,7 +101,8 @@ def run_upgrade(self, from_kafka_version):
                                            version=KafkaVersion(from_kafka_version))
         self.consumer = ConsoleConsumer(self.test_context, self.num_consumers, self.kafka,
                                         self.topic, new_consumer=True, consumer_timeout_ms=30000,
-                                        message_validator=is_int, version=KafkaVersion(from_kafka_version))
+                                        message_validator=is_int, version=KafkaVersion(from_kafka_version),
+                                        consumer_properties=consumer_group.maybe_set_group_protocol(group_protocol))
         self.run_produce_consume_validate(core_test_action=lambda: self.perform_version_change(from_kafka_version))
         cluster_id = self.kafka.cluster_id()
         assert cluster_id is not None
@@ -112,13 +113,21 @@ def run_upgrade(self, from_kafka_version):
     @matrix(from_kafka_version=[str(LATEST_3_1), str(LATEST_3_2), str(LATEST_3_3), str(LATEST_3_4), str(LATEST_3_5), str(DEV_BRANCH)], 
             use_new_coordinator=[True, False], 
             metadata_quorum=[combined_kraft])
-    def test_combined_mode_upgrade(self, from_kafka_version, metadata_quorum, use_new_coordinator=False):
-        self.run_upgrade(from_kafka_version)
+    @matrix(from_kafka_version=[str(LATEST_3_7), str(DEV_BRANCH)],
+            use_new_coordinator=[True],
+            metadata_quorum=[combined_kraft],
+            group_protocol=consumer_group.all_group_protocols)
+    def test_combined_mode_upgrade(self, from_kafka_version, metadata_quorum, use_new_coordinator=False, group_protocol=None):
+        self.run_upgrade(from_kafka_version, group_protocol)
 
     @cluster(num_nodes=8)
-    @matrix(from_kafka_version=[str(LATEST_3_1), str(LATEST_3_2), str(LATEST_3_3), str(LATEST_3_4), str(LATEST_3_5), str(DEV_BRANCH)], 
-            use_new_coordinator=[True, False], 
+    @matrix(from_kafka_version=[str(LATEST_3_1), str(LATEST_3_2), str(LATEST_3_3), str(LATEST_3_4), str(LATEST_3_5), str(DEV_BRANCH)],
+            use_new_coordinator=[True, False],
             metadata_quorum=[isolated_kraft])
-    def test_isolated_mode_upgrade(self, from_kafka_version, metadata_quorum, use_new_coordinator=False):
-        self.run_upgrade(from_kafka_version)
+    @matrix(from_kafka_version=[str(LATEST_3_7), str(DEV_BRANCH)],
+            use_new_coordinator=[True],
+            metadata_quorum=[isolated_kraft],
+            group_protocol=consumer_group.all_group_protocols)
+    def test_isolated_mode_upgrade(self, from_kafka_version, metadata_quorum, use_new_coordinator=False, group_protocol=None):
+        self.run_upgrade(from_kafka_version, group_protocol)
 

From 6569a354e616e2f4fbf50e25660b5da906b3e764 Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Wed, 3 Apr 2024 10:13:03 -0700
Subject: [PATCH 248/258] =?UTF-8?q?KAFKA-16438:=20Update=20consumer=5Ftest?=
 =?UTF-8?q?.py=E2=80=99s=20static=20tests=20to=20support=20KIP-848?=
 =?UTF-8?q?=E2=80=99s=20group=20protocol=20config=20(#15627)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Migrated the following tests for the new consumer:

- test_fencing_static_consumer
- test_static_consumer_bounce
- test_static_consumer_persisted_after_rejoin

Reviewers: Walker Carlson <wcarlson@apache.org>
---
 tests/kafkatest/tests/client/consumer_test.py | 29 ++++++++++---------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/tests/kafkatest/tests/client/consumer_test.py b/tests/kafkatest/tests/client/consumer_test.py
index d4d6af9f2aad0..89c58e5f973e6 100644
--- a/tests/kafkatest/tests/client/consumer_test.py
+++ b/tests/kafkatest/tests/client/consumer_test.py
@@ -194,7 +194,7 @@ def test_consumer_bounce(self, clean_shutdown, bounce_mode, metadata_quorum=quor
         static_membership=[True, False],
         bounce_mode=["all", "rolling"],
         num_bounces=[5],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
@@ -203,9 +203,10 @@ def test_consumer_bounce(self, clean_shutdown, bounce_mode, metadata_quorum=quor
         bounce_mode=["all", "rolling"],
         num_bounces=[5],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_static_consumer_bounce(self, clean_shutdown, static_membership, bounce_mode, num_bounces, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_static_consumer_bounce(self, clean_shutdown, static_membership, bounce_mode, num_bounces, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify correct static consumer behavior when the consumers in the group are restarted. In order to make
         sure the behavior of static members are different from dynamic ones, we take both static and dynamic
@@ -226,7 +227,7 @@ def test_static_consumer_bounce(self, clean_shutdown, static_membership, bounce_
         self.await_produced_messages(producer)
 
         self.session_timeout_sec = 60
-        consumer = self.setup_consumer(self.TOPIC, static_membership=static_membership)
+        consumer = self.setup_consumer(self.TOPIC, static_membership=static_membership, group_protocol=group_protocol)
 
         consumer.start()
         self.await_all_members(consumer)
@@ -268,15 +269,16 @@ def test_static_consumer_bounce(self, clean_shutdown, static_membership, bounce_
     @cluster(num_nodes=7)
     @matrix(
         bounce_mode=["all", "rolling"],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         bounce_mode=["all", "rolling"],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_static_consumer_persisted_after_rejoin(self, bounce_mode, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_static_consumer_persisted_after_rejoin(self, bounce_mode, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify that the updated member.id(updated_member_id) caused by static member rejoin would be persisted. If not,
         after the brokers rolling bounce, the migrated group coordinator would load the stale persisted member.id and
@@ -291,7 +293,7 @@ def test_static_consumer_persisted_after_rejoin(self, bounce_mode, metadata_quor
         producer.start()
         self.await_produced_messages(producer)
         self.session_timeout_sec = 60
-        consumer = self.setup_consumer(self.TOPIC, static_membership=True)
+        consumer = self.setup_consumer(self.TOPIC, static_membership=True, group_protocol=group_protocol)
         consumer.start()
         self.await_all_members(consumer)
 
@@ -309,16 +311,17 @@ def test_static_consumer_persisted_after_rejoin(self, bounce_mode, metadata_quor
     @matrix(
         num_conflict_consumers=[1, 2],
         fencing_stage=["stable", "all"],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         num_conflict_consumers=[1, 2],
         fencing_stage=["stable", "all"],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_fencing_static_consumer(self, num_conflict_consumers, fencing_stage, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_fencing_static_consumer(self, num_conflict_consumers, fencing_stage, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Verify correct static consumer behavior when there are conflicting consumers with same group.instance.id.
 
@@ -335,10 +338,10 @@ def test_fencing_static_consumer(self, num_conflict_consumers, fencing_stage, me
         self.await_produced_messages(producer)
 
         self.session_timeout_sec = 60
-        consumer = self.setup_consumer(self.TOPIC, static_membership=True)
+        consumer = self.setup_consumer(self.TOPIC, static_membership=True, group_protocol=group_protocol)
 
         self.num_consumers = num_conflict_consumers
-        conflict_consumer = self.setup_consumer(self.TOPIC, static_membership=True)
+        conflict_consumer = self.setup_consumer(self.TOPIC, static_membership=True, group_protocol=group_protocol)
 
         # wait original set of consumer to stable stage before starting conflict members.
         if fencing_stage == "stable":

From 6bb9caced0f34ea27efd91e9c4978f75809367ec Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Wed, 3 Apr 2024 10:13:14 -0700
Subject: [PATCH 249/258] =?UTF-8?q?KAFKA-16440:=20Update=20security=5Ftest?=
 =?UTF-8?q?.py=20to=20support=20KIP-848=E2=80=99s=20group=20protocol=20con?=
 =?UTF-8?q?fig=20(#15628)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the setup_consumer method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Walker Carlson <wcarlson@apache.org>
---
 tests/kafkatest/tests/core/security_test.py | 31 +++++++++++----------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/tests/kafkatest/tests/core/security_test.py b/tests/kafkatest/tests/core/security_test.py
index 8b7c05de8ba20..47aba24dbd693 100644
--- a/tests/kafkatest/tests/core/security_test.py
+++ b/tests/kafkatest/tests/core/security_test.py
@@ -19,7 +19,7 @@
 from ducktape.utils.util import wait_until
 from ducktape.errors import TimeoutError
 
-from kafkatest.services.kafka import quorum
+from kafkatest.services.kafka import quorum, consumer_group
 from kafkatest.services.security.security_config import SecurityConfig
 from kafkatest.services.security.security_config import SslStores
 from kafkatest.tests.end_to_end import EndToEndTest
@@ -61,28 +61,30 @@ def producer_consumer_have_expected_error(self, error):
     @matrix(
         security_protocol=['PLAINTEXT'],
         interbroker_security_protocol=['SSL'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         security_protocol=['PLAINTEXT'],
         interbroker_security_protocol=['SSL'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
     @matrix(
         security_protocol=['SSL'],
         interbroker_security_protocol=['PLAINTEXT'],
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         security_protocol=['SSL'],
         interbroker_security_protocol=['PLAINTEXT'],
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_client_ssl_endpoint_validation_failure(self, security_protocol, interbroker_security_protocol, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_client_ssl_endpoint_validation_failure(self, security_protocol, interbroker_security_protocol, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Test that invalid hostname in certificate results in connection failures.
         When security_protocol=SSL, client SSL handshakes are expected to fail due to hostname verification failure.
@@ -120,11 +122,11 @@ def test_client_ssl_endpoint_validation_failure(self, security_protocol, interbr
             # the inter-broker security protocol using TLS with a hostname verification failure
             # doesn't impact a producer in case of a single broker with a KRaft Controller,
             # so confirm that this is in fact the observed behavior
-            self.create_and_start_clients(log_level="INFO")
+            self.create_and_start_clients(log_level="INFO", group_protocol=group_protocol)
             self.run_validation()
         else:
             # We need more verbose logging to catch the expected errors
-            self.create_and_start_clients(log_level="DEBUG")
+            self.create_and_start_clients(log_level="DEBUG", group_protocol=group_protocol)
 
             try:
                 wait_until(lambda: self.producer.num_acked > 0, timeout_sec=30)
@@ -143,26 +145,27 @@ def test_client_ssl_endpoint_validation_failure(self, security_protocol, interbr
 
             SecurityConfig.ssl_stores.valid_hostname = True
             self.kafka.restart_cluster()
-            self.create_and_start_clients(log_level="INFO")
+            self.create_and_start_clients(log_level="INFO", group_protocol=group_protocol)
             self.run_validation()
 
-    def create_and_start_clients(self, log_level):
+    def create_and_start_clients(self, log_level, group_protocol):
         self.create_producer(log_level=log_level)
         self.producer.start()
 
-        self.create_consumer(log_level=log_level)
+        self.create_consumer(log_level=log_level, group_protocol=group_protocol)
         self.consumer.start()
 
     @cluster(num_nodes=2)
     @matrix(
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_quorum_ssl_endpoint_validation_failure(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_quorum_ssl_endpoint_validation_failure(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         Test that invalid hostname in ZooKeeper or KRaft Controller results in broker inability to start.
         """

From c7ef80bb6c25efbc59c1df8de370d4f893cd8a3b Mon Sep 17 00:00:00 2001
From: Kirk True <kirk@kirktrue.pro>
Date: Wed, 3 Apr 2024 10:13:26 -0700
Subject: [PATCH 250/258] =?UTF-8?q?KAFKA-16439:=20Update=20replication=5Fr?=
 =?UTF-8?q?eplica=5Ffailure=5Ftest.py=20to=20support=20KIP-848=E2=80=99s?=
 =?UTF-8?q?=20group=20protocol=20config=20(#15629)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added a new optional group_protocol parameter to the test methods, then passed that down to the setup_consumer method.

Unfortunately, because the new consumer can only be used with the new coordinator, this required a new @matrix block instead of adding the group_protocol=["classic", "consumer"] to the existing blocks 😢

Reviewers: Walker Carlson <wcarlson@apache.org>
---
 .../tests/core/replication_replica_failure_test.py    | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/kafkatest/tests/core/replication_replica_failure_test.py b/tests/kafkatest/tests/core/replication_replica_failure_test.py
index 66cd1ecac1d8f..f17f6a2d26cae 100644
--- a/tests/kafkatest/tests/core/replication_replica_failure_test.py
+++ b/tests/kafkatest/tests/core/replication_replica_failure_test.py
@@ -20,7 +20,7 @@
 from ducktape.mark import parametrize
 from ducktape.mark.resource import cluster
 
-from kafkatest.services.kafka import quorum
+from kafkatest.services.kafka import quorum, consumer_group
 from kafkatest.tests.end_to_end import EndToEndTest
 from kafkatest.services.kafka import config_property
 from kafkatest.services.trogdor.network_partition_fault_spec import NetworkPartitionFaultSpec
@@ -38,14 +38,15 @@ def __init__(self, test_context):
 
     @cluster(num_nodes=7)
     @matrix(
-        metadata_quorum=[quorum.zk],
+        metadata_quorum=[quorum.zk, quorum.isolated_kraft],
         use_new_coordinator=[False]
     )
     @matrix(
         metadata_quorum=[quorum.isolated_kraft],
-        use_new_coordinator=[True, False]
+        use_new_coordinator=[True],
+        group_protocol=consumer_group.all_group_protocols
     )
-    def test_replication_with_replica_failure(self, metadata_quorum=quorum.zk, use_new_coordinator=False):
+    def test_replication_with_replica_failure(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None):
         """
         This test verifies that replication shrinks the ISR when a replica is not fetching anymore.
         It also verifies that replication provides simple durability guarantees by checking that data acked by
@@ -90,7 +91,7 @@ def test_replication_with_replica_failure(self, metadata_quorum=quorum.zk, use_n
         self.create_producer()
         self.producer.start()
 
-        self.create_consumer()
+        self.create_consumer(group_protocol=group_protocol)
         self.consumer.start()
 
         self.await_startup()

From ef7f823136c2b11ad1a1c62a325d2f5552553a54 Mon Sep 17 00:00:00 2001
From: Apoorv Mittal <amittal@confluent.io>
Date: Thu, 4 Apr 2024 10:56:08 +0100
Subject: [PATCH 251/258] KAFKA-16359: Corrected manifest file for
 kafka-clients (#15532)

The issue KAFKA-16359 reported inclusion of kafka-clients runtime dependencies in MANIFEST.MF Class-Path.

The root cause is the issue defined here with the usage of shadow plugin.

Looking into the specifics of plugin and documentation, specifies that any dependency marked as shadow will be treated as following by the shadow plugin:

    1. Adds the dependency as runtime dependency in resultant pom.xml - code here
    2. Adds the dependency as Class-Path in MANIFEST.MF as well - code here

Resolution

We do need the runtime dependencies available in the pom (1 above) but not on manifest (2 above). Also there is no clear way to separate the behaviour as both above tasks relies on shadow configuration.

To fix, I have defined another custom configuration named shadowed which is later used to populate the correct pom (the code is similar to what shadow plugin does to populate pom for runtime dependencies).

Though this might seem like a workaround, but I think that's the only way to fix the issue. I have checked other SDKs which suffered with same issue and went with similar route to populate pom.

Reviewers: Luke Chen <showuon@gmail.com>, Reviewers: Mickael Maison <mickael.maison@gmail.com>, Gaurav Narula <gaurav_narula2@apple.com>
---
 build.gradle | 34 ++++++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/build.gradle b/build.gradle
index 14ec93512b692..56f81689cc088 100644
--- a/build.gradle
+++ b/build.gradle
@@ -305,6 +305,24 @@ subprojects {
           } else {
             apply plugin: 'com.github.johnrengelman.shadow'
             project.shadow.component(mavenJava)
+
+            // Fix for avoiding inclusion of runtime dependencies marked as 'shadow' in MANIFEST Class-Path.
+            // https://github.com/johnrengelman/shadow/issues/324
+            afterEvaluate {
+              pom.withXml { xml ->
+                if (xml.asNode().get('dependencies') == null) {
+                  xml.asNode().appendNode('dependencies')
+                }
+                def dependenciesNode = xml.asNode().get('dependencies').get(0)
+                project.configurations.shadowed.allDependencies.each {
+                  def dependencyNode = dependenciesNode.appendNode('dependency')
+                  dependencyNode.appendNode('groupId', it.group)
+                  dependencyNode.appendNode('artifactId', it.name)
+                  dependencyNode.appendNode('version', it.version)
+                  dependencyNode.appendNode('scope', 'runtime')
+                }
+              }
+            }
           }
 
           afterEvaluate {
@@ -1403,6 +1421,7 @@ project(':clients') {
 
   configurations {
     generator
+    shadowed
   }
 
   dependencies {
@@ -1413,10 +1432,10 @@ project(':clients') {
     implementation libs.opentelemetryProto
 
     // libraries which should be added as runtime dependencies in generated pom.xml should be defined here:
-    shadow libs.zstd
-    shadow libs.lz4
-    shadow libs.snappy
-    shadow libs.slf4jApi
+    shadowed libs.zstd
+    shadowed libs.lz4
+    shadowed libs.snappy
+    shadowed libs.slf4jApi
 
     compileOnly libs.jacksonDatabind // for SASL/OAUTHBEARER bearer token parsing
     compileOnly libs.jacksonJDK8Datatypes
@@ -1469,10 +1488,9 @@ project(':clients') {
 
     // dependencies excluded from the final jar, since they are declared as runtime dependencies
     dependencies {
-      exclude(dependency(libs.snappy))
-      exclude(dependency(libs.zstd))
-      exclude(dependency(libs.lz4))
-      exclude(dependency(libs.slf4jApi))
+      project.configurations.shadowed.allDependencies.each {
+        exclude(dependency(it.group + ':' + it.name + ':' + it.version))
+      }
       // exclude proto files from the jar
       exclude "**/opentelemetry/proto/**/*.proto"
       exclude "**/google/protobuf/*.proto"

From 376e9e20dbf7c7aeb6f6f666d47932c445eb6bd1 Mon Sep 17 00:00:00 2001
From: Calvin Liu <83986057+CalvinConfluent@users.noreply.github.com>
Date: Thu, 4 Apr 2024 06:12:05 -0700
Subject: [PATCH 252/258] KAFKA-15586: Clean shutdown detection - server side
 (#14706)

If the broker registers with the same broker epoch as the previous session, it is recognized as a clean shutdown. Otherwise, it is an unclean shutdown. This replica will be removed from any ELR.

Reviewers: Artem Livshits <alivshits@confluent.io>, David Arthur <mumrah@gmail.com>
---
 .../kafka/controller/BrokersToElrs.java       | 162 ++++++++++++++++++
 .../controller/ClusterControlManager.java     |  27 ++-
 .../controller/PartitionChangeBuilder.java    |   8 +
 .../kafka/controller/QuorumController.java    |   5 +
 .../controller/ReplicationControlManager.java |  97 +++++++++--
 .../kafka/controller/BrokerToElrsTest.java    |  74 ++++++++
 .../controller/ClusterControlManagerTest.java |  13 ++
 .../ProducerIdControlManagerTest.java         |   1 +
 .../controller/QuorumControllerTest.java      | 138 +++++++++++++++
 .../controller/QuorumControllerTestEnv.java   |   4 +
 .../ReplicationControlManagerTest.java        |  83 +++++++++
 11 files changed, 590 insertions(+), 22 deletions(-)
 create mode 100644 metadata/src/main/java/org/apache/kafka/controller/BrokersToElrs.java
 create mode 100644 metadata/src/test/java/org/apache/kafka/controller/BrokerToElrsTest.java

diff --git a/metadata/src/main/java/org/apache/kafka/controller/BrokersToElrs.java b/metadata/src/main/java/org/apache/kafka/controller/BrokersToElrs.java
new file mode 100644
index 0000000000000..2ea5e3f21d952
--- /dev/null
+++ b/metadata/src/main/java/org/apache/kafka/controller/BrokersToElrs.java
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.controller;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.metadata.Replicas;
+import org.apache.kafka.timeline.SnapshotRegistry;
+import org.apache.kafka.timeline.TimelineHashMap;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+
+import static org.apache.kafka.metadata.Replicas.NONE;
+
+public class BrokersToElrs {
+    private final SnapshotRegistry snapshotRegistry;
+
+    // It maps from the broker id to the topic id partitions if the partition has ELR.
+    private final TimelineHashMap<Integer, TimelineHashMap<Uuid, int[]>> elrMembers;
+
+    BrokersToElrs(SnapshotRegistry snapshotRegistry) {
+        this.snapshotRegistry = snapshotRegistry;
+        this.elrMembers = new TimelineHashMap<>(snapshotRegistry, 0);
+    }
+
+    /**
+     * Update our records of a partition's ELR.
+     *
+     * @param topicId       The topic ID of the partition.
+     * @param partitionId   The partition ID of the partition.
+     * @param prevElr       The previous ELR, or null if the partition is new.
+     * @param nextElr       The new ELR, or null if the partition is being removed.
+     */
+
+    void update(Uuid topicId, int partitionId, int[] prevElr, int[] nextElr) {
+        int[] prev;
+        if (prevElr == null) {
+            prev = NONE;
+        } else {
+            prev = Replicas.clone(prevElr);
+            Arrays.sort(prev);
+        }
+        int[] next;
+        if (nextElr == null) {
+            next = NONE;
+        } else {
+            next = Replicas.clone(nextElr);
+            Arrays.sort(next);
+        }
+
+        int i = 0, j = 0;
+        while (true) {
+            if (i == prev.length) {
+                if (j == next.length) {
+                    break;
+                }
+                int newReplica = next[j];
+                add(newReplica, topicId, partitionId);
+                j++;
+            } else if (j == next.length) {
+                int prevReplica = prev[i];
+                remove(prevReplica, topicId, partitionId);
+                i++;
+            } else {
+                int prevReplica = prev[i];
+                int newReplica = next[j];
+                if (prevReplica < newReplica) {
+                    remove(prevReplica, topicId, partitionId);
+                    i++;
+                } else if (prevReplica > newReplica) {
+                    add(newReplica, topicId, partitionId);
+                    j++;
+                } else {
+                    i++;
+                    j++;
+                }
+            }
+        }
+    }
+
+    void removeTopicEntryForBroker(Uuid topicId, int brokerId) {
+        Map<Uuid, int[]> topicMap = elrMembers.get(brokerId);
+        if (topicMap != null) {
+            topicMap.remove(topicId);
+        }
+    }
+
+    private void add(int brokerId, Uuid topicId, int newPartition) {
+        TimelineHashMap<Uuid, int[]> topicMap = elrMembers.get(brokerId);
+        if (topicMap == null) {
+            topicMap = new TimelineHashMap<>(snapshotRegistry, 0);
+            elrMembers.put(brokerId, topicMap);
+        }
+        int[] partitions = topicMap.get(topicId);
+        int[] newPartitions;
+        if (partitions == null) {
+            newPartitions = new int[1];
+        } else {
+            newPartitions = new int[partitions.length + 1];
+            System.arraycopy(partitions, 0, newPartitions, 0, partitions.length);
+        }
+        newPartitions[newPartitions.length - 1] = newPartition;
+        topicMap.put(topicId, newPartitions);
+    }
+
+    private void remove(int brokerId, Uuid topicId, int removedPartition) {
+        TimelineHashMap<Uuid, int[]> topicMap = elrMembers.get(brokerId);
+        if (topicMap == null) {
+            throw new RuntimeException("Broker " + brokerId + " has no elrMembers " +
+                    "entry, so we can't remove " + topicId + ":" + removedPartition);
+        }
+        int[] partitions = topicMap.get(topicId);
+        if (partitions == null) {
+            throw new RuntimeException("Broker " + brokerId + " has no " +
+                    "entry in elrMembers for topic " + topicId);
+        }
+        if (partitions.length == 1) {
+            if (partitions[0] != removedPartition) {
+                throw new RuntimeException("Broker " + brokerId + " has no " +
+                        "entry in elrMembers for " + topicId + ":" + removedPartition);
+            }
+            topicMap.remove(topicId);
+            if (topicMap.isEmpty()) {
+                elrMembers.remove(brokerId);
+            }
+        } else {
+            int[] newPartitions = new int[partitions.length - 1];
+            int j = 0;
+            for (int i = 0; i < partitions.length; i++) {
+                int partition = partitions[i];
+                if (partition != removedPartition) {
+                    newPartitions[j++] = partition;
+                }
+            }
+            topicMap.put(topicId, newPartitions);
+        }
+    }
+
+    BrokersToIsrs.PartitionsOnReplicaIterator partitionsWithBrokerInElr(int brokerId) {
+        Map<Uuid, int[]> topicMap = elrMembers.get(brokerId);
+        if (topicMap == null) {
+            topicMap = Collections.emptyMap();
+        }
+        return new BrokersToIsrs.PartitionsOnReplicaIterator(topicMap, false);
+    }
+}
diff --git a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java
index 98d64c54835ed..f0bd98776bc38 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java
@@ -91,6 +91,7 @@ static class Builder {
         private ReplicaPlacer replicaPlacer = null;
         private FeatureControlManager featureControl = null;
         private boolean zkMigrationEnabled = false;
+        private BrokerUncleanShutdownHandler brokerUncleanShutdownHandler = null;
 
         Builder setLogContext(LogContext logContext) {
             this.logContext = logContext;
@@ -132,6 +133,11 @@ Builder setZkMigrationEnabled(boolean zkMigrationEnabled) {
             return this;
         }
 
+        Builder setBrokerUncleanShutdownHandler(BrokerUncleanShutdownHandler brokerUncleanShutdownHandler) {
+            this.brokerUncleanShutdownHandler = brokerUncleanShutdownHandler;
+            return this;
+        }
+
         ClusterControlManager build() {
             if (logContext == null) {
                 logContext = new LogContext();
@@ -148,6 +154,9 @@ ClusterControlManager build() {
             if (featureControl == null) {
                 throw new RuntimeException("You must specify FeatureControlManager");
             }
+            if (brokerUncleanShutdownHandler == null) {
+                throw new RuntimeException("You must specify BrokerUncleanShutdownHandler");
+            }
             return new ClusterControlManager(logContext,
                 clusterId,
                 time,
@@ -155,7 +164,8 @@ ClusterControlManager build() {
                 sessionTimeoutNs,
                 replicaPlacer,
                 featureControl,
-                zkMigrationEnabled
+                zkMigrationEnabled,
+                brokerUncleanShutdownHandler
             );
         }
     }
@@ -247,6 +257,8 @@ boolean check() {
      */
     private final boolean zkMigrationEnabled;
 
+    private BrokerUncleanShutdownHandler brokerUncleanShutdownHandler;
+
     /**
      * Maps controller IDs to controller registrations.
      */
@@ -265,7 +277,8 @@ private ClusterControlManager(
         long sessionTimeoutNs,
         ReplicaPlacer replicaPlacer,
         FeatureControlManager featureControl,
-        boolean zkMigrationEnabled
+        boolean zkMigrationEnabled,
+        BrokerUncleanShutdownHandler brokerUncleanShutdownHandler
     ) {
         this.logContext = logContext;
         this.clusterId = clusterId;
@@ -281,6 +294,7 @@ private ClusterControlManager(
         this.zkMigrationEnabled = zkMigrationEnabled;
         this.controllerRegistrations = new TimelineHashMap<>(snapshotRegistry, 0);
         this.directoryToBroker = new TimelineHashMap<>(snapshotRegistry, 0);
+        this.brokerUncleanShutdownHandler = brokerUncleanShutdownHandler;
     }
 
     ReplicaPlacer replicaPlacer() {
@@ -336,10 +350,11 @@ public ControllerResult<BrokerRegistrationReply> registerBroker(
                 ", but got cluster ID " + request.clusterId());
         }
         int brokerId = request.brokerId();
+        List<ApiMessageAndVersion> records = new ArrayList<>();
         BrokerRegistration existing = brokerRegistrations.get(brokerId);
         if (version < 2 || existing == null || request.previousBrokerEpoch() != existing.epoch()) {
-            // TODO(KIP-966): Update the ELR if the broker has an unclean shutdown.
             log.debug("Received an unclean shutdown request");
+            brokerUncleanShutdownHandler.addRecordsForShutdown(request.brokerId(), records);
         }
         if (existing != null) {
             if (heartbeatManager.hasValidSession(brokerId)) {
@@ -410,7 +425,6 @@ public ControllerResult<BrokerRegistrationReply> registerBroker(
 
         heartbeatManager.register(brokerId, record.fenced());
 
-        List<ApiMessageAndVersion> records = new ArrayList<>();
         records.add(new ApiMessageAndVersion(record, featureControl.metadataVersion().
             registerBrokerRecordVersion()));
         return ControllerResult.atomicOf(records, new BrokerRegistrationReply(brokerEpoch));
@@ -780,4 +794,9 @@ public Entry<Integer, Map<String, VersionRange>> next() {
             }
         };
     }
+
+    @FunctionalInterface
+    interface BrokerUncleanShutdownHandler {
+        void addRecordsForShutdown(int brokerId, List<ApiMessageAndVersion> records);
+    }
 }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/PartitionChangeBuilder.java b/metadata/src/main/java/org/apache/kafka/controller/PartitionChangeBuilder.java
index c2a0bbc492888..7f1b2cb6d17a0 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/PartitionChangeBuilder.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/PartitionChangeBuilder.java
@@ -493,6 +493,10 @@ private void maybeUpdateLastKnownLeader(PartitionChangeRecord record) {
             partition.lastKnownElr[0] != partition.leader)) {
             // Only update the last known leader when the first time the partition becomes leaderless.
             record.setLastKnownElr(Arrays.asList(partition.leader));
+        } else if ((record.leader() >= 0 || (partition.leader != NO_LEADER && record.leader() != NO_LEADER))
+            && partition.lastKnownElr.length > 0) {
+            // Clear the LastKnownElr field if the partition will have or continues to have a valid leader.
+            record.setLastKnownElr(Collections.emptyList());
         }
     }
 
@@ -517,6 +521,10 @@ private void maybeUpdateRecordElr(PartitionChangeRecord record) {
             targetLastKnownElr = Replicas.toList(partition.lastKnownElr);
         }
 
+        // If the last known ELR is expected to store the last known leader, the lastKnownElr field should be updated
+        // later in maybeUpdateLastKnownLeader.
+        if (useLastKnownLeaderInBalancedRecovery) return;
+
         if (!targetLastKnownElr.equals(Replicas.toList(partition.lastKnownElr))) {
             record.setLastKnownElr(targetLastKnownElr);
         }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java b/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java
index e83badc2bf36d..557547be13c2d 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java
@@ -1861,6 +1861,7 @@ private QuorumController(
             setReplicaPlacer(replicaPlacer).
             setFeatureControlManager(featureControl).
             setZkMigrationEnabled(zkMigrationEnabled).
+            setBrokerUncleanShutdownHandler(this::handleUncleanBrokerShutdown).
             build();
         this.producerIdControlManager = new ProducerIdControlManager.Builder().
             setLogContext(logContext).
@@ -2355,4 +2356,8 @@ void setNewNextWriteOffset(long newNextWriteOffset) {
             offsetControl.setNextWriteOffset(newNextWriteOffset);
         });
     }
+
+    void handleUncleanBrokerShutdown(int brokerId, List<ApiMessageAndVersion> records) {
+        replicationControl.handleBrokerUncleanShutdown(brokerId, records);
+    }
 }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java
index 759e3dfe5c4c5..28e767797ae61 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java
@@ -105,6 +105,7 @@
 
 import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -376,6 +377,12 @@ static Map<String, String> translateCreationConfigs(CreateableTopicConfigCollect
      */
     private final BrokersToIsrs brokersToIsrs;
 
+    /**
+     * A map of broker IDs to the partitions that the broker is in the ELR for.
+     * Note that, a broker should not be in both brokersToIsrs and brokersToElrs.
+     */
+    private final BrokersToElrs brokersToElrs;
+
     /**
      * A map from topic IDs to the partitions in the topic which are reassigning.
      */
@@ -424,6 +431,7 @@ private ReplicationControlManager(
         this.topicsWithCollisionChars = new TimelineHashMap<>(snapshotRegistry, 0);
         this.topics = new TimelineHashMap<>(snapshotRegistry, 0);
         this.brokersToIsrs = new BrokersToIsrs(snapshotRegistry);
+        this.brokersToElrs = new BrokersToElrs(snapshotRegistry);
         this.reassigningTopics = new TimelineHashMap<>(snapshotRegistry, 0);
         this.imbalancedPartitions = new TimelineHashSet<>(snapshotRegistry, 0);
         this.directoriesToPartitions = new TimelineHashMap<>(snapshotRegistry, 0);
@@ -471,8 +479,7 @@ public void replay(PartitionRecord record) {
             log.info("Replayed PartitionRecord for new partition {} and {}.", description,
                     newPartInfo);
             topicInfo.parts.put(record.partitionId(), newPartInfo);
-            brokersToIsrs.update(record.topicId(), record.partitionId(), null,
-                newPartInfo.isr, NO_LEADER, newPartInfo.leader);
+            updatePartitionInfo(record.topicId(), record.partitionId(), null, newPartInfo);
             updatePartitionDirectories(record.topicId(), record.partitionId(), null, newPartInfo.directories);
             updateReassigningTopicsIfNeeded(record.topicId(), record.partitionId(),
                     false,  isReassignmentInProgress(newPartInfo));
@@ -481,8 +488,7 @@ public void replay(PartitionRecord record) {
                     newPartInfo);
             newPartInfo.maybeLogPartitionChange(log, description, prevPartInfo);
             topicInfo.parts.put(record.partitionId(), newPartInfo);
-            brokersToIsrs.update(record.topicId(), record.partitionId(), prevPartInfo.isr,
-                newPartInfo.isr, prevPartInfo.leader, newPartInfo.leader);
+            updatePartitionInfo(record.topicId(), record.partitionId(), prevPartInfo, newPartInfo);
             updatePartitionDirectories(record.topicId(), record.partitionId(), prevPartInfo.directories, newPartInfo.directories);
             updateReassigningTopicsIfNeeded(record.topicId(), record.partitionId(),
                     isReassignmentInProgress(prevPartInfo), isReassignmentInProgress(newPartInfo));
@@ -528,9 +534,7 @@ public void replay(PartitionChangeRecord record) {
         updateReassigningTopicsIfNeeded(record.topicId(), record.partitionId(),
                 isReassignmentInProgress(prevPartitionInfo), isReassignmentInProgress(newPartitionInfo));
         topicInfo.parts.put(record.partitionId(), newPartitionInfo);
-        brokersToIsrs.update(record.topicId(), record.partitionId(),
-            prevPartitionInfo.isr, newPartitionInfo.isr, prevPartitionInfo.leader,
-            newPartitionInfo.leader);
+        updatePartitionInfo(record.topicId(), record.partitionId(), prevPartitionInfo, newPartitionInfo);
         updatePartitionDirectories(record.topicId(), record.partitionId(), prevPartitionInfo.directories, newPartitionInfo.directories);
         String topicPart = topicInfo.name + "-" + record.partitionId() + " with topic ID " +
             record.topicId();
@@ -582,6 +586,10 @@ public void replay(RemoveTopicRecord record) {
                 updatePartitionDirectories(topic.id, partitionId, partition.directories, null);
             }
 
+            for (int elrMember : partition.elr) {
+                brokersToElrs.removeTopicEntryForBroker(topic.id, elrMember);
+            }
+
             imbalancedPartitions.remove(new TopicIdPartition(record.topicId(), partitionId));
         }
         brokersToIsrs.removeTopicEntryForBroker(topic.id, NO_LEADER);
@@ -997,6 +1005,11 @@ BrokersToIsrs brokersToIsrs() {
         return brokersToIsrs;
     }
 
+    // VisibleForTesting
+    BrokersToElrs brokersToElrs() {
+        return brokersToElrs;
+    }
+
     // VisibleForTesting
     TimelineHashSet<TopicIdPartition> imbalancedPartitions() {
         return imbalancedPartitions;
@@ -1291,7 +1304,7 @@ void handleBrokerFenced(int brokerId, List<ApiMessageAndVersion> records) {
         if (brokerRegistration == null) {
             throw new RuntimeException("Can't find broker registration for broker " + brokerId);
         }
-        generateLeaderAndIsrUpdates("handleBrokerFenced", brokerId, NO_LEADER, records,
+        generateLeaderAndIsrUpdates("handleBrokerFenced", brokerId, NO_LEADER, NO_LEADER, records,
             brokersToIsrs.partitionsWithBrokerInIsr(brokerId));
         if (featureControl.metadataVersion().isBrokerRegistrationChangeRecordSupported()) {
             records.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
@@ -1308,7 +1321,7 @@ void handleBrokerFenced(int brokerId, List<ApiMessageAndVersion> records) {
     /**
      * Generate the appropriate records to handle a broker being unregistered.
      *
-     * First, we remove this broker from any ISR. Then we generate an
+     * First, we remove this broker from any ISR or ELR. Then we generate an
      * UnregisterBrokerRecord.
      *
      * @param brokerId      The broker id.
@@ -1317,8 +1330,10 @@ void handleBrokerFenced(int brokerId, List<ApiMessageAndVersion> records) {
      */
     void handleBrokerUnregistered(int brokerId, long brokerEpoch,
                                   List<ApiMessageAndVersion> records) {
-        generateLeaderAndIsrUpdates("handleBrokerUnregistered", brokerId, NO_LEADER, records,
+        generateLeaderAndIsrUpdates("handleBrokerUnregistered", brokerId, NO_LEADER, NO_LEADER, records,
             brokersToIsrs.partitionsWithBrokerInIsr(brokerId));
+        generateLeaderAndIsrUpdates("handleBrokerUnregistered", brokerId, NO_LEADER, NO_LEADER, records,
+            brokersToElrs.partitionsWithBrokerInElr(brokerId));
         records.add(new ApiMessageAndVersion(new UnregisterBrokerRecord().
             setBrokerId(brokerId).setBrokerEpoch(brokerEpoch),
             (short) 0));
@@ -1345,7 +1360,7 @@ void handleBrokerUnfenced(int brokerId, long brokerEpoch, List<ApiMessageAndVers
             records.add(new ApiMessageAndVersion(new UnfenceBrokerRecord().setId(brokerId).
                 setEpoch(brokerEpoch), (short) 0));
         }
-        generateLeaderAndIsrUpdates("handleBrokerUnfenced", NO_LEADER, brokerId, records,
+        generateLeaderAndIsrUpdates("handleBrokerUnfenced", NO_LEADER, brokerId, NO_LEADER, records,
             brokersToIsrs.partitionsWithNoLeader());
     }
 
@@ -1369,7 +1384,21 @@ void handleBrokerInControlledShutdown(int brokerId, long brokerEpoch, List<ApiMe
                 (short) 1));
         }
         generateLeaderAndIsrUpdates("enterControlledShutdown[" + brokerId + "]",
-            brokerId, NO_LEADER, records, brokersToIsrs.partitionsWithBrokerInIsr(brokerId));
+            brokerId, NO_LEADER, NO_LEADER, records, brokersToIsrs.partitionsWithBrokerInIsr(brokerId));
+    }
+
+    /**
+     * Create partition change records to remove replicas from any ISR or ELR for brokers doing unclean shutdown.
+     *
+     * @param brokerId      The broker id.
+     * @param records       The record list to append to.
+     */
+    void handleBrokerUncleanShutdown(int brokerId, List<ApiMessageAndVersion> records) {
+        if (!featureControl.metadataVersion().isElrSupported()) return;
+        generateLeaderAndIsrUpdates("handleBrokerUncleanShutdown", NO_LEADER, NO_LEADER, brokerId, records,
+            brokersToIsrs.partitionsWithBrokerInIsr(brokerId));
+        generateLeaderAndIsrUpdates("handleBrokerUncleanShutdown", NO_LEADER, NO_LEADER, brokerId, records,
+            brokersToElrs.partitionsWithBrokerInElr(brokerId));
     }
 
     /**
@@ -1399,7 +1428,7 @@ void handleDirectoriesOffline(
                         Collections.emptyIterator() : parts.iterator();
                 generateLeaderAndIsrUpdates(
                         "handleDirectoriesOffline[" + brokerId + ":" + newOfflineDir + "]",
-                        brokerId, NO_LEADER, records, iterator);
+                        brokerId, NO_LEADER, NO_LEADER, records, iterator);
             }
             List<Uuid> newOnlineDirs = registration.directoryDifference(offlineDirs);
             records.add(new ApiMessageAndVersion(new BrokerRegistrationChangeRecord().
@@ -1807,7 +1836,7 @@ void validateManualPartitionAssignment(
     }
 
     /**
-     * Iterate over a sequence of partitions and generate ISR changes and/or leader
+     * Iterate over a sequence of partitions and generate ISR/ELR changes and/or leader
      * changes if necessary.
      *
      * @param context           A human-readable context string used in log4j logging.
@@ -1815,12 +1844,17 @@ void validateManualPartitionAssignment(
      *                          broker to remove from the ISR and leadership, otherwise.
      * @param brokerToAdd       NO_LEADER if no broker is being added; the ID of the
      *                          broker which is now eligible to be a leader, otherwise.
+     * @param brokerWithUncleanShutdown
+     *                          NO_LEADER if no broker has unclean shutdown; the ID of the
+     *                          broker which is now removed from the ISR, ELR and
+     *                          leadership, otherwise.
      * @param records           A list of records which we will append to.
      * @param iterator          The iterator containing the partitions to examine.
      */
     void generateLeaderAndIsrUpdates(String context,
                                      int brokerToRemove,
                                      int brokerToAdd,
+                                     int brokerWithUncleanShutdown,
                                      List<ApiMessageAndVersion> records,
                                      Iterator<TopicIdPartition> iterator) {
         int oldSize = records.size();
@@ -1837,8 +1871,13 @@ void generateLeaderAndIsrUpdates(String context,
         // from the target ISR, but we need to exclude it here too, to handle the case
         // where there is an unclean leader election which chooses a leader from outside
         // the ISR.
+        //
+        // If the caller passed a valid broker ID for brokerWithUncleanShutdown, rather than
+        // passing NO_LEADER, this node should not be an acceptable leader. We also exclude
+        // brokerWithUncleanShutdown from ELR and ISR.
         IntPredicate isAcceptableLeader =
-            r -> (r != brokerToRemove) && (r == brokerToAdd || clusterControl.isActive(r));
+            r -> (r != brokerToRemove && r != brokerWithUncleanShutdown)
+                && (r == brokerToAdd || clusterControl.isActive(r));
 
         while (iterator.hasNext()) {
             TopicIdPartition topicIdPart = iterator.next();
@@ -1865,11 +1904,14 @@ void generateLeaderAndIsrUpdates(String context,
             if (configurationControl.uncleanLeaderElectionEnabledForTopic(topic.name)) {
                 builder.setElection(PartitionChangeBuilder.Election.UNCLEAN);
             }
+            if (brokerWithUncleanShutdown != NO_LEADER) {
+                builder.setUncleanShutdownReplicas(Arrays.asList(brokerWithUncleanShutdown));
+            }
 
-            // Note: if brokerToRemove was passed as NO_LEADER, this is a no-op (the new
+            // Note: if brokerToRemove and brokerWithUncleanShutdown were passed as NO_LEADER, this is a no-op (the new
             // target ISR will be the same as the old one).
             builder.setTargetIsr(Replicas.toList(
-                Replicas.copyWithout(partition.isr, brokerToRemove)));
+                Replicas.copyWithout(partition.isr, new int[] {brokerToRemove, brokerWithUncleanShutdown})));
 
             builder.setDefaultDirProvider(clusterDescriber)
                     .build().ifPresent(records::add);
@@ -2145,7 +2187,7 @@ ControllerResult<AssignReplicasToDirsResponseData> handleAssignReplicasToDirs(As
             response.directories().add(resDir);
         }
         if (!leaderAndIsrUpdates.isEmpty()) {
-            generateLeaderAndIsrUpdates("offline-dir-assignment", brokerId, NO_LEADER, records, leaderAndIsrUpdates.iterator());
+            generateLeaderAndIsrUpdates("offline-dir-assignment", brokerId, NO_LEADER, NO_LEADER, records, leaderAndIsrUpdates.iterator());
         }
         return ControllerResult.of(records, response);
     }
@@ -2240,6 +2282,25 @@ private void updatePartitionDirectories(
         }
     }
 
+    private void updatePartitionInfo(
+        Uuid topicId,
+        Integer partitionId,
+        PartitionRegistration prevPartInfo,
+        PartitionRegistration newPartInfo
+    ) {
+        HashSet<Integer> validationSet = new HashSet<>();
+        Arrays.stream(newPartInfo.isr).forEach(validationSet::add);
+        Arrays.stream(newPartInfo.elr).forEach(validationSet::add);
+        if (validationSet.size() != newPartInfo.isr.length + newPartInfo.elr.length) {
+            log.error("{}-{} has overlapping ISR={} and ELR={}", topics.get(topicId).name, partitionId,
+                Arrays.toString(newPartInfo.isr), partitionId, Arrays.toString(newPartInfo.elr));
+        }
+        brokersToIsrs.update(topicId, partitionId, prevPartInfo == null ? null : prevPartInfo.isr,
+            newPartInfo.isr, prevPartInfo == null ? NO_LEADER : prevPartInfo.leader, newPartInfo.leader);
+        brokersToElrs.update(topicId, partitionId, prevPartInfo == null ? null : prevPartInfo.elr,
+            newPartInfo.elr);
+    }
+
     private static final class IneligibleReplica {
         private final int replicaId;
         private final String reason;
diff --git a/metadata/src/test/java/org/apache/kafka/controller/BrokerToElrsTest.java b/metadata/src/test/java/org/apache/kafka/controller/BrokerToElrsTest.java
new file mode 100644
index 0000000000000..890a298531275
--- /dev/null
+++ b/metadata/src/test/java/org/apache/kafka/controller/BrokerToElrsTest.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.controller;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.server.common.TopicIdPartition;
+import org.apache.kafka.timeline.SnapshotRegistry;
+import org.junit.jupiter.api.Test;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class BrokerToElrsTest {
+    private static final Uuid[] UUIDS = new Uuid[] {
+            Uuid.fromString("z5XgH_fQSAK3-RYoF2ymgw"),
+            Uuid.fromString("U52uRe20RsGI0RvpcTx33Q")
+    };
+
+    private static Set<TopicIdPartition> toSet(TopicIdPartition... partitions) {
+        HashSet<TopicIdPartition> set = new HashSet<>();
+        for (TopicIdPartition partition : partitions) {
+            set.add(partition);
+        }
+        return set;
+    }
+
+    private static Set<TopicIdPartition> toSet(BrokersToIsrs.PartitionsOnReplicaIterator iterator) {
+        HashSet<TopicIdPartition> set = new HashSet<>();
+        while (iterator.hasNext()) {
+            set.add(iterator.next());
+        }
+        return set;
+    }
+
+    @Test
+    public void testIterator() {
+        SnapshotRegistry snapshotRegistry = new SnapshotRegistry(new LogContext());
+        BrokersToElrs brokersToElrs = new BrokersToElrs(snapshotRegistry);
+        assertEquals(toSet(), toSet(brokersToElrs.partitionsWithBrokerInElr(1)));
+        brokersToElrs.update(UUIDS[0], 0, null, new int[] {1, 2, 3});
+        brokersToElrs.update(UUIDS[1], 1, null, new int[] {2, 3, 4});
+        assertEquals(toSet(new TopicIdPartition(UUIDS[0], 0)),
+            toSet(brokersToElrs.partitionsWithBrokerInElr(1)));
+        assertEquals(toSet(new TopicIdPartition(UUIDS[0], 0),
+                new TopicIdPartition(UUIDS[1], 1)),
+            toSet(brokersToElrs.partitionsWithBrokerInElr(2)));
+        assertEquals(toSet(new TopicIdPartition(UUIDS[1], 1)),
+            toSet(brokersToElrs.partitionsWithBrokerInElr(4)));
+        assertEquals(toSet(), toSet(brokersToElrs.partitionsWithBrokerInElr(5)));
+        brokersToElrs.update(UUIDS[1], 2, null, new int[] {3, 2, 1});
+        assertEquals(toSet(new TopicIdPartition(UUIDS[0], 0),
+                new TopicIdPartition(UUIDS[1], 1),
+                new TopicIdPartition(UUIDS[1], 2)),
+            toSet(brokersToElrs.partitionsWithBrokerInElr(2)));
+    }
+}
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java
index cc3a7aa3935be..20c2b7c690997 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java
@@ -95,6 +95,7 @@ public void testReplay(MetadataVersion metadataVersion) {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         assertFalse(clusterControl.isUnfenced(0));
@@ -156,6 +157,7 @@ public void testReplayRegisterBrokerRecord() {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
 
         assertFalse(clusterControl.isUnfenced(0));
@@ -208,6 +210,7 @@ public void testReplayBrokerRegistrationChangeRecord() {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
 
         assertFalse(clusterControl.isUnfenced(0));
@@ -262,6 +265,7 @@ public void testRegistrationWithIncorrectClusterId() {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         assertThrows(InconsistentClusterIdException.class, () ->
@@ -301,6 +305,7 @@ public void testRegisterBrokerRecordVersion(MetadataVersion metadataVersion) {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
 
@@ -362,6 +367,7 @@ public void testUnregister() {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         clusterControl.replay(brokerRecord, 100L);
@@ -400,6 +406,7 @@ public void testPlaceReplicas(int numUsableBrokers) {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         for (int i = 0; i < numUsableBrokers; i++) {
@@ -463,6 +470,7 @@ public void testRegistrationsToRecords(MetadataVersion metadataVersion) {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         assertFalse(clusterControl.isUnfenced(0));
@@ -541,6 +549,7 @@ public void testRegistrationWithUnsupportedMetadataVersion() {
                 setTime(new MockTime(0, 0, 0)).
                 setSnapshotRegistry(snapshotRegistry).
                 setFeatureControlManager(featureControl).
+                setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
                 build();
         clusterControl.activate();
 
@@ -584,6 +593,7 @@ public void testRegisterControlWithOlderMetadataVersion() {
         ClusterControlManager clusterControl = new ClusterControlManager.Builder().
             setClusterId("fPZv1VBsRFmnlRvmGcOW9w").
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
         clusterControl.activate();
         assertEquals("The current MetadataVersion is too old to support controller registrations.",
@@ -596,6 +606,7 @@ public void testRegisterWithDuplicateDirectoryId() {
         ClusterControlManager clusterControl = new ClusterControlManager.Builder().
                 setClusterId("QzZZEtC7SxucRM29Xdzijw").
                 setFeatureControlManager(new FeatureControlManager.Builder().build()).
+                setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
                 build();
         RegisterBrokerRecord brokerRecord = new RegisterBrokerRecord().setBrokerEpoch(100).setBrokerId(0).setLogDirs(asList(
                 Uuid.fromString("yJGxmjfbQZSVFAlNM3uXZg"),
@@ -637,6 +648,7 @@ public void testHasOnlineDir() {
         ClusterControlManager clusterControl = new ClusterControlManager.Builder().
                 setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                 setFeatureControlManager(new FeatureControlManager.Builder().build()).
+                setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
                 build();
         clusterControl.activate();
         registerNewBrokerWithDirs(clusterControl, 1, asList(Uuid.fromString("dir1SEbpRuG1dcpTRGOvJw"), Uuid.fromString("dir2xaEwR2m3JHTiy7PWwA")));
@@ -655,6 +667,7 @@ public void testDefaultDir() {
         ClusterControlManager clusterControl = new ClusterControlManager.Builder().
                 setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                 setFeatureControlManager(new FeatureControlManager.Builder().build()).
+                setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
                 build();
         clusterControl.activate();
         RegisterBrokerRecord brokerRecord = new RegisterBrokerRecord().setBrokerEpoch(100).setBrokerId(1).setLogDirs(Collections.emptyList());
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
index d8c3770f85af0..a58f194b1f293 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ProducerIdControlManagerTest.java
@@ -56,6 +56,7 @@ public void setUp() {
             setSnapshotRegistry(snapshotRegistry).
             setSessionTimeoutNs(1000).
             setFeatureControlManager(featureControl).
+            setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
             build();
 
         clusterControl.activate();
diff --git a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
index e907514e05b90..1b18c9648de4f 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTest.java
@@ -56,6 +56,7 @@
 import org.apache.kafka.common.metadata.RegisterControllerRecord;
 import org.apache.kafka.common.metadata.UnfenceBrokerRecord;
 import org.apache.kafka.common.metadata.ZkMigrationStateRecord;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.requests.AlterPartitionRequest;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
@@ -147,6 +148,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static java.util.Collections.singletonList;
 import static java.util.function.Function.identity;
 import static org.apache.kafka.clients.admin.AlterConfigOp.OpType.SET;
 import static org.apache.kafka.common.config.ConfigResource.Type.BROKER;
@@ -155,11 +157,13 @@
 import static org.apache.kafka.controller.ConfigurationControlManagerTest.SCHEMA;
 import static org.apache.kafka.controller.ConfigurationControlManagerTest.entry;
 import static org.apache.kafka.controller.ControllerRequestContextUtil.ANONYMOUS_CONTEXT;
+import static org.apache.kafka.controller.ControllerRequestContextUtil.anonymousContextFor;
 import static org.apache.kafka.controller.QuorumControllerIntegrationTestUtils.brokerFeatures;
 import static org.apache.kafka.controller.QuorumControllerIntegrationTestUtils.forceRenounce;
 import static org.apache.kafka.controller.QuorumControllerIntegrationTestUtils.pause;
 import static org.apache.kafka.controller.QuorumControllerIntegrationTestUtils.registerBrokersAndUnfence;
 import static org.apache.kafka.controller.QuorumControllerIntegrationTestUtils.sendBrokerHeartbeatToUnfenceBrokers;
+import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -361,6 +365,140 @@ public void testFenceMultipleBrokers() throws Throwable {
         }
     }
 
+    @Test
+    public void testUncleanShutdownBroker() throws Throwable {
+        List<Integer> allBrokers = Arrays.asList(1, 2, 3);
+        short replicationFactor = (short) allBrokers.size();
+        long sessionTimeoutMillis = 500;
+
+        try (
+            LocalLogManagerTestEnv logEnv = new LocalLogManagerTestEnv.Builder(1).
+                build();
+            QuorumControllerTestEnv controlEnv = new QuorumControllerTestEnv.Builder(logEnv).
+                setControllerBuilderInitializer(controllerBuilder -> {
+                    controllerBuilder.setConfigSchema(SCHEMA);
+                }).
+                setSessionTimeoutMillis(OptionalLong.of(sessionTimeoutMillis)).
+
+                setBootstrapMetadata(BootstrapMetadata.fromVersion(MetadataVersion.IBP_3_8_IV0, "test-provided bootstrap ELR enabled")).
+                build()
+        ) {
+            ListenerCollection listeners = new ListenerCollection();
+            listeners.add(new Listener().setName("PLAINTEXT").setHost("localhost").setPort(9092));
+            QuorumController active = controlEnv.activeController();
+            Map<Integer, Long> brokerEpochs = new HashMap<>();
+
+            for (Integer brokerId : allBrokers) {
+                CompletableFuture<BrokerRegistrationReply> reply = active.registerBroker(
+                    anonymousContextFor(ApiKeys.BROKER_REGISTRATION),
+                    new BrokerRegistrationRequestData().
+                        setBrokerId(brokerId).
+                        setClusterId(active.clusterId()).
+                        setFeatures(brokerFeatures(MetadataVersion.IBP_3_0_IV1, MetadataVersion.IBP_3_8_IV0)).
+                        setIncarnationId(Uuid.randomUuid()).
+                        setLogDirs(Collections.singletonList(Uuid.randomUuid())).
+                        setListeners(listeners));
+                brokerEpochs.put(brokerId, reply.get().epoch());
+            }
+
+            // Brokers are only registered and should still be fenced
+            allBrokers.forEach(brokerId -> {
+                assertFalse(active.clusterControl().isUnfenced(brokerId),
+                    "Broker " + brokerId + " should have been fenced");
+            });
+
+            // Unfence all brokers and create a topic foo
+            sendBrokerHeartbeatToUnfenceBrokers(active, allBrokers, brokerEpochs);
+            CreateTopicsRequestData createTopicsRequestData = new CreateTopicsRequestData().setTopics(
+                new CreatableTopicCollection(Collections.singleton(
+                    new CreatableTopic().setName("foo").setNumPartitions(1).
+                        setReplicationFactor(replicationFactor)).iterator()));
+            CreateTopicsResponseData createTopicsResponseData = active.createTopics(
+                ANONYMOUS_CONTEXT, createTopicsRequestData,
+                Collections.singleton("foo")).get();
+            assertEquals(Errors.NONE, Errors.forCode(createTopicsResponseData.topics().find("foo").errorCode()));
+            Uuid topicIdFoo = createTopicsResponseData.topics().find("foo").topicId();
+            ConfigRecord configRecord = new ConfigRecord()
+                .setResourceType(ConfigResource.Type.TOPIC.id())
+                .setResourceName("foo")
+                .setName(org.apache.kafka.common.config.TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG)
+                .setValue("2");
+            RecordTestUtils.replayAll(active.configurationControl(), singletonList(new ApiMessageAndVersion(configRecord, (short) 0)));
+
+            // Fence all the brokers
+            TestUtils.waitForCondition(() -> {
+                    for (Integer brokerId : allBrokers) {
+                        if (active.clusterControl().isUnfenced(brokerId)) {
+                            return false;
+                        }
+                    }
+                    return true;
+                }, sessionTimeoutMillis * 30,
+                "Fencing of brokers did not process within expected time"
+            );
+
+            // Verify the isr and elr for the topic partition
+            PartitionRegistration partition = active.replicationControl().getPartition(topicIdFoo, 0);
+            assertEquals(1, partition.lastKnownElr.length, partition.toString());
+            int[] lastKnownElr = partition.lastKnownElr;
+            assertEquals(0, partition.isr.length, partition.toString());
+            assertEquals(NO_LEADER, partition.leader, partition.toString());
+
+            // The ELR set is not determined.
+            assertEquals(2, partition.elr.length, partition.toString());
+            int brokerToUncleanShutdown, brokerToBeTheLeader;
+
+            // lastKnownElr stores the last known leader.
+            if (lastKnownElr[0] == partition.elr[0]) {
+                brokerToUncleanShutdown = partition.elr[0];
+                brokerToBeTheLeader = partition.elr[1];
+            } else {
+                brokerToUncleanShutdown = partition.elr[1];
+                brokerToBeTheLeader = partition.elr[0];
+            }
+
+            // Unclean shutdown should remove the ELR members.
+            active.registerBroker(
+                anonymousContextFor(ApiKeys.BROKER_REGISTRATION),
+                new BrokerRegistrationRequestData().
+                    setBrokerId(brokerToUncleanShutdown).
+                    setClusterId(active.clusterId()).
+                    setFeatures(brokerFeatures(MetadataVersion.IBP_3_0_IV1, MetadataVersion.IBP_3_8_IV0)).
+                    setIncarnationId(Uuid.randomUuid()).
+                    setLogDirs(Collections.singletonList(Uuid.randomUuid())).
+                    setListeners(listeners)).get();
+            partition = active.replicationControl().getPartition(topicIdFoo, 0);
+            assertArrayEquals(new int[]{brokerToBeTheLeader}, partition.elr, partition.toString());
+
+            // Unclean shutdown should not remove the last known ELR members.
+            active.registerBroker(
+                anonymousContextFor(ApiKeys.BROKER_REGISTRATION),
+                new BrokerRegistrationRequestData().
+                    setBrokerId(lastKnownElr[0]).
+                    setClusterId(active.clusterId()).
+                    setFeatures(brokerFeatures(MetadataVersion.IBP_3_0_IV1, MetadataVersion.IBP_3_8_IV0)).
+                    setIncarnationId(Uuid.randomUuid()).
+                    setLogDirs(Collections.singletonList(Uuid.randomUuid())).
+                    setListeners(listeners)).get();
+            partition = active.replicationControl().getPartition(topicIdFoo, 0);
+            assertArrayEquals(lastKnownElr, partition.lastKnownElr, partition.toString());
+
+            // Unfence the last one in the ELR, it should be elected.
+            sendBrokerHeartbeatToUnfenceBrokers(active, Arrays.asList(brokerToBeTheLeader), brokerEpochs);
+            TestUtils.waitForCondition(() -> {
+                    return active.clusterControl().isUnfenced(brokerToBeTheLeader);
+                }, sessionTimeoutMillis * 3,
+                "Broker should be unfenced."
+            );
+
+            partition = active.replicationControl().getPartition(topicIdFoo, 0);
+            assertArrayEquals(new int[]{brokerToBeTheLeader}, partition.isr, partition.toString());
+            assertEquals(0, partition.elr.length, partition.toString());
+            assertEquals(0, partition.lastKnownElr.length, partition.toString());
+            assertEquals(brokerToBeTheLeader, partition.leader, partition.toString());
+        }
+    }
+
     @Test
     public void testBalancePartitionLeaders() throws Throwable {
         List<Integer> allBrokers = Arrays.asList(1, 2, 3);
diff --git a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTestEnv.java b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTestEnv.java
index f340d27925bb7..885173638d05b 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTestEnv.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/QuorumControllerTestEnv.java
@@ -49,6 +49,7 @@ public static class Builder {
         private Consumer<QuorumController.Builder> controllerBuilderInitializer = __ -> { };
         private OptionalLong sessionTimeoutMillis = OptionalLong.empty();
         private OptionalLong leaderImbalanceCheckIntervalNs = OptionalLong.empty();
+        private boolean eligibleLeaderReplicasEnabled = false;
         private BootstrapMetadata bootstrapMetadata = BootstrapMetadata.
                 fromVersion(MetadataVersion.latestTesting(), "test-provided version");
 
@@ -82,6 +83,7 @@ public QuorumControllerTestEnv build() throws Exception {
                 controllerBuilderInitializer,
                 sessionTimeoutMillis,
                 leaderImbalanceCheckIntervalNs,
+                bootstrapMetadata.metadataVersion().isElrSupported(),
                 bootstrapMetadata);
         }
     }
@@ -91,6 +93,7 @@ private QuorumControllerTestEnv(
         Consumer<QuorumController.Builder> controllerBuilderInitializer,
         OptionalLong sessionTimeoutMillis,
         OptionalLong leaderImbalanceCheckIntervalNs,
+        boolean eligibleLeaderReplicasEnabled,
         BootstrapMetadata bootstrapMetadata
     ) throws Exception {
         this.logEnv = logEnv;
@@ -112,6 +115,7 @@ private QuorumControllerTestEnv(
                 fatalFaultHandlers.put(nodeId, fatalFaultHandler);
                 MockFaultHandler nonFatalFaultHandler = new MockFaultHandler("nonFatalFaultHandler");
                 builder.setNonFatalFaultHandler(nonFatalFaultHandler);
+                builder.setEligibleLeaderReplicasEnabled(eligibleLeaderReplicasEnabled);
                 nonFatalFaultHandlers.put(nodeId, fatalFaultHandler);
                 controllerBuilderInitializer.accept(builder);
                 this.controllers.add(builder.build());
diff --git a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
index 7e04818627050..f68cf459dea88 100644
--- a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
+++ b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java
@@ -244,6 +244,7 @@ private ReplicationControlTestContext(
                 setSessionTimeoutNs(TimeUnit.MILLISECONDS.convert(BROKER_SESSION_TIMEOUT_MS, TimeUnit.NANOSECONDS)).
                 setReplicaPlacer(new StripedReplicaPlacer(random)).
                 setFeatureControlManager(featureControl).
+                setBrokerUncleanShutdownHandler(this::handleUncleanBrokerShutdown).
                 build();
 
             this.replicationControl = new ReplicationControlManager.Builder().
@@ -259,6 +260,10 @@ private ReplicationControlTestContext(
             clusterControl.activate();
         }
 
+        void handleUncleanBrokerShutdown(int brokerId, List<ApiMessageAndVersion> records) {
+            replicationControl.handleBrokerUncleanShutdown(brokerId, records);
+        }
+
         CreatableTopicResult createTestTopic(String name,
                                              int numPartitions,
                                              short replicationFactor,
@@ -383,6 +388,14 @@ void registerBrokersWithDirs(Object... brokerIdsAndDirs) throws Exception {
             }
         }
 
+        void handleBrokersUncleanShutdown(Integer... brokerIds) throws Exception {
+            List<ApiMessageAndVersion> records = new ArrayList<>();
+            for (int brokerId : brokerIds) {
+                replicationControl.handleBrokerUncleanShutdown(brokerId, records);
+            }
+            replay(records);
+        }
+
         void alterPartition(
             TopicIdPartition topicIdPartition,
             int leaderId,
@@ -1013,6 +1026,42 @@ public void testEligibleLeaderReplicas_BrokerFence() throws Exception {
         assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
     }
 
+    @Test
+    public void testEligibleLeaderReplicas_DeleteTopic() throws Exception {
+        ReplicationControlTestContext ctx = new ReplicationControlTestContext.Builder().setIsElrEnabled(true).build();
+        ReplicationControlManager replicationControl = ctx.replicationControl;
+        ctx.registerBrokers(0, 1, 2);
+        ctx.unfenceBrokers(0, 1, 2);
+        CreatableTopicResult createTopicResult = ctx.createTestTopic("foo",
+            new int[][] {new int[] {0, 1, 2}});
+
+        TopicIdPartition topicIdPartition = new TopicIdPartition(createTopicResult.topicId(), 0);
+        assertEquals(OptionalInt.of(0), ctx.currentLeader(topicIdPartition));
+        long brokerEpoch = ctx.currentBrokerEpoch(0);
+        ctx.alterTopicConfig("foo", TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "2");
+
+        // Change ISR to {0}.
+        PartitionData shrinkIsrRequest = newAlterPartition(
+            replicationControl, topicIdPartition, isrWithDefaultEpoch(0), LeaderRecoveryState.RECOVERED);
+
+        ControllerResult<AlterPartitionResponseData> shrinkIsrResult = sendAlterPartition(
+            replicationControl, 0, brokerEpoch, topicIdPartition.topicId(), shrinkIsrRequest);
+        AlterPartitionResponseData.PartitionData shrinkIsrResponse = assertAlterPartitionResponse(
+            shrinkIsrResult, topicIdPartition, NONE);
+        assertConsistentAlterPartitionResponse(replicationControl, topicIdPartition, shrinkIsrResponse);
+        PartitionRegistration partition = replicationControl.getPartition(topicIdPartition.topicId(),
+            topicIdPartition.partitionId());
+        assertTrue(Arrays.equals(new int[]{1, 2}, partition.elr), partition.toString());
+        assertTrue(Arrays.equals(new int[]{}, partition.lastKnownElr), partition.toString());
+        assertTrue(replicationControl.brokersToElrs().partitionsWithBrokerInElr(1).hasNext());
+
+        ControllerRequestContext deleteTopicsRequestContext = anonymousContextFor(ApiKeys.DELETE_TOPICS);
+        ctx.deleteTopic(deleteTopicsRequestContext, createTopicResult.topicId());
+
+        assertFalse(replicationControl.brokersToElrs().partitionsWithBrokerInElr(1).hasNext());
+        assertFalse(replicationControl.brokersToIsrs().partitionsWithBrokerInIsr(0).hasNext());
+    }
+
     @Test
     public void testEligibleLeaderReplicas_EffectiveMinIsr() throws Exception {
         ReplicationControlTestContext ctx = new ReplicationControlTestContext.Builder().setIsElrEnabled(true).build();
@@ -1058,6 +1107,40 @@ public void testEligibleLeaderReplicas_CleanElection() throws Exception {
         assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
     }
 
+    @Test
+    public void testEligibleLeaderReplicas_UncleanShutdown() throws Exception {
+        ReplicationControlTestContext ctx = new ReplicationControlTestContext.Builder()
+            .setIsElrEnabled(true)
+            .build();
+        ReplicationControlManager replicationControl = ctx.replicationControl;
+        ctx.registerBrokers(0, 1, 2, 3);
+        ctx.unfenceBrokers(0, 1, 2, 3);
+        CreatableTopicResult createTopicResult = ctx.createTestTopic("foo",
+                new int[][] {new int[] {0, 1, 2, 3}});
+
+        TopicIdPartition topicIdPartition = new TopicIdPartition(createTopicResult.topicId(), 0);
+        assertEquals(OptionalInt.of(0), ctx.currentLeader(topicIdPartition));
+        ctx.alterTopicConfig("foo", TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "3");
+
+        ctx.fenceBrokers(Utils.mkSet(1, 2, 3));
+
+        PartitionRegistration partition = replicationControl.getPartition(topicIdPartition.topicId(), topicIdPartition.partitionId());
+        assertTrue(Arrays.equals(new int[]{2, 3}, partition.elr), partition.toString());
+        assertTrue(Arrays.equals(new int[]{}, partition.lastKnownElr), partition.toString());
+
+        // An unclean shutdown ELR member should be kicked out of ELR.
+        ctx.handleBrokersUncleanShutdown(3);
+        partition = replicationControl.getPartition(topicIdPartition.topicId(), topicIdPartition.partitionId());
+        assertTrue(Arrays.equals(new int[]{2}, partition.elr), partition.toString());
+        assertTrue(Arrays.equals(new int[]{}, partition.lastKnownElr), partition.toString());
+
+        // An unclean shutdown last ISR member should be recognized as the last known leader.
+        ctx.handleBrokersUncleanShutdown(0);
+        partition = replicationControl.getPartition(topicIdPartition.topicId(), topicIdPartition.partitionId());
+        assertTrue(Arrays.equals(new int[]{2}, partition.elr), partition.toString());
+        assertTrue(Arrays.equals(new int[]{0}, partition.lastKnownElr), partition.toString());
+    }
+
     @ParameterizedTest
     @ApiKeyVersionsSource(apiKey = ApiKeys.ALTER_PARTITION)
     public void testAlterPartitionHandleUnknownTopicIdOrName(short version) throws Exception {

From e8624ff128701d5b47a43062111dd3c8200bbb48 Mon Sep 17 00:00:00 2001
From: ShivsundarR <shr@confluent.io>
Date: Fri, 5 Apr 2024 16:29:17 +0530
Subject: [PATCH 253/258] Minor changes according to latest trunk

---
 .../clients/consumer/internals/ShareConsumerImpl.java     | 8 ++++----
 .../consumer/internals/ShareHeartbeatRequestManager.java  | 4 ++--
 .../consumer/internals/ShareMembershipManager.java        | 4 ++--
 .../internals/events/ApplicationEventProcessor.java       | 4 ++--
 .../events/ShareUnsubscribeApplicationEvent.java          | 2 +-
 .../clients/consumer/internals/ShareConsumerImplTest.java | 8 ++++----
 .../kafka/coordinator/group/GroupMetadataManager.java     | 8 +++-----
 7 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java
index 68bdcfa71559f..e596cf9197fe3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java
@@ -33,7 +33,7 @@
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.PollEvent;
 import org.apache.kafka.clients.consumer.internals.events.ShareLeaveOnCloseApplicationEvent;
@@ -121,7 +121,7 @@ public BackgroundEventProcessor(final LogContext logContext,
 
         /**
          * Process the events, if any, that were produced by the {@link ConsumerNetworkThread network thread}.
-         * It is possible that {@link org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent an error}
+         * It is possible that {@link org.apache.kafka.clients.consumer.internals.events.ErrorEvent an error}
          * could occur when processing the events. In such cases, the processor will take a reference to the first
          * error, continue to process the remaining events, and then throw the first error that occurred.
          */
@@ -150,13 +150,13 @@ public boolean process() {
         @Override
         public void process(final BackgroundEvent event) {
             if (event.type() == BackgroundEvent.Type.ERROR) {
-                process((ErrorBackgroundEvent) event);
+                process((ErrorEvent) event);
             } else {
                 throw new IllegalArgumentException("Background event type " + event.type() + " was not expected");
             }
         }
 
-        private void process(final ErrorBackgroundEvent event) {
+        private void process(final ErrorEvent event) {
             throw event.error();
         }
     }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java
index a7c1978a2cf03..35ed07d1cd559 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java
@@ -21,7 +21,7 @@
 import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.metrics.HeartbeatMetricsManager;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.RetriableException;
@@ -402,7 +402,7 @@ private void logInfo(final String message,
     }
 
     private void handleFatalFailure(Throwable error) {
-        backgroundEventHandler.add(new ErrorBackgroundEvent(error));
+        backgroundEventHandler.add(new ErrorEvent(error));
         shareMembershipManager.transitionToFatal();
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java
index e965491c21f4e..c05eee4893a28 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java
@@ -23,7 +23,7 @@
 import org.apache.kafka.clients.consumer.internals.Utils.TopicIdPartitionComparator;
 import org.apache.kafka.clients.consumer.internals.Utils.TopicPartitionComparator;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.Uuid;
@@ -208,7 +208,7 @@ public class ShareMembershipManager implements RequestManager {
 
     /**
      * Serves as the conduit by which we can report events to the application thread. This is needed as we send
-     * {@link ErrorBackgroundEvent errors} to the application thread.
+     * {@link ErrorEvent errors} to the application thread.
      */
     private final BackgroundEventHandler backgroundEventHandler;
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
index 0aececac971c8..a27c56da99bab 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
@@ -336,7 +336,7 @@ private void process(final ShareUnsubscribeApplicationEvent event) {
         ShareMembershipManager membershipManager = requestManagers.shareHeartbeatRequestManager.get().membershipManager();
         CompletableFuture<Void> future = membershipManager.leaveGroup();
         // The future will be completed on heartbeat sent
-        event.chain(future);
+        future.whenComplete(complete(event.future()));
     }
 
     private void process(final ShareLeaveOnCloseApplicationEvent event) {
@@ -352,7 +352,7 @@ private void process(final ShareLeaveOnCloseApplicationEvent event) {
         log.debug("Leaving group before closing");
         CompletableFuture<Void> future = membershipManager.leaveGroup();
         // The future will be completed on heartbeat sent
-        event.chain(future);
+        future.whenComplete(complete(event.future()));
     }
 
     private <T> BiConsumer<? super T, ? super Throwable> complete(final CompletableFuture<T> b) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
index cb1f3cf7c1460..4676d70e717fc 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
@@ -28,6 +28,6 @@
  */
 public class ShareUnsubscribeApplicationEvent extends CompletableApplicationEvent<Void> {
     public ShareUnsubscribeApplicationEvent() {
-        super(Type.SHARE_UNSUBSCRIBE, new Timer());
+        super(Type.SHARE_UNSUBSCRIBE, Integer.MAX_VALUE);
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java
index 9a20be27cf94b..c550b5be6f5c0 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java
@@ -19,7 +19,7 @@
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
-import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.ShareLeaveOnCloseApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ShareSubscriptionChangeApplicationEvent;
@@ -225,7 +225,7 @@ public void testBackgroundError() {
         consumer = newConsumer(config);
 
         final KafkaException expectedException = new KafkaException("Nobody expects the Spanish Inquisition");
-        final ErrorBackgroundEvent errorBackgroundEvent = new ErrorBackgroundEvent(expectedException);
+        final ErrorEvent errorBackgroundEvent = new ErrorEvent(expectedException);
         backgroundEventQueue.add(errorBackgroundEvent);
         consumer.subscribe(Collections.singletonList("t1"));
         final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
@@ -240,10 +240,10 @@ public void testMultipleBackgroundErrors() {
         consumer = newConsumer(config);
 
         final KafkaException expectedException1 = new KafkaException("Nobody expects the Spanish Inquisition");
-        final ErrorBackgroundEvent errorBackgroundEvent1 = new ErrorBackgroundEvent(expectedException1);
+        final ErrorEvent errorBackgroundEvent1 = new ErrorEvent(expectedException1);
         backgroundEventQueue.add(errorBackgroundEvent1);
         final KafkaException expectedException2 = new KafkaException("Spam, Spam, Spam");
-        final ErrorBackgroundEvent errorBackgroundEvent2 = new ErrorBackgroundEvent(expectedException2);
+        final ErrorEvent errorBackgroundEvent2 = new ErrorEvent(expectedException2);
         backgroundEventQueue.add(errorBackgroundEvent2);
         consumer.subscribe(Collections.singletonList("t1"));
         final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index 6947f236e9a54..42728c2ca8a45 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -65,10 +65,10 @@
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.common.CurrentAssignmentBuilder;
 import org.apache.kafka.coordinator.group.common.TargetAssignmentBuilder;
-import org.apache.kafka.coordinator.group.consumer.CurrentAssignmentBuilder;
+import org.apache.kafka.coordinator.group.common.CurrentAssignmentBuilder;
 import org.apache.kafka.coordinator.group.consumer.MemberState;
-import org.apache.kafka.coordinator.group.consumer.TargetAssignmentBuilder;
-import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
+import org.apache.kafka.coordinator.group.common.TargetAssignmentBuilder;
+import org.apache.kafka.coordinator.group.common.TopicMetadata;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
@@ -2533,10 +2533,8 @@ public void replay(
                     (ConsumerGroupMember) oldMember)
                     .setMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                     .setPreviousMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
-                    .setTargetMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                     .setAssignedPartitions(Collections.emptyMap())
                     .setPartitionsPendingRevocation(Collections.emptyMap())
-                    .setPartitionsPendingAssignment(Collections.emptyMap())
                     .build();
                 group.updateMember(newMember);
             } else {

From 539ba958b5dca451318eb82a4481aa8f6406a378 Mon Sep 17 00:00:00 2001
From: ShivsundarR <shr@confluent.io>
Date: Fri, 5 Apr 2024 18:44:13 +0530
Subject: [PATCH 254/258] Updated Tuple2->Entry changes

---
 .../internals/MembershipManagerImplTest.java  | 12 --------
 .../kafka/server/SharePartitionManager.java   | 28 ++++++++++---------
 .../server/SharePartitionManagerTest.java     |  8 +++---
 .../kafka/tools/ShareGroupsCommand.java       | 17 ++++++-----
 .../reassign/ReassignPartitionsCommand.java   |  1 -
 .../ReassignPartitionsIntegrationTest.java    |  1 -
 .../reassign/ReassignPartitionsUnitTest.java  |  1 -
 7 files changed, 29 insertions(+), 39 deletions(-)

diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
index ada9fbb7a3706..cef608a902d1f 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
@@ -2043,18 +2043,6 @@ private static SortedSet<TopicPartition> topicPartitions(Map<Uuid, SortedSet<Int
         return topicPartitions;
     }
 
-    private static SortedSet<TopicPartition> topicPartitions(Map<Uuid, SortedSet<Integer>> topicIdMap, Map<Uuid, String> topicIdNames) {
-        SortedSet<TopicPartition> topicPartitions = new TreeSet<>(new Utils.TopicPartitionComparator());
-
-        for (Uuid topicId : topicIdMap.keySet()) {
-            for (int partition : topicIdMap.get(topicId)) {
-                topicPartitions.add(new TopicPartition(topicIdNames.get(topicId), partition));
-            }
-        }
-
-        return topicPartitions;
-    }
-
     private SortedSet<TopicPartition> topicPartitions(String topicName, int... partitions) {
         SortedSet<TopicPartition> topicPartitions = new TreeSet<>(new Utils.TopicPartitionComparator());
 
diff --git a/core/src/main/java/kafka/server/SharePartitionManager.java b/core/src/main/java/kafka/server/SharePartitionManager.java
index 389b98a526abe..1c7521984bc39 100644
--- a/core/src/main/java/kafka/server/SharePartitionManager.java
+++ b/core/src/main/java/kafka/server/SharePartitionManager.java
@@ -36,6 +36,7 @@
 import org.apache.kafka.storage.internals.log.FetchPartitionData;
 import org.slf4j.Logger;
 
+import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -44,6 +45,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
@@ -165,7 +167,7 @@ public void maybeProcessFetchQueue() {
                 shareFetchPartitionData.fetchParams,
                 CollectionConverters.asScala(
                     topicPartitionData.entrySet().stream().map(entry ->
-                        new Tuple2<>(entry.getKey(), entry.getValue())).collect(Collectors.toList())
+                        new Tuple2(entry.getKey(), entry.getValue())).collect(Collectors.toList())
                 ),
                 QuotaFactory.UnboundedQuota$.MODULE$,
                 responsePartitionData -> {
@@ -552,11 +554,11 @@ public enum ModifiedTopicIdPartitionType {
 
     // Helper class to return the erroneous partitions and valid partition data
     public static class ErroneousAndValidPartitionData {
-        private final List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous;
-        private final List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions;
+        private final List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous;
+        private final List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions;
 
-        public ErroneousAndValidPartitionData(List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous,
-                                              List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions) {
+        public ErroneousAndValidPartitionData(List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous,
+                                              List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions) {
             this.erroneous = erroneous;
             this.validTopicIdPartitions = validTopicIdPartitions;
         }
@@ -566,18 +568,18 @@ public ErroneousAndValidPartitionData(Map<TopicIdPartition, ShareFetchRequest.Sh
             validTopicIdPartitions = new ArrayList<>();
             shareFetchData.forEach((topicIdPartition, sharePartitionData) -> {
                 if (topicIdPartition.topic() == null) {
-                    erroneous.add(new Tuple2<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
+                    erroneous.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
                 } else {
-                    validTopicIdPartitions.add(new Tuple2<>(topicIdPartition, sharePartitionData));
+                    validTopicIdPartitions.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, sharePartitionData));
                 }
             });
         }
 
-        public List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous() {
+        public List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous() {
             return erroneous;
         }
 
-        public List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions() {
+        public List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions() {
             return validTopicIdPartitions;
         }
     }
@@ -813,8 +815,8 @@ ErroneousAndValidPartitionData getErroneousAndValidTopicIdPartitions() {
             if (!isSubsequent) {
                 return new ErroneousAndValidPartitionData(shareFetchData);
             } else {
-                List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous = new ArrayList<>();
-                List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> valid = new ArrayList<>();
+                List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous = new ArrayList<>();
+                List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> valid = new ArrayList<>();
                 // Take the session lock and iterate over all the cached partitions.
                 synchronized (session) {
                     session.partitionMap.forEach(cachedSharePartition -> {
@@ -822,9 +824,9 @@ ErroneousAndValidPartitionData getErroneousAndValidTopicIdPartitions() {
                                 TopicPartition(cachedSharePartition.topic, cachedSharePartition.partition));
                         ShareFetchRequest.SharePartitionData reqData = cachedSharePartition.reqData();
                         if (topicIdPartition.topic() == null) {
-                            erroneous.add(new Tuple2<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
+                            erroneous.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
                         } else {
-                            valid.add(new Tuple2<>(topicIdPartition, reqData));
+                            valid.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, reqData));
                         }
                     });
                     return new ErroneousAndValidPartitionData(erroneous, valid);
diff --git a/core/src/test/java/kafka/server/SharePartitionManagerTest.java b/core/src/test/java/kafka/server/SharePartitionManagerTest.java
index e31bf0838aea1..712232648b641 100644
--- a/core/src/test/java/kafka/server/SharePartitionManagerTest.java
+++ b/core/src/test/java/kafka/server/SharePartitionManagerTest.java
@@ -813,10 +813,10 @@ private void assertErroneousAndValidTopicIdPartitions(SharePartitionManager.Erro
                                                     List<TopicIdPartition> expectedErroneous, List<TopicIdPartition> expectedValid) {
         List<TopicIdPartition> actualErroneousPartitions = new ArrayList<>();
         List<TopicIdPartition> actualValidPartitions = new ArrayList<>();
-        erroneousAndValidPartitionData.erroneous().forEach(topicIdPartitionPartitionDataTuple2 ->
-                actualErroneousPartitions.add(topicIdPartitionPartitionDataTuple2._1));
-        erroneousAndValidPartitionData.validTopicIdPartitions().forEach(topicIdPartitionPartitionDataTuple2 ->
-                actualValidPartitions.add(topicIdPartitionPartitionDataTuple2._1));
+        erroneousAndValidPartitionData.erroneous().forEach(topicIdPartitionPartitionDataEntry ->
+                actualErroneousPartitions.add(topicIdPartitionPartitionDataEntry.getKey()));
+        erroneousAndValidPartitionData.validTopicIdPartitions().forEach(topicIdPartitionPartitionDataEntry ->
+                actualValidPartitions.add(topicIdPartitionPartitionDataEntry.getKey()));
         assertEquals(expectedErroneous, actualErroneousPartitions);
         assertEquals(expectedValid, actualValidPartitions);
     }
diff --git a/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java b/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
index 4bfe8098f4333..d3845b9eeaeb2 100644
--- a/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
@@ -32,8 +32,10 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.util.CommandLineUtils;
+import org.apache.kafka.timeline.TimelineHashMap;
 
 import java.io.IOException;
+import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -42,6 +44,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
@@ -127,7 +130,7 @@ public void listGroups() throws ExecutionException, InterruptedException {
             ? Collections.emptySet()
             : shareGroupStatesFromString(stateValue);
         List<ShareGroupListing> listings = listShareGroupsWithState(states);
-        printGroupStates(listings.stream().map(e -> new Tuple2<>(e.groupId(), e.state().toString())).collect(Collectors.toList()));
+        printGroupStates(listings.stream().map(e -> new AbstractMap.SimpleImmutableEntry<String, String>(e.groupId(), e.state().toString()) {}).collect(Collectors.toList()));
       } else
         listShareGroups().forEach(System.out::println);
     }
@@ -149,17 +152,17 @@ List<ShareGroupListing> listShareGroupsWithState(Set<ShareGroupState> states) th
       return new ArrayList<>(result.all().get());
     }
 
-    private void printGroupStates(List<Tuple2<String, String>> groupsAndStates) {
+    private void printGroupStates(List<Entry<String, String>> groupsAndStates) {
       // find proper columns width
       int maxGroupLen = 15;
-      for (Tuple2<String, String> tuple : groupsAndStates) {
-        String groupId = tuple.v1;
+      for (Entry<String, String> entry : groupsAndStates) {
+        String groupId = entry.getKey();
         maxGroupLen = Math.max(maxGroupLen, groupId.length());
       }
       System.out.printf("%" + (-maxGroupLen) + "s %s", "GROUP", "STATE");
-      for (Tuple2<String, String> tuple : groupsAndStates) {
-        String groupId = tuple.v1;
-        String state = tuple.v2;
+      for (Entry<String, String> entry : groupsAndStates) {
+        String groupId = entry.getKey();
+        String state = entry.getValue();
         System.out.printf("%" + (-maxGroupLen) + "s %s", groupId, state);
       }
     }
diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
index 1e5978a1f9f4f..a23a25e6f57b5 100644
--- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
@@ -49,7 +49,6 @@
 import org.apache.kafka.server.util.json.JsonValue;
 import org.apache.kafka.tools.TerseException;
 import org.apache.kafka.tools.ToolsUtils;
-import org.apache.kafka.tools.Tuple2;
 
 import java.io.IOException;
 import java.util.AbstractMap.SimpleImmutableEntry;
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
index e6b24b9d62ca4..0348aa24ffe9f 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsIntegrationTest.java
@@ -43,7 +43,6 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.tools.TerseException;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
index d2301b5f36c6b..c6f145d9a5598 100644
--- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
@@ -31,7 +31,6 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.common.AdminCommandFailedException;
 import org.apache.kafka.server.common.AdminOperationException;
-import org.apache.kafka.tools.Tuple2;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;

From dcdf5b4e7ba1ff5aa3927b2e2f5b5f8187e47776 Mon Sep 17 00:00:00 2001
From: Apoorv Mittal <amittal@confluent.io>
Date: Wed, 10 Apr 2024 16:16:59 +0100
Subject: [PATCH 255/258] Resolving merge conflicts

---
 .../kafka/coordinator/group/GroupMember.java  |  44 +----
 .../group/GroupMetadataManager.java           |  78 ++++-----
 .../coordinator/group/RecordHelpers.java      |   2 +-
 .../common/CurrentAssignmentBuilder.java      |  35 +++-
 .../{consumer => common}/MemberState.java     |   2 +-
 .../group/consumer/ConsumerGroup.java         |   2 +-
 .../group/consumer/ConsumerGroupMember.java   | 128 +--------------
 .../group/share/ShareGroupMember.java         |  27 ++--
 .../kafka/coordinator/group/Assertions.java   |   6 -
 .../group/GroupMetadataManagerTest.java       | 151 ++----------------
 .../GroupMetadataManagerTestContext.java      |  95 +++++++++--
 .../coordinator/group/RecordHelpersTest.java  |   2 +-
 .../common/CurrentAssignmentBuilderTest.java  |  28 ++--
 .../group/consumer/ConsumerGroupTest.java     |   1 +
 14 files changed, 196 insertions(+), 405 deletions(-)
 rename group-coordinator/src/main/java/org/apache/kafka/coordinator/group/{consumer => common}/MemberState.java (97%)

diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMember.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMember.java
index fba7e16bb418a..512cfc9e009b7 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMember.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMember.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.coordinator.group;
 
 import org.apache.kafka.common.Uuid;
-import org.apache.kafka.coordinator.group.common.CurrentAssignmentBuilder;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 
 import java.util.Collections;
@@ -32,27 +32,6 @@
  */
 public abstract class GroupMember {
 
-    /**
-     * The various states that a member can be in. For their definition,
-     * refer to the documentation of {{@link CurrentAssignmentBuilder}}.
-     */
-    public enum MemberState {
-        REVOKING("revoking"),
-        ASSIGNING("assigning"),
-        STABLE("stable");
-
-        private final String name;
-
-        MemberState(String name) {
-            this.name = name;
-        }
-
-        @Override
-        public String toString() {
-            return name;
-        }
-    }
-
     /**
      * The member id.
      */
@@ -68,13 +47,6 @@ public String toString() {
      */
     protected int previousMemberEpoch;
 
-    /**
-     * The next member epoch. This corresponds to the target
-     * assignment epoch used to compute the current assigned,
-     * revoking and assigning partitions.
-     */
-    protected int targetMemberEpoch;
-
     /**
      * The instance id provided by the member.
      */
@@ -136,13 +108,6 @@ public int previousMemberEpoch() {
         return previousMemberEpoch;
     }
 
-    /**
-     * @return The target member epoch.
-     */
-    public int targetMemberEpoch() {
-        return targetMemberEpoch;
-    }
-
     /**
      * @return The instance id.
      */
@@ -206,4 +171,11 @@ public static Map<Uuid, Set<Integer>> assignmentFromTopicPartitions(
             ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions::topicId,
             topicPartitions -> Collections.unmodifiableSet(new HashSet<>(topicPartitions.partitions()))));
     }
+
+    /**
+     * @return True if the member is in the Stable state and at the desired epoch.
+     */
+    public boolean isReconciledTo(int targetAssignmentEpoch) {
+        return state == MemberState.STABLE && memberEpoch == targetAssignmentEpoch;
+    }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index c6f906676cf5c..fc5aa45b9518b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -65,10 +65,7 @@
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.common.CurrentAssignmentBuilder;
 import org.apache.kafka.coordinator.group.common.TargetAssignmentBuilder;
-import org.apache.kafka.coordinator.group.common.CurrentAssignmentBuilder;
-import org.apache.kafka.coordinator.group.consumer.MemberState;
-import org.apache.kafka.coordinator.group.common.TargetAssignmentBuilder;
-import org.apache.kafka.coordinator.group.common.TopicMetadata;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
@@ -765,7 +762,28 @@ ConsumerGroup getOrMaybeCreatePersistedConsumerGroup(
         String groupId,
         boolean createIfNotExists
     ) throws GroupIdNotFoundException {
-        return (ConsumerGroup) getOrMaybeCreateGroup(groupId, createIfNotExists, CONSUMER);
+        return (ConsumerGroup) getOrMaybeCreatePersistedGroup(groupId, createIfNotExists, CONSUMER);
+    }
+
+    /**
+     * The method should be called on the replay path.
+     * Gets or maybe creates a consumer group and updates the groups map if a new group is created.
+     *
+     * @param groupId           The group id.
+     * @param createIfNotExists A boolean indicating whether the group should be
+     *                          created if it does not exist.
+     *
+     * @return A ConsumerGroup.
+     * @throws IllegalStateException if the group does not exist and createIfNotExists is false or
+     *                               if the group is not a consumer group.
+     * Package private for testing.
+     */
+    AbstractGroup getOrMaybeCreatePersistedGroup(
+        String groupId,
+        boolean createIfNotExists,
+        GroupType groupType
+    ) throws GroupIdNotFoundException {
+        return getOrMaybeCreateGroup(groupId, createIfNotExists, groupType);
     }
 
     /**
@@ -807,7 +825,7 @@ private AbstractGroup getOrMaybeCreateGroup(
         Group group = groups.get(groupId);
 
         if (group == null && !createIfNotExists) {
-            throw new GroupIdNotFoundException(String.format("Group %s not found.", groupId));
+            throw new IllegalStateException(String.format("Group %s not found.", groupId));
         }
 
         if (group == null) {
@@ -894,31 +912,6 @@ public ClassicGroup classicGroup(
         }
     }
 
-    /**
-     * Gets a consumer group by committed offset.
-     *
-     * @param groupId           The group id.
-     * @param committedOffset   A specified committed offset corresponding to this shard.
-     *
-     * @return A ConsumerGroup.
-     * @throws GroupIdNotFoundException if the group does not exist or is not a consumer group.
-     */
-    public ConsumerGroup consumerGroup(
-        String groupId,
-        long committedOffset
-    ) throws GroupIdNotFoundException {
-        Group group = group(groupId, committedOffset);
-
-        if (group.type() == CONSUMER) {
-            return (ConsumerGroup) group;
-        } else {
-            // We don't support upgrading/downgrading between protocols at the moment so
-            // we throw an exception if a group exists with the wrong type.
-            throw new GroupIdNotFoundException(String.format("Group %s is not a consumer group.",
-                groupId));
-        }
-    }
-
     /**
      * Gets a share group by committed offset.
      *
@@ -1233,7 +1226,7 @@ private void throwIfStaticMemberIsUnknown(ConsumerGroupMember staticMember, Stri
         }
     }
 
-    private ConsumerGroupHeartbeatResponseData.Assignment createConsumerGroupResponseAssignment(
+    private ConsumerGroupHeartbeatResponseData.Assignment createResponseAssignment(
         ConsumerGroupMember member
     ) {
         return new ConsumerGroupHeartbeatResponseData.Assignment()
@@ -1311,7 +1304,7 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGr
         // Get or create the consumer group.
         boolean createIfNotExists = memberEpoch == 0;
         final ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, createIfNotExists, records);
-        throwIfConsumerGroupIsFull(group, memberId);
+        throwIfGroupIsFull(group, memberId);
 
         // Get or create the member.
         if (memberId.isEmpty()) memberId = Uuid.randomUuid().toString();
@@ -1523,7 +1516,7 @@ private ConsumerGroupMember maybeReconcile(
             return member;
         }
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        ConsumerGroupMember updatedMember = (ConsumerGroupMember) new CurrentAssignmentBuilder(member)
             .withTargetAssignment(targetAssignmentEpoch, targetAssignment)
             .withCurrentPartitionEpoch(currentPartitionEpoch)
             .withOwnedTopicPartitions(ownedTopicPartitions)
@@ -1677,7 +1670,7 @@ private CoordinatorResult<ShareGroupHeartbeatResponseData, Record> shareGroupHea
         // 3. Reconcile the member's assignment with the target assignment. This is only required if
         // a new target assignment has been installed.
         boolean assignmentUpdated = false;
-        if (updatedMember.targetMemberEpoch() != targetAssignmentEpoch) {
+        if (!member.isReconciledTo(targetAssignmentEpoch)) {
             ShareGroupMember prevMember = updatedMember;
             updatedMember = (ShareGroupMember) new CurrentAssignmentBuilder(updatedMember)
                     .withTargetAssignment(targetAssignmentEpoch, targetAssignment)
@@ -2016,7 +2009,7 @@ private void scheduleConsumerGroupRebalanceTimeout(
         int memberEpoch,
         int rebalanceTimeoutMs
     ) {
-        String key = consumerGroupRebalanceTimeoutKey(groupId, memberId);
+        String key = groupRebalanceTimeoutKey(groupId, memberId);
         timer.schedule(key, rebalanceTimeoutMs, TimeUnit.MILLISECONDS, true, () -> {
             try {
                 ConsumerGroup group = consumerGroup(groupId);
@@ -2054,7 +2047,7 @@ private void cancelConsumerGroupRebalanceTimeout(
         String groupId,
         String memberId
     ) {
-        timer.cancel(consumerGroupRebalanceTimeoutKey(groupId, memberId));
+        timer.cancel(groupRebalanceTimeoutKey(groupId, memberId));
     }
 
     /**
@@ -2190,7 +2183,7 @@ public void replay(
         String groupId = key.groupId();
         String memberId = key.memberId();
 
-        ConsumerGroup group = getOrMaybeCreateConsumerGroup(groupId, value != null);
+        ConsumerGroup group = getOrMaybeCreatePersistedConsumerGroup(groupId, value != null);
         Set<String> oldSubscribedTopicNames = new HashSet<>(group.subscribedTopicNames());
 
         if (value != null) {
@@ -2315,7 +2308,7 @@ public void replay(
             Group group = groups.get(groupId);
             if (group == null) {
                 // Safe check as tombstone record should be replayed after group creation.
-                throw new GroupIdNotFoundException(String.format("Group %s not found.", groupId));
+                throw new IllegalStateException(String.format("Group %s not found.", groupId));
             }
             // Group can either be share or consumer group only hence safe to cast.
             AbstractGroup abstractGroup = (AbstractGroup) group;
@@ -2410,7 +2403,7 @@ public void replay(
     ) {
         String groupId = key.groupId();
         String memberId = key.memberId();
-        AbstractGroup group = getOrMaybeCreateGroup(groupId, false, groupType);
+        AbstractGroup group = getOrMaybeCreatePersistedGroup(groupId, false, groupType);
 
         if (value != null) {
             group.updateTargetAssignment(memberId, Assignment.fromRecord(value));
@@ -2449,7 +2442,7 @@ public void replay(
         GroupType groupType
     ) {
         String groupId = key.groupId();
-        AbstractGroup group = getOrMaybeCreateGroup(groupId, false, groupType);
+        AbstractGroup group = getOrMaybeCreatePersistedGroup(groupId, false, groupType);
 
         if (value != null) {
             group.setTargetAssignmentEpoch(value.assignmentEpoch());
@@ -2521,7 +2514,6 @@ public void replay(
                     (ShareGroupMember) oldMember)
                     .setMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                     .setPreviousMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
-                    .setTargetMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                     .setAssignedPartitions(Collections.emptyMap())
                     .build();
                 group.updateMember(newMember);
@@ -2694,7 +2686,7 @@ public static String groupSessionTimeoutKey(String groupId, String memberId) {
         return "session-timeout-" + groupId + "-" + memberId;
     }
 
-    public static String consumerGroupRebalanceTimeoutKey(String groupId, String memberId) {
+    public static String groupRebalanceTimeoutKey(String groupId, String memberId) {
         return "rebalance-timeout-" + groupId + "-" + memberId;
     }
 
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
index 1a5341df4b3fb..402aa4c2cd57a 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/RecordHelpers.java
@@ -388,7 +388,7 @@ public static Record newCurrentAssignmentRecord(
                 new ConsumerGroupCurrentMemberAssignmentValue()
                     .setMemberEpoch(member.memberEpoch())
                     .setPreviousMemberEpoch(member.previousMemberEpoch())
-                    .setTargetMemberEpoch(member.targetMemberEpoch())
+                    .setState(member.state().value())
                     .setAssignedPartitions(toTopicPartitions(member.assignedPartitions())),
                 (short) 0
             )
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilder.java
index 78c401ec07ab7..ebf7429a2bd9a 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilder.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilder.java
@@ -21,7 +21,6 @@
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
 import org.apache.kafka.coordinator.group.GroupMember;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
-import org.apache.kafka.coordinator.group.consumer.MemberState;
 import org.apache.kafka.coordinator.group.share.ShareGroupMember;
 
 import java.util.Collections;
@@ -143,7 +142,7 @@ else if (member instanceof ShareGroupMember)
      *
      * @return A new ConsumerGroupMember or the current one.
      */
-    public ConsumerGroupMember build() {
+    public ConsumerGroupMember buildConsumerGroupMember(ConsumerGroupMember member) {
         switch (member.state()) {
             case STABLE:
                 // When the member is in the STABLE state, we verify if a newer
@@ -220,6 +219,30 @@ public ConsumerGroupMember build() {
         return member;
     }
 
+    /**
+     * The state machine for share group has just one state: STABLE. This state means that the member
+     * has received all its assigned partitions. However, the member may still need to update its
+     * metadata to reflect the new target assignment.
+     *
+     * @param member The share group member to reconcile.
+     *
+     * @return A new ShareGroupMember.
+     */
+    private ShareGroupMember buildShareGroupMember(ShareGroupMember member) {
+        // A new target assignment has been installed, we need to restart
+        // the reconciliation loop from the beginning.
+        if (targetAssignmentEpoch != member.memberEpoch()) {
+            // We transition to the target epoch. The transition to the new state is done
+            // when the member is updated.
+            return new ShareGroupMember.Builder(member)
+                .setAssignedPartitions(targetAssignment.partitions())
+                .setMemberEpoch(targetAssignmentEpoch)
+                .build();
+        }
+
+        return member;
+    }
+
     /**
      * Computes the next assignment.
      *
@@ -279,7 +302,7 @@ private ConsumerGroupMember computeNextAssignment(
             // epoch and requests the revocation of those partitions. It transitions to
             // the UNREVOKED_PARTITIONS state to wait until the client acknowledges the
             // revocation of the partitions.
-            return new ConsumerGroupMember.Builder(member)
+            return new ConsumerGroupMember.Builder((ConsumerGroupMember) member)
                 .setState(MemberState.UNREVOKED_PARTITIONS)
                 .updateMemberEpoch(memberEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
@@ -295,7 +318,7 @@ private ConsumerGroupMember computeNextAssignment(
                 .computeIfAbsent(topicId, __ -> new HashSet<>())
                 .addAll(partitions));
             MemberState newState = hasUnreleasedPartitions ? MemberState.UNRELEASED_PARTITIONS : MemberState.STABLE;
-            return new ConsumerGroupMember.Builder(member)
+            return new ConsumerGroupMember.Builder((ConsumerGroupMember) member)
                 .setState(newState)
                 .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
@@ -305,7 +328,7 @@ private ConsumerGroupMember computeNextAssignment(
             // If there are no partitions to be revoked nor to be assigned but some
             // partitions are not available yet, the member transitions to the target
             // epoch, to the UNRELEASED_PARTITIONS state and waits.
-            return new ConsumerGroupMember.Builder(member)
+            return new ConsumerGroupMember.Builder((ConsumerGroupMember) member)
                 .setState(MemberState.UNRELEASED_PARTITIONS)
                 .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
@@ -314,7 +337,7 @@ private ConsumerGroupMember computeNextAssignment(
         } else {
             // Otherwise, the member transitions to the target epoch and to the
             // STABLE state.
-            return new ConsumerGroupMember.Builder(member)
+            return new ConsumerGroupMember.Builder((ConsumerGroupMember) member)
                 .setState(MemberState.STABLE)
                 .updateMemberEpoch(targetAssignmentEpoch)
                 .setAssignedPartitions(newAssignedPartitions)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/MemberState.java
similarity index 97%
rename from group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java
rename to group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/MemberState.java
index 3f3237bfe0b12..68a63132b1a1c 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/MemberState.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/common/MemberState.java
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.coordinator.group.consumer;
+package org.apache.kafka.coordinator.group.common;
 
 import java.util.HashMap;
 import java.util.Map;
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
index 705ef3e8d45e9..352ddcbd7604b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroup.java
@@ -347,7 +347,7 @@ protected void maybeUpdateGroupState() {
             newState = ASSIGNING;
         } else {
             for (GroupMember member : members.values()) {
-                if (member.targetMemberEpoch() != targetAssignmentEpoch.get() || member.state() != ConsumerGroupMember.MemberState.STABLE) {
+                if (!member.isReconciledTo(targetAssignmentEpoch.get())) {
                     newState = RECONCILING;
                     break;
                 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
index c6f5e8783bae3..300724ae43e2e 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupMember.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.coordinator.group.GroupMember;
 import org.apache.kafka.coordinator.group.Utils;
 import org.apache.kafka.coordinator.group.common.Assignment;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue;
 import org.apache.kafka.image.TopicsImage;
@@ -34,7 +35,6 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 /**
  * ConsumerGroupMember contains all the information related to a member
@@ -233,55 +233,6 @@ public ConsumerGroupMember build() {
         }
     }
 
-    /**
-     * The member id.
-     */
-    private final String memberId;
-
-    /**
-     * The current member epoch.
-     */
-    private final int memberEpoch;
-
-    /**
-     * The previous member epoch.
-     */
-    private final int previousMemberEpoch;
-
-    /**
-     * The member state.
-     */
-    private final MemberState state;
-
-    /**
-     * The instance id provided by the member.
-     */
-    private final String instanceId;
-
-    /**
-     * The rack id provided by the member.
-     */
-    private final String rackId;
-
-    /**
-     * The rebalance timeout provided by the member.
-     */
-    private final int rebalanceTimeoutMs;
-
-    /**
-     * The client id reported by the member.
-     */
-    private final String clientId;
-
-    /**
-     * The host reported by the member.
-     */
-    private final String clientHost;
-
-    /**
-     * The list of subscriptions (topic names) configured by the member.
-     */
-    private final List<String> subscribedTopicNames;
 
     /**
      * The subscription pattern configured by the member.
@@ -335,69 +286,6 @@ private ConsumerGroupMember(
         this.partitionsPendingRevocation = partitionsPendingRevocation;
     }
 
-    /**
-     * @return The member id.
-     */
-    public String memberId() {
-        return memberId;
-    }
-
-    /**
-     * @return The current member epoch.
-     */
-    public int memberEpoch() {
-        return memberEpoch;
-    }
-
-    /**
-     * @return The previous member epoch.
-     */
-    public int previousMemberEpoch() {
-        return previousMemberEpoch;
-    }
-
-    /**
-     * @return The instance id.
-     */
-    public String instanceId() {
-        return instanceId;
-    }
-
-    /**
-     * @return The rack id.
-     */
-    public String rackId() {
-        return rackId;
-    }
-
-    /**
-     * @return The rebalance timeout in millis.
-     */
-    public int rebalanceTimeoutMs() {
-        return rebalanceTimeoutMs;
-    }
-
-    /**
-     * @return The client id.
-     */
-    public String clientId() {
-        return clientId;
-    }
-
-    /**
-     * @return The client host.
-     */
-    public String clientHost() {
-        return clientHost;
-    }
-
-    /**
-     * @return The list of subscribed topic names.
-     */
-    public List<String> subscribedTopicNames() {
-        return subscribedTopicNames;
-    }
-
     /**
      * @return The regular expression based subscription.
      */
@@ -412,20 +300,6 @@ public Optional<String> serverAssignorName() {
         return Optional.ofNullable(serverAssignorName);
     }
 
-    /**
-     * @return The current state.
-     */
-    public MemberState state() {
-        return state;
-    }
-
-    /**
-     * @return True if the member is in the Stable state and at the desired epoch.
-     */
-    public boolean isReconciledTo(int targetAssignmentEpoch) {
-        return state == MemberState.STABLE && memberEpoch == targetAssignmentEpoch;
-    }
-
     /**
      * @return The set of assigned partitions.
      */
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/share/ShareGroupMember.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/share/ShareGroupMember.java
index b82fe7ba0b2cd..34fd3687459e0 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/share/ShareGroupMember.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/share/ShareGroupMember.java
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.message.ShareGroupDescribeResponseData;
 import org.apache.kafka.coordinator.group.GroupMember;
 import org.apache.kafka.coordinator.group.Utils;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ShareGroupMemberMetadataValue;
 import org.apache.kafka.image.TopicsImage;
@@ -42,7 +43,7 @@ public class ShareGroupMember extends GroupMember {
   /**
    * A builder that facilitates the creation of a new member or the update of
    * an existing one.
-   *
+   * <p>
    * Please refer to the javadoc of {{@link ShareGroupMember}} for the
    * definition of the fields.
    */
@@ -50,7 +51,7 @@ public static class Builder {
     private final String memberId;
     private int memberEpoch = 0;
     private int previousMemberEpoch = -1;
-    private int targetMemberEpoch = 0;
+    private MemberState state = MemberState.STABLE;
     private String rackId = null;
     private int rebalanceTimeoutMs = -1;
     private String clientId = "";
@@ -68,7 +69,6 @@ public Builder(ShareGroupMember member) {
       this.memberId = member.memberId();
       this.memberEpoch = member.memberEpoch;
       this.previousMemberEpoch = member.previousMemberEpoch;
-      this.targetMemberEpoch = member.targetMemberEpoch;
       this.rackId = member.rackId;
       this.rebalanceTimeoutMs = member.rebalanceTimeoutMs;
       this.clientId = member.clientId;
@@ -87,11 +87,6 @@ public Builder setPreviousMemberEpoch(int previousMemberEpoch) {
       return this;
     }
 
-    public Builder setTargetMemberEpoch(int targetMemberEpoch) {
-      this.targetMemberEpoch = targetMemberEpoch;
-      return this;
-    }
-
     public Builder setRackId(String rackId) {
       this.rackId = rackId;
       return this;
@@ -134,6 +129,11 @@ public Builder maybeUpdateSubscribedTopicNames(Optional<List<String>> subscribed
       return this;
     }
 
+    public Builder setState(MemberState state) {
+      this.state = state;
+      return this;
+    }
+
     public Builder setAssignedPartitions(Map<Uuid, Set<Integer>> assignedPartitions) {
       this.assignedPartitions = assignedPartitions;
       return this;
@@ -152,7 +152,7 @@ public Builder updateWith(
         ConsumerGroupCurrentMemberAssignmentValue record) {
       setMemberEpoch(record.memberEpoch());
       setPreviousMemberEpoch(record.previousMemberEpoch());
-      setTargetMemberEpoch(record.targetMemberEpoch());
+      setState(MemberState.fromValue(record.state()));
       setAssignedPartitions(assignmentFromTopicPartitions(record.assignedPartitions()));
       return this;
     }
@@ -162,13 +162,12 @@ public ShareGroupMember build() {
               memberId,
               memberEpoch,
               previousMemberEpoch,
-              targetMemberEpoch,
               rackId,
               rebalanceTimeoutMs,
               clientId,
               clientHost,
               subscribedTopicNames,
-              MemberState.STABLE,
+              state,
               assignedPartitions
       );
     }
@@ -178,7 +177,6 @@ private ShareGroupMember(
           String memberId,
           int memberEpoch,
           int previousMemberEpoch,
-          int targetMemberEpoch,
           String rackId,
           int rebalanceTimeoutMs,
           String clientId,
@@ -190,7 +188,6 @@ private ShareGroupMember(
     this.memberId = memberId;
     this.memberEpoch = memberEpoch;
     this.previousMemberEpoch = previousMemberEpoch;
-    this.targetMemberEpoch = targetMemberEpoch;
     this.rackId = rackId;
     this.rebalanceTimeoutMs = rebalanceTimeoutMs;
     this.clientId = clientId;
@@ -206,7 +203,6 @@ private ShareGroupMember(
   public String currentAssignmentSummary() {
     return "CurrentAssignment(memberEpoch=" + memberEpoch +
             ", previousMemberEpoch=" + previousMemberEpoch +
-            ", targetMemberEpoch=" + targetMemberEpoch +
             ", state=" + state +
             ", assignedPartitions=" + assignedPartitions +
             ')';
@@ -219,7 +215,6 @@ public boolean equals(Object o) {
     ShareGroupMember that = (ShareGroupMember) o;
     return memberEpoch == that.memberEpoch
             && previousMemberEpoch == that.previousMemberEpoch
-            && targetMemberEpoch == that.targetMemberEpoch
             && rebalanceTimeoutMs == that.rebalanceTimeoutMs
             && Objects.equals(memberId, that.memberId)
             && Objects.equals(rackId, that.rackId)
@@ -234,7 +229,6 @@ public int hashCode() {
     int result = memberId != null ? memberId.hashCode() : 0;
     result = 31 * result + memberEpoch;
     result = 31 * result + previousMemberEpoch;
-    result = 31 * result + targetMemberEpoch;
     result = 31 * result + Objects.hashCode(rackId);
     result = 31 * result + rebalanceTimeoutMs;
     result = 31 * result + Objects.hashCode(clientId);
@@ -250,7 +244,6 @@ public String toString() {
             "memberId='" + memberId + '\'' +
             ", memberEpoch=" + memberEpoch +
             ", previousMemberEpoch=" + previousMemberEpoch +
-            ", targetMemberEpoch=" + targetMemberEpoch +
             ", rackId='" + rackId + '\'' +
             ", rebalanceTimeoutMs=" + rebalanceTimeoutMs +
             ", clientId='" + clientId + '\'' +
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
index 68c1390abcafd..9dc355669eaf7 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/Assertions.java
@@ -194,18 +194,12 @@ private static void assertApiMessageAndVersionEquals(
 
             assertEquals(expectedValue.memberEpoch(), actualValue.memberEpoch());
             assertEquals(expectedValue.previousMemberEpoch(), actualValue.previousMemberEpoch());
-            assertEquals(expectedValue.targetMemberEpoch(), actualValue.targetMemberEpoch());
-            assertEquals(expectedValue.error(), actualValue.error());
-            assertEquals(expectedValue.metadataVersion(), actualValue.metadataVersion());
-            assertEquals(expectedValue.metadataBytes(), actualValue.metadataBytes());
 
             // We transform those to Maps before comparing them.
             assertEquals(fromTopicPartitions(expectedValue.assignedPartitions()),
                 fromTopicPartitions(actualValue.assignedPartitions()));
             assertEquals(fromTopicPartitions(expectedValue.partitionsPendingRevocation()),
                 fromTopicPartitions(actualValue.partitionsPendingRevocation()));
-            assertEquals(fromTopicPartitions(expectedValue.partitionsPendingAssignment()),
-                fromTopicPartitions(actualValue.partitionsPendingAssignment()));
         } else if (actual.message() instanceof ConsumerGroupPartitionMetadataValue) {
             // The order of the racks stored in the PartitionMetadata of the ConsumerGroupPartitionMetadataValue
             // is not always guaranteed. Therefore, we need a special comparator.
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index 819ea91e6d2c4..d1fa788009cc3 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -67,13 +67,10 @@
 import org.apache.kafka.coordinator.group.classic.ClassicGroupState;
 import org.apache.kafka.coordinator.group.common.Assignment;
 import org.apache.kafka.coordinator.group.common.TopicMetadata;
-import org.apache.kafka.coordinator.group.classic.ClassicGroupState;
-import org.apache.kafka.coordinator.group.consumer.Assignment;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
-import org.apache.kafka.coordinator.group.consumer.MemberState;
-import org.apache.kafka.coordinator.group.consumer.TopicMetadata;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.GroupMetadataValue;
 import org.apache.kafka.coordinator.group.classic.ClassicGroup;
 import org.apache.kafka.coordinator.group.classic.ClassicGroupMember;
@@ -117,12 +114,10 @@
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkAssignment;
 import static org.apache.kafka.coordinator.group.AssignmentTestUtil.mkTopicAssignment;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.appendGroupMetadataErrorToResponseError;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRebalanceTimeoutKey;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupSessionTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.groupRebalanceTimeoutKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.EMPTY_RESULT;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupSyncKey;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.groupRevocationTimeoutKey;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.groupSessionTimeoutKey;
 import static org.apache.kafka.coordinator.group.RecordHelpersTest.mkMapOfPartitionRacks;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupMember.EMPTY_ASSIGNMENT;
@@ -2014,16 +2009,7 @@ public void testReconciliationProcess() {
             new ConsumerGroupHeartbeatResponseData()
                 .setMemberId(memberId1)
                 .setMemberEpoch(11)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Arrays.asList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Arrays.asList(0, 1)),
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(barTopicId)
-                            .setPartitions(Collections.singletonList(0))
-                    ))),
+                .setHeartbeatIntervalMs(5000),
             result.response()
         );
 
@@ -2796,7 +2782,7 @@ public void testOnNewMetadataImage() {
         });
 
         Collections.singletonList("group5").forEach(groupId -> {
-            ConsumerGroup group = context.groupMetadataManager.getOrMaybeCreateConsumerGroup(groupId, false);
+            ConsumerGroup group = context.groupMetadataManager.consumerGroup(groupId);
             assertFalse(group.hasMetadataExpired(context.time.milliseconds()));
         });
 
@@ -3143,78 +3129,6 @@ public void testRebalanceTimeoutLifecycle() {
             result.response()
         );
 
-        // Verify that there is a revocation timeout.
-        context.assertRevocationTimeout(groupId, memberId1, 12000);
-
-        assertEquals(
-            Collections.emptyList(),
-            context.sleep(result.response().heartbeatIntervalMs())
-        );
-
-        // Prepare next assignment.
-        assignor.prepareGroupAssignment(new GroupAssignment(
-            new HashMap<String, MemberAssignment>() {
-                {
-                    put(memberId1, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 0)
-                    )));
-                    put(memberId2, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 2)
-                    )));
-                    put(memberId3, new MemberAssignment(mkAssignment(
-                        mkTopicAssignment(fooTopicId, 1)
-                    )));
-                }
-            }
-        ));
-
-        // Member 3 joins the group.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId3)
-                .setMemberEpoch(0)
-                .setRebalanceTimeoutMs(90000)
-                .setSubscribedTopicNames(Collections.singletonList("foo"))
-                .setTopicPartitions(Collections.emptyList()));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId3)
-                .setMemberEpoch(3)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()),
-            result.response()
-        );
-
-        assertEquals(
-            Collections.emptyList(),
-            context.sleep(result.response().heartbeatIntervalMs())
-        );
-
-        // Member 1 heartbeats and re-transitions to revoking. The revocation timeout
-        // is re-scheduled.
-        result = context.consumerGroupHeartbeat(
-            new ConsumerGroupHeartbeatRequestData()
-                .setGroupId(groupId)
-                .setMemberId(memberId1)
-                .setMemberEpoch(1)
-                .setRebalanceTimeoutMs(90000)
-                .setSubscribedTopicNames(Collections.singletonList("foo")));
-
-        assertResponseEquals(
-            new ConsumerGroupHeartbeatResponseData()
-                .setMemberId(memberId1)
-                .setMemberEpoch(1)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Collections.singletonList(0))))),
-            result.response()
-        );
-
         // Verify that there is a revocation timeout. Keep a reference
         // to the timeout for later.
         ScheduledTimeout<Void, Record> scheduledTimeout =
@@ -3238,13 +3152,8 @@ public void testRebalanceTimeoutLifecycle() {
         assertResponseEquals(
             new ConsumerGroupHeartbeatResponseData()
                 .setMemberId(memberId1)
-                .setMemberEpoch(3)
-                .setHeartbeatIntervalMs(5000)
-                .setAssignment(new ConsumerGroupHeartbeatResponseData.Assignment()
-                    .setTopicPartitions(Collections.singletonList(
-                        new ConsumerGroupHeartbeatResponseData.TopicPartitions()
-                            .setTopicId(fooTopicId)
-                            .setPartitions(Collections.singletonList(0))))),
+                .setMemberEpoch(2)
+                .setHeartbeatIntervalMs(5000),
             result.response()
         );
 
@@ -3379,7 +3288,7 @@ public void testRebalanceTimeoutExpiration() {
         // Verify the expired timeout.
         assertEquals(
             Collections.singletonList(new ExpiredTimeout<Void, Record>(
-                consumerGroupRebalanceTimeoutKey(groupId, memberId1),
+                groupRebalanceTimeoutKey(groupId, memberId1),
                 new CoordinatorResult<>(
                     Arrays.asList(
                         RecordHelpers.newCurrentAssignmentTombstoneRecord(groupId, memberId1),
@@ -3447,7 +3356,7 @@ public void testOnLoaded() {
         assertNotNull(context.timer.timeout(groupSessionTimeoutKey("foo", "foo-2")));
 
         // foo-1 should also have a revocation timeout in place.
-        assertNotNull(context.timer.timeout(consumerGroupRebalanceTimeoutKey("foo", "foo-1")));
+        assertNotNull(context.timer.timeout(groupRebalanceTimeoutKey("foo", "foo-1")));
     }
 
     @Test
@@ -5293,7 +5202,7 @@ public void testReplaceStaticMemberInStableStateSucceeds(
         assertEquals(7000, newMember.rebalanceTimeoutMs());
         assertEquals(6000, newMember.sessionTimeoutMs());
         assertEquals(GroupMetadataManagerTestContext.toProtocols("range", "roundrobin"), newMember.supportedProtocols());
-        assertEquals(1, group.size());
+        assertEquals(1, group.numMembers());
         assertEquals(1, group.generationId());
         assertTrue(group.isInState(STABLE));
     }
@@ -5390,7 +5299,7 @@ public void testNewMemberTimeoutCompletion() throws Exception {
 
         // Make sure the NewMemberTimeout is not still in effect, and the member is not kicked
         GroupMetadataManagerTestContext.assertNoOrEmptyResult(context.sleep(context.classicGroupNewMemberJoinTimeoutMs));
-        assertEquals(1, group.size());
+        assertEquals(1, group.numMembers());
 
         // Member should be removed as heartbeat expires. The group is now empty.
         List<ExpiredTimeout<Void, Record>> timeouts = context.sleep(5000);
@@ -9328,7 +9237,6 @@ public void testConsumerGroupDelete() {
         GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder()
             .withConsumerGroup(new ConsumerGroupBuilder(groupId, 10))
             .build();
-        context.groupMetadataManager.getOrMaybeCreateConsumerGroup("group-id", true);
 
         List<Record> expectedRecords = Arrays.asList(
             RecordHelpers.newTargetAssignmentEpochTombstoneRecord(groupId),
@@ -9679,7 +9587,7 @@ public void testShareGroupUnknownGroupId() {
             .withShareGroupAssignor(assignor)
             .build();
 
-        assertThrows(GroupIdNotFoundException.class, () ->
+        assertThrows(IllegalStateException.class, () ->
             context.shareGroupHeartbeat(
                 new ShareGroupHeartbeatRequestData()
                     .setGroupId(groupId)
@@ -9786,7 +9694,7 @@ public void testShareGroupMemberJoinsEmptyGroupWithAssignments() {
             )))
         ));
 
-        assertThrows(GroupIdNotFoundException.class, () ->
+        assertThrows(IllegalStateException.class, () ->
             context.groupMetadataManager.getOrMaybeCreateShareGroup(groupId, false));
 
         CoordinatorResult<ShareGroupHeartbeatResponseData, Record> result = context.shareGroupHeartbeat(
@@ -9943,7 +9851,7 @@ public void testShareGroupDelete() {
             RecordHelpers.newGroupEpochTombstoneRecord("share-group-id")
         );
         List<Record> records = new ArrayList<>();
-        context.groupMetadataManager.deleteGroup("share-group-id", records);
+        context.groupMetadataManager.createGroupTombstoneRecords("share-group-id", records);
         assertEquals(expectedRecords, records);
     }
 
@@ -9983,7 +9891,6 @@ public void testShareGroupStates() {
         context.replay(RecordHelpers.newCurrentAssignmentRecord(groupId, new ShareGroupMember.Builder(memberId1)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2)))
             .build()), GroupType.SHARE);
 
@@ -9992,42 +9899,12 @@ public void testShareGroupStates() {
         context.replay(RecordHelpers.newCurrentAssignmentRecord(groupId, new ShareGroupMember.Builder(memberId1)
             .setMemberEpoch(11)
             .setPreviousMemberEpoch(10)
-            .setTargetMemberEpoch(11)
             .setAssignedPartitions(mkAssignment(mkTopicAssignment(fooTopicId, 1, 2, 3)))
             .build()), GroupType.SHARE);
 
         assertEquals(ShareGroup.ShareGroupState.STABLE, context.shareGroupState(groupId));
     }
 
-    private static void checkJoinGroupResponse(
-        JoinGroupResponseData expectedResponse,
-        JoinGroupResponseData actualResponse,
-        ClassicGroup group,
-        ClassicGroupState expectedState,
-        Set<String> expectedGroupInstanceIds
-    ) {
-        assertEquals(expectedResponse, actualResponse);
-        assertTrue(group.isInState(expectedState));
-
-        Set<String> groupInstanceIds = actualResponse.members()
-                                                     .stream()
-                                                     .map(JoinGroupResponseMember::groupInstanceId)
-                                                     .collect(Collectors.toSet());
-
-        assertEquals(expectedGroupInstanceIds, groupInstanceIds);
-    }
-
-    private static List<JoinGroupResponseMember> toJoinResponseMembers(ClassicGroup group) {
-        List<JoinGroupResponseMember> members = new ArrayList<>();
-        String protocolName = group.protocolName().get();
-        group.allMembers().forEach(member -> members.add(
-            new JoinGroupResponseMember()
-                .setMemberId(member.memberId())
-                .setGroupInstanceId(member.groupInstanceId().orElse(""))
-                .setMetadata(member.metadata(protocolName))
-        ));
-
-        return members;
     @Test
     public void testConsumerGroupHeartbeatWithNonEmptyClassicGroup() {
         String classicGroupId = "classic-group-id";
@@ -10190,7 +10067,7 @@ public void testClassicGroupOnUnloadedEmptyAndPreparingRebalance() throws Except
         assertFalse(joinResult1.joinFuture.isDone());
         assertFalse(joinResult2.joinFuture.isDone());
         assertTrue(preparingGroup.isInState(PREPARING_REBALANCE));
-        assertEquals(2, preparingGroup.size());
+        assertEquals(2, preparingGroup.numMembers());
 
         context.onUnloaded();
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
index 7fd1137341fc3..135bd31917552 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTestContext.java
@@ -30,6 +30,9 @@
 import org.apache.kafka.common.message.LeaveGroupRequestData;
 import org.apache.kafka.common.message.LeaveGroupResponseData;
 import org.apache.kafka.common.message.ListGroupsResponseData;
+import org.apache.kafka.common.message.ShareGroupDescribeResponseData;
+import org.apache.kafka.common.message.ShareGroupHeartbeatRequestData;
+import org.apache.kafka.common.message.ShareGroupHeartbeatResponseData;
 import org.apache.kafka.common.message.SyncGroupRequestData;
 import org.apache.kafka.common.message.SyncGroupResponseData;
 import org.apache.kafka.common.network.ClientInformation;
@@ -43,11 +46,12 @@
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.coordinator.group.Group.GroupType;
 import org.apache.kafka.coordinator.group.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.classic.ClassicGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroup;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupBuilder;
-import org.apache.kafka.coordinator.group.consumer.MemberState;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
@@ -66,6 +70,7 @@
 import org.apache.kafka.coordinator.group.generated.ShareGroupMemberMetadataValue;
 import org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetricsShard;
 import org.apache.kafka.coordinator.group.runtime.CoordinatorResult;
+import org.apache.kafka.coordinator.group.share.ShareGroup;
 import org.apache.kafka.image.MetadataImage;
 import org.apache.kafka.server.common.ApiMessageAndVersion;
 import org.apache.kafka.server.common.MetadataVersion;
@@ -87,8 +92,8 @@
 import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.EMPTY_RESULT;
 import static org.apache.kafka.coordinator.group.GroupMetadataManager.classicGroupHeartbeatKey;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupRebalanceTimeoutKey;
-import static org.apache.kafka.coordinator.group.GroupMetadataManager.consumerGroupSessionTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.groupRebalanceTimeoutKey;
+import static org.apache.kafka.coordinator.group.GroupMetadataManager.groupSessionTimeoutKey;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.COMPLETING_REBALANCE;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.DEAD;
 import static org.apache.kafka.coordinator.group.classic.ClassicGroupState.EMPTY;
@@ -346,6 +351,7 @@ public static class Builder {
         final private SnapshotRegistry snapshotRegistry = new SnapshotRegistry(logContext);
         private MetadataImage metadataImage;
         private List<PartitionAssignor> consumerGroupAssignors = Collections.singletonList(new MockPartitionAssignor("range"));
+        private PartitionAssignor shareGroupAssignor = new MockPartitionAssignor("share");
         final private List<ConsumerGroupBuilder> consumerGroupBuilders = new ArrayList<>();
         private int consumerGroupMaxSize = Integer.MAX_VALUE;
         private int consumerGroupMetadataRefreshIntervalMs = Integer.MAX_VALUE;
@@ -366,6 +372,11 @@ public Builder withAssignors(List<PartitionAssignor> assignors) {
             return this;
         }
 
+        public Builder withShareGroupAssignor(PartitionAssignor shareGroupAssignor) {
+            this.shareGroupAssignor = shareGroupAssignor;
+            return this;
+        }
+
         public Builder withConsumerGroup(ConsumerGroupBuilder builder) {
             this.consumerGroupBuilders.add(builder);
             return this;
@@ -420,6 +431,7 @@ public GroupMetadataManagerTestContext build() {
                     .withConsumerGroupSessionTimeout(45000)
                     .withConsumerGroupMaxSize(consumerGroupMaxSize)
                     .withConsumerGroupAssignors(consumerGroupAssignors)
+                    .withShareGroupAssignor(shareGroupAssignor)
                     .withConsumerGroupMetadataRefreshIntervalMs(consumerGroupMetadataRefreshIntervalMs)
                     .withClassicGroupMaxSize(classicGroupMaxSize)
                     .withClassicGroupMinSessionTimeoutMs(classicGroupMinSessionTimeoutMs)
@@ -489,6 +501,14 @@ public ConsumerGroup.ConsumerGroupState consumerGroupState(
             .state();
     }
 
+    public ShareGroup.ShareGroupState shareGroupState(
+        String groupId
+    ) {
+        return groupMetadataManager
+            .getOrMaybeCreateShareGroup(groupId, false)
+            .state();
+    }
+
     public MemberState consumerGroupMemberState(
         String groupId,
         String memberId
@@ -527,6 +547,34 @@ public CoordinatorResult<ConsumerGroupHeartbeatResponseData, Record> consumerGro
         return result;
     }
 
+    public CoordinatorResult<ShareGroupHeartbeatResponseData, Record> shareGroupHeartbeat(
+        ShareGroupHeartbeatRequestData request
+    ) {
+        RequestContext context = new RequestContext(
+            new RequestHeader(
+                ApiKeys.SHARE_GROUP_HEARTBEAT,
+                ApiKeys.SHARE_GROUP_HEARTBEAT.latestVersion(),
+                "client",
+                0
+            ),
+            "1",
+            InetAddress.getLoopbackAddress(),
+            KafkaPrincipal.ANONYMOUS,
+            ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+            SecurityProtocol.PLAINTEXT,
+            ClientInformation.EMPTY,
+            false
+        );
+
+        CoordinatorResult<ShareGroupHeartbeatResponseData, Record> result = groupMetadataManager.shareGroupHeartbeat(
+            context,
+            request
+        );
+
+        result.records().forEach(this::replay);
+        return result;
+    }
+
     public List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> sleep(long ms) {
         time.sleep(ms);
         List<MockCoordinatorTimer.ExpiredTimeout<Void, Record>> timeouts = timer.poll();
@@ -544,7 +592,7 @@ public void assertSessionTimeout(
         long delayMs
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
+            timer.timeout(groupSessionTimeoutKey(groupId, memberId));
         assertNotNull(timeout);
         assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
     }
@@ -554,7 +602,7 @@ public void assertNoSessionTimeout(
         String memberId
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupSessionTimeoutKey(groupId, memberId));
+            timer.timeout(groupSessionTimeoutKey(groupId, memberId));
         assertNull(timeout);
     }
 
@@ -564,7 +612,7 @@ public MockCoordinatorTimer.ScheduledTimeout<Void, Record> assertRebalanceTimeou
         long delayMs
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupRebalanceTimeoutKey(groupId, memberId));
+            timer.timeout(groupRebalanceTimeoutKey(groupId, memberId));
         assertNotNull(timeout);
         assertEquals(time.milliseconds() + delayMs, timeout.deadlineMs);
         return timeout;
@@ -575,7 +623,7 @@ public void assertNoRebalanceTimeout(
         String memberId
     ) {
         MockCoordinatorTimer.ScheduledTimeout<Void, Record> timeout =
-            timer.timeout(consumerGroupRebalanceTimeoutKey(groupId, memberId));
+            timer.timeout(groupRebalanceTimeoutKey(groupId, memberId));
         assertNull(timeout);
     }
 
@@ -838,7 +886,7 @@ public RebalanceResult staticMembersJoinAndRebalance(
         assertEquals(Errors.NONE.code(), followerJoinResult.joinFuture.get().errorCode());
         assertEquals(1, leaderJoinResult.joinFuture.get().generationId());
         assertEquals(1, followerJoinResult.joinFuture.get().generationId());
-        assertEquals(2, group.size());
+        assertEquals(2, group.numMembers());
         assertEquals(1, group.generationId());
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
@@ -888,7 +936,7 @@ public RebalanceResult staticMembersJoinAndRebalance(
         assertEquals(Errors.NONE.code(), followerSyncResult.syncFuture.get().errorCode());
         assertTrue(group.isInState(STABLE));
 
-        assertEquals(2, group.size());
+        assertEquals(2, group.numMembers());
         assertEquals(1, group.generationId());
 
         return new RebalanceResult(
@@ -1003,7 +1051,7 @@ public PendingMemberGroupResult setupGroupWithPendingMember(ClassicGroup group)
         assertTrue(followerJoinResult.records.isEmpty());
         assertFalse(followerJoinResult.joinFuture.isDone());
         assertTrue(group.isInState(PREPARING_REBALANCE));
-        assertEquals(2, group.size());
+        assertEquals(2, group.numMembers());
         assertEquals(1, group.numPendingJoinMembers());
 
         return new PendingMemberGroupResult(
@@ -1040,7 +1088,7 @@ public void verifySessionExpiration(ClassicGroup group, int timeoutMs) {
         assertEquals(expectedRecords, timeouts.get(timeoutsSize - 1).result.records());
         assertNoOrEmptyResult(timeouts.subList(0, timeoutsSize - 1));
         assertTrue(group.isInState(EMPTY));
-        assertEquals(0, group.size());
+        assertEquals(0, group.numMembers());
     }
 
     public HeartbeatResponseData sendClassicGroupHeartbeat(
@@ -1082,6 +1130,10 @@ public List<DescribeGroupsResponseData.DescribedGroup> describeGroups(List<Strin
         return groupMetadataManager.describeGroups(groupIds, lastCommittedOffset);
     }
 
+    public List<ShareGroupDescribeResponseData.DescribedGroup> shareGroupDescribe(List<String> groupIds) {
+        return groupMetadataManager.shareGroupDescribe(groupIds, lastCommittedOffset);
+    }
+
     public void verifyHeartbeat(
         String groupId,
         JoinGroupResponseData joinResponse,
@@ -1159,7 +1211,7 @@ public List<JoinGroupResponseData> joinWithNMembers(
             return null;
         }).collect(Collectors.toList());
 
-        assertEquals(numMembers, group.size());
+        assertEquals(numMembers, group.numMembers());
         assertTrue(group.isInState(COMPLETING_REBALANCE));
 
         return joinResponses;
@@ -1210,6 +1262,13 @@ private ApiMessage messageOrNull(ApiMessageAndVersion apiMessageAndVersion) {
 
     public void replay(
         Record record
+    ) {
+        replay(record, GroupType.CONSUMER);
+    }
+
+    public void replay(
+        Record record,
+        GroupType groupType
     ) {
         ApiMessageAndVersion key = record.key();
         ApiMessageAndVersion value = record.value();
@@ -1243,28 +1302,32 @@ public void replay(
             case ConsumerGroupPartitionMetadataKey.HIGHEST_SUPPORTED_VERSION:
                 groupMetadataManager.replay(
                     (ConsumerGroupPartitionMetadataKey) key.message(),
-                    (ConsumerGroupPartitionMetadataValue) messageOrNull(value)
+                    (ConsumerGroupPartitionMetadataValue) messageOrNull(value),
+                    groupType
                 );
                 break;
 
             case ConsumerGroupTargetAssignmentMemberKey.HIGHEST_SUPPORTED_VERSION:
                 groupMetadataManager.replay(
                     (ConsumerGroupTargetAssignmentMemberKey) key.message(),
-                    (ConsumerGroupTargetAssignmentMemberValue) messageOrNull(value)
+                    (ConsumerGroupTargetAssignmentMemberValue) messageOrNull(value),
+                    groupType
                 );
                 break;
 
             case ConsumerGroupTargetAssignmentMetadataKey.HIGHEST_SUPPORTED_VERSION:
                 groupMetadataManager.replay(
                     (ConsumerGroupTargetAssignmentMetadataKey) key.message(),
-                    (ConsumerGroupTargetAssignmentMetadataValue) messageOrNull(value)
+                    (ConsumerGroupTargetAssignmentMetadataValue) messageOrNull(value),
+                    groupType
                 );
                 break;
 
             case ConsumerGroupCurrentMemberAssignmentKey.HIGHEST_SUPPORTED_VERSION:
                 groupMetadataManager.replay(
                     (ConsumerGroupCurrentMemberAssignmentKey) key.message(),
-                    (ConsumerGroupCurrentMemberAssignmentValue) messageOrNull(value)
+                    (ConsumerGroupCurrentMemberAssignmentValue) messageOrNull(value),
+                    groupType
                 );
                 break;
 
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
index e56247a547586..fb3dc582e92a0 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/RecordHelpersTest.java
@@ -28,7 +28,7 @@
 import org.apache.kafka.coordinator.group.classic.ClassicGroupState;
 import org.apache.kafka.coordinator.group.common.TopicMetadata;
 import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
-import org.apache.kafka.coordinator.group.consumer.MemberState;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue;
 import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey;
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilderTest.java
index c1db6f228900b..fe608f14590e4 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/common/CurrentAssignmentBuilderTest.java
@@ -14,11 +14,13 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.coordinator.group.consumer;
+package org.apache.kafka.coordinator.group.common;
 
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.errors.FencedMemberEpochException;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData;
+import org.apache.kafka.coordinator.group.GroupMember;
+import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember;
 import org.junit.jupiter.api.Test;
 
 import java.util.Arrays;
@@ -45,7 +47,7 @@ public void testStableToStable() {
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3),
                 mkTopicAssignment(topicId2, 4, 5, 6))))
@@ -79,7 +81,7 @@ public void testStableToStableWithNewPartitions() {
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 4, 5, 6, 7))))
@@ -113,7 +115,7 @@ public void testStableToUnrevokedPartitions() {
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 5, 6, 7))))
@@ -150,7 +152,7 @@ public void testStableToUnreleasedPartitions() {
                 mkTopicAssignment(topicId2, 4, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 4, 5, 6, 7))))
@@ -187,7 +189,7 @@ public void testUnrevokedPartitionsToStable() {
                 mkTopicAssignment(topicId2, 4)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3),
                 mkTopicAssignment(topicId2, 5, 6))))
@@ -288,7 +290,7 @@ public void testUnrevokedPartitionsToUnrevokedPartitions() {
                 mkTopicAssignment(topicId2, 4)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6))))
@@ -332,7 +334,7 @@ public void testUnrevokedPartitionsToUnreleasedPartitions() {
                 mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 5, 6, 7))))
@@ -373,7 +375,7 @@ public void testUnreleasedPartitionsToStable() {
                 mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3),
                 mkTopicAssignment(topicId2, 5, 6))))
@@ -407,7 +409,7 @@ public void testUnreleasedPartitionsToStableWithNewPartitions() {
                 mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 5, 6, 7))))
@@ -441,7 +443,7 @@ public void testUnreleasedPartitionsToUnreleasedPartitions() {
                 mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(11, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 2, 3, 4),
                 mkTopicAssignment(topicId2, 5, 6, 7))))
@@ -465,7 +467,7 @@ public void testUnreleasedPartitionsToUnrevokedPartitions() {
                 mkTopicAssignment(topicId2, 5, 6)))
             .build();
 
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6))))
@@ -515,7 +517,7 @@ public void testUnknownState() {
             .build());
 
         // Then the member rejoins with no owned partitions.
-        ConsumerGroupMember updatedMember = new CurrentAssignmentBuilder(member)
+        GroupMember updatedMember = new CurrentAssignmentBuilder(member)
             .withTargetAssignment(12, new Assignment(mkAssignment(
                 mkTopicAssignment(topicId1, 3),
                 mkTopicAssignment(topicId2, 6))))
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
index 71813b9f47ff7..3f39f07962eb3 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/consumer/ConsumerGroupTest.java
@@ -29,6 +29,7 @@
 import org.apache.kafka.coordinator.group.OffsetAndMetadata;
 import org.apache.kafka.coordinator.group.OffsetExpirationCondition;
 import org.apache.kafka.coordinator.group.OffsetExpirationConditionImpl;
+import org.apache.kafka.coordinator.group.common.MemberState;
 import org.apache.kafka.coordinator.group.common.TopicMetadata;
 import org.apache.kafka.coordinator.group.metrics.GroupCoordinatorMetricsShard;
 import org.apache.kafka.image.MetadataImage;

From 9468e73ee81bd93d29f49edb386eeac3ed5abbb4 Mon Sep 17 00:00:00 2001
From: Apoorv Mittal <amittal@confluent.io>
Date: Thu, 11 Apr 2024 10:16:34 +0100
Subject: [PATCH 256/258] Resolving merge conflicts

---
 .../kafka/server/SharePartitionManager.java   | 28 +++++++++----------
 .../scala/kafka/server/BrokerServer.scala     |  1 +
 .../main/scala/kafka/server/KafkaApis.scala   |  1 +
 .../server/SharePartitionManagerTest.java     |  8 +++---
 .../kafka/api/BaseConsumerTest.scala          |  3 +-
 .../kafka/api/PlaintextConsumerTest.scala     |  8 +-----
 6 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/core/src/main/java/kafka/server/SharePartitionManager.java b/core/src/main/java/kafka/server/SharePartitionManager.java
index 1c7521984bc39..389b98a526abe 100644
--- a/core/src/main/java/kafka/server/SharePartitionManager.java
+++ b/core/src/main/java/kafka/server/SharePartitionManager.java
@@ -36,7 +36,6 @@
 import org.apache.kafka.storage.internals.log.FetchPartitionData;
 import org.slf4j.Logger;
 
-import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -45,7 +44,6 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
@@ -167,7 +165,7 @@ public void maybeProcessFetchQueue() {
                 shareFetchPartitionData.fetchParams,
                 CollectionConverters.asScala(
                     topicPartitionData.entrySet().stream().map(entry ->
-                        new Tuple2(entry.getKey(), entry.getValue())).collect(Collectors.toList())
+                        new Tuple2<>(entry.getKey(), entry.getValue())).collect(Collectors.toList())
                 ),
                 QuotaFactory.UnboundedQuota$.MODULE$,
                 responsePartitionData -> {
@@ -554,11 +552,11 @@ public enum ModifiedTopicIdPartitionType {
 
     // Helper class to return the erroneous partitions and valid partition data
     public static class ErroneousAndValidPartitionData {
-        private final List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous;
-        private final List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions;
+        private final List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous;
+        private final List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions;
 
-        public ErroneousAndValidPartitionData(List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous,
-                                              List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions) {
+        public ErroneousAndValidPartitionData(List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous,
+                                              List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions) {
             this.erroneous = erroneous;
             this.validTopicIdPartitions = validTopicIdPartitions;
         }
@@ -568,18 +566,18 @@ public ErroneousAndValidPartitionData(Map<TopicIdPartition, ShareFetchRequest.Sh
             validTopicIdPartitions = new ArrayList<>();
             shareFetchData.forEach((topicIdPartition, sharePartitionData) -> {
                 if (topicIdPartition.topic() == null) {
-                    erroneous.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
+                    erroneous.add(new Tuple2<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
                 } else {
-                    validTopicIdPartitions.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, sharePartitionData));
+                    validTopicIdPartitions.add(new Tuple2<>(topicIdPartition, sharePartitionData));
                 }
             });
         }
 
-        public List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous() {
+        public List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous() {
             return erroneous;
         }
 
-        public List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions() {
+        public List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> validTopicIdPartitions() {
             return validTopicIdPartitions;
         }
     }
@@ -815,8 +813,8 @@ ErroneousAndValidPartitionData getErroneousAndValidTopicIdPartitions() {
             if (!isSubsequent) {
                 return new ErroneousAndValidPartitionData(shareFetchData);
             } else {
-                List<Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous = new ArrayList<>();
-                List<Entry<TopicIdPartition, ShareFetchRequest.SharePartitionData>> valid = new ArrayList<>();
+                List<Tuple2<TopicIdPartition, ShareFetchResponseData.PartitionData>> erroneous = new ArrayList<>();
+                List<Tuple2<TopicIdPartition, ShareFetchRequest.SharePartitionData>> valid = new ArrayList<>();
                 // Take the session lock and iterate over all the cached partitions.
                 synchronized (session) {
                     session.partitionMap.forEach(cachedSharePartition -> {
@@ -824,9 +822,9 @@ ErroneousAndValidPartitionData getErroneousAndValidTopicIdPartitions() {
                                 TopicPartition(cachedSharePartition.topic, cachedSharePartition.partition));
                         ShareFetchRequest.SharePartitionData reqData = cachedSharePartition.reqData();
                         if (topicIdPartition.topic() == null) {
-                            erroneous.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
+                            erroneous.add(new Tuple2<>(topicIdPartition, ShareFetchResponse.partitionResponse(topicIdPartition, Errors.UNKNOWN_TOPIC_ID)));
                         } else {
-                            valid.add(new AbstractMap.SimpleImmutableEntry<>(topicIdPartition, reqData));
+                            valid.add(new Tuple2<>(topicIdPartition, reqData));
                         }
                     });
                     return new ErroneousAndValidPartitionData(erroneous, valid);
diff --git a/core/src/main/scala/kafka/server/BrokerServer.scala b/core/src/main/scala/kafka/server/BrokerServer.scala
index 57dd5b6d6629d..5dd865055fc65 100644
--- a/core/src/main/scala/kafka/server/BrokerServer.scala
+++ b/core/src/main/scala/kafka/server/BrokerServer.scala
@@ -24,6 +24,7 @@ import kafka.log.LogManager
 import kafka.log.remote.RemoteLogManager
 import kafka.network.{DataPlaneAcceptor, SocketServer}
 import kafka.raft.KafkaRaftManager
+import kafka.server.SharePartitionManager.ShareSessionCache
 import kafka.server.metadata._
 import kafka.utils.CoreUtils
 import org.apache.kafka.common.config.ConfigException
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index b8558432960b4..ba34354d540db 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -1291,6 +1291,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     // Regular Kafka consumers need READ permission on each partition they are fetching.
     val partitionDatas = new mutable.ArrayBuffer[(TopicIdPartition, ShareFetchRequest.SharePartitionData)]
     val erroneousAndValidPartitionData : ErroneousAndValidPartitionData = shareFetchContext.getErroneousAndValidTopicIdPartitions
+
     erroneousAndValidPartitionData.erroneous.forEach {
       erroneousData => erroneous += erroneousData
     }
diff --git a/core/src/test/java/kafka/server/SharePartitionManagerTest.java b/core/src/test/java/kafka/server/SharePartitionManagerTest.java
index 712232648b641..8cd2bf51c12aa 100644
--- a/core/src/test/java/kafka/server/SharePartitionManagerTest.java
+++ b/core/src/test/java/kafka/server/SharePartitionManagerTest.java
@@ -813,10 +813,10 @@ private void assertErroneousAndValidTopicIdPartitions(SharePartitionManager.Erro
                                                     List<TopicIdPartition> expectedErroneous, List<TopicIdPartition> expectedValid) {
         List<TopicIdPartition> actualErroneousPartitions = new ArrayList<>();
         List<TopicIdPartition> actualValidPartitions = new ArrayList<>();
-        erroneousAndValidPartitionData.erroneous().forEach(topicIdPartitionPartitionDataEntry ->
-                actualErroneousPartitions.add(topicIdPartitionPartitionDataEntry.getKey()));
-        erroneousAndValidPartitionData.validTopicIdPartitions().forEach(topicIdPartitionPartitionDataEntry ->
-                actualValidPartitions.add(topicIdPartitionPartitionDataEntry.getKey()));
+        erroneousAndValidPartitionData.erroneous().forEach(topicIdPartitionPartitionDataTuple2 ->
+            actualErroneousPartitions.add(topicIdPartitionPartitionDataTuple2._1));
+        erroneousAndValidPartitionData.validTopicIdPartitions().forEach(topicIdPartitionPartitionDataTuple2 ->
+            actualValidPartitions.add(topicIdPartitionPartitionDataTuple2._1));
         assertEquals(expectedErroneous, actualErroneousPartitions);
         assertEquals(expectedValid, actualValidPartitions);
     }
diff --git a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
index 816edc40bea2f..ba795c4b66195 100644
--- a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
@@ -30,6 +30,7 @@ import org.junit.jupiter.params.provider.{Arguments, MethodSource}
 import java.util
 import java.util.{Map, Properties}
 import java.util.concurrent.atomic.AtomicInteger
+import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
 /**
@@ -96,7 +97,7 @@ abstract class BaseConsumerTest extends AbstractConsumerTest {
     assertEquals(1, listener.callsToAssigned)
 
     // get metadata for the topic
-    var parts: Seq[PartitionInfo] = null
+    var parts: mutable.Buffer[PartitionInfo] = null
     while (parts == null)
       parts = consumer.partitionsFor(Topic.GROUP_METADATA_TOPIC_NAME).asScala
     assertEquals(1, parts.size)
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index 651ca35ffe9cd..8adee2ce4597b 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -15,13 +15,9 @@ package kafka.api
 import java.time.Duration
 import java.util
 import java.util.Arrays.asList
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.locks.ReentrantLock
-import java.util.regex.Pattern
-import java.util.{Locale, Optional, Properties}
+import java.util.{Locale, Properties}
 import kafka.api.BaseConsumerTest.SerializerImpl
 import kafka.api.BaseConsumerTest.DeserializerImpl
-import java.util.{Locale, Properties}
 import kafka.server.{KafkaBroker, QuotaType}
 import kafka.utils.{TestInfoUtils, TestUtils}
 import org.apache.kafka.clients.admin.{NewPartitions, NewTopic}
@@ -29,9 +25,7 @@ import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.{KafkaException, MetricName, TopicPartition}
 import org.apache.kafka.common.config.TopicConfig
-import org.apache.kafka.common.errors.{InvalidGroupIdException, InvalidTopicException, UnsupportedAssignorException}
 import org.apache.kafka.common.errors.{InvalidGroupIdException, InvalidTopicException}
-import org.apache.kafka.common.header.Headers
 import org.apache.kafka.common.record.{CompressionType, TimestampType}
 import org.apache.kafka.common.serialization._
 import org.apache.kafka.test.{MockConsumerInterceptor, MockProducerInterceptor}

From 91140c647c3d42764b59701e44b96a63a9357c17 Mon Sep 17 00:00:00 2001
From: Apoorv Mittal <amittal@confluent.io>
Date: Thu, 11 Apr 2024 11:11:48 +0100
Subject: [PATCH 257/258] Fixing build

---
 .../java/org/apache/kafka/clients/admin/KafkaAdminClient.java   | 1 -
 .../internals/events/ShareUnsubscribeApplicationEvent.java      | 2 --
 2 files changed, 3 deletions(-)

diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
index 5cccf546214fb..5f79255385351 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
@@ -59,7 +59,6 @@
 import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.ConsumerGroupState;
-import org.apache.kafka.common.GroupType;
 import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.GroupType;
 import org.apache.kafka.common.KafkaException;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
index 4676d70e717fc..7378f5ea0bde3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeApplicationEvent.java
@@ -17,8 +17,6 @@
 
 package org.apache.kafka.clients.consumer.internals.events;
 
-import org.apache.kafka.common.utils.Timer;
-
 /**
  * Application event triggered when a user calls the unsubscribe API. This will make the consumer
  * release all its assignments and send a heartbeat request to leave the share group.

From 81738beed24b2ebd52809c77641d85943c239c55 Mon Sep 17 00:00:00 2001
From: Apoorv Mittal <amittal@confluent.io>
Date: Thu, 11 Apr 2024 14:11:16 +0100
Subject: [PATCH 258/258] Fixing build

---
 checkstyle/import-control.xml                                 | 2 ++
 .../main/java/org/apache/kafka/tools/ShareGroupsCommand.java  | 3 +--
 .../java/org/apache/kafka/tools/consumer/ConsoleConsumer.java | 2 --
 .../apache/kafka/tools/consumer/ConsoleConsumerOptions.java   | 1 -
 .../org/apache/kafka/tools/consumer/ConsoleConsumerTest.java  | 4 ----
 .../kafka/tools/consumer/group/ConsumerGroupCommandTest.java  | 1 -
 6 files changed, 3 insertions(+), 10 deletions(-)

diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index d8c8d6a20c7a4..244b5a3bf2c85 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -327,6 +327,8 @@
 
       <subpackage name="group">
         <allow pkg="kafka.api"/>
+        <allow pkg="kafka.security"/>
+        <allow pkg="kafka.zk" />
         <allow pkg="org.apache.kafka.tools"/>
         <allow pkg="org.apache.kafka.server.config" />
         <allow pkg="scala"/>
diff --git a/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java b/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
index d3845b9eeaeb2..20290247fb65a 100644
--- a/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
+++ b/tools/src/main/java/org/apache/kafka/tools/ShareGroupsCommand.java
@@ -32,7 +32,6 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.server.util.CommandLineUtils;
-import org.apache.kafka.timeline.TimelineHashMap;
 
 import java.io.IOException;
 import java.util.AbstractMap;
@@ -130,7 +129,7 @@ public void listGroups() throws ExecutionException, InterruptedException {
             ? Collections.emptySet()
             : shareGroupStatesFromString(stateValue);
         List<ShareGroupListing> listings = listShareGroupsWithState(states);
-        printGroupStates(listings.stream().map(e -> new AbstractMap.SimpleImmutableEntry<String, String>(e.groupId(), e.state().toString()) {}).collect(Collectors.toList()));
+        printGroupStates(listings.stream().map(e -> new AbstractMap.SimpleImmutableEntry<String, String>(e.groupId(), e.state().toString()) { }).collect(Collectors.toList()));
       } else
         listShareGroups().forEach(System.out::println);
     }
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
index eaf3f5af2b2a1..bb5ab1443ed49 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java
@@ -22,8 +22,6 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
 import java.util.concurrent.CountDownLatch;
 import java.util.regex.Pattern;
 import java.util.Collections;
diff --git a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
index d60c6f4e9ac27..aa37919515450 100644
--- a/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
+++ b/tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumerOptions.java
@@ -34,7 +34,6 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Properties;
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
index 9195f61873605..f4fa6ac3be221 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java
@@ -24,7 +24,6 @@
 import org.apache.kafka.common.MessageFormatter;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.requests.ListOffsetsRequest;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.util.MockTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -37,9 +36,6 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
 import java.util.regex.Pattern;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
diff --git a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
index 1da0817a9255e..e1daa37f1976b 100644
--- a/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/consumer/group/ConsumerGroupCommandTest.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.tools.consumer.group;
 
 import kafka.api.BaseConsumerTest;
-import kafka.admin.ConsumerGroupCommand;
 import kafka.server.KafkaConfig;
 import kafka.utils.TestUtils;
 import org.apache.kafka.clients.admin.AdminClientConfig;